Spaces:
Sleeping
Sleeping
| import spaces | |
| import gradio as gr | |
| from PIL import Image | |
| from src.tryon_pipeline import StableDiffusionXLInpaintPipeline as TryonPipeline | |
| from src.unet_hacked_garmnet import UNet2DConditionModel as UNet2DConditionModel_ref | |
| from src.unet_hacked_tryon import UNet2DConditionModel | |
| from transformers import ( | |
| CLIPImageProcessor, | |
| CLIPVisionModelWithProjection, | |
| CLIPTextModel, | |
| CLIPTextModelWithProjection, | |
| ) | |
| from diffusers import DDPMScheduler,AutoencoderKL | |
| from typing import List | |
| import torch | |
| import os | |
| import io | |
| import warnings | |
| import requests | |
| from transformers import AutoTokenizer | |
| import numpy as np | |
| from utils_mask import get_mask_location | |
| from torchvision import transforms | |
| import apply_net | |
| from preprocess.humanparsing.run_parsing import Parsing | |
| from preprocess.openpose.run_openpose import OpenPose | |
| from detectron2.data.detection_utils import convert_PIL_to_numpy,_apply_exif_orientation | |
| from torchvision.transforms.functional import to_pil_image | |
| # import pillow_heif # HEIC μ΄λ―Έμ§ μ²λ¦¬μ© (μμ΄ν° 촬μ μ¬μ§ ν¬λ§·) | |
| from urllib.parse import urlparse | |
| # zeroGPU νκ²½μμ compile μ¬μ© μ¬λΆ | |
| is_compile_for_zeroGPU = False # True: compile μ¬μ©, False: compile μ¬μ© μ ν¨ | |
| # SSL κ²½κ³ μ΅μ | |
| warnings.filterwarnings("ignore", message=".*OpenSSL.*") | |
| warnings.filterwarnings("ignore", category=UserWarning, module="urllib3") | |
| # requests μΈμ μ€μ | |
| session = requests.Session() | |
| session.verify = False # SSL κ²μ¦ λΉνμ±ν (κ°λ° νκ²½μ©) | |
| def pil_to_binary_mask(pil_image, threshold=0): | |
| np_image = np.array(pil_image) | |
| grayscale_image = Image.fromarray(np_image).convert("L") | |
| binary_mask = np.array(grayscale_image) > threshold | |
| mask = np.zeros(binary_mask.shape, dtype=np.uint8) | |
| for i in range(binary_mask.shape[0]): | |
| for j in range(binary_mask.shape[1]): | |
| if binary_mask[i,j] == True : | |
| mask[i,j] = 1 | |
| mask = (mask*255).astype(np.uint8) | |
| output_mask = Image.fromarray(mask) | |
| return output_mask | |
| print("=" * 60) | |
| print("Starting GENAI-VTON Application Initialization") | |
| print("=" * 60) | |
| base_path = 'yisol/IDM-VTON' | |
| example_path = os.path.join(os.path.dirname(__file__), 'example') | |
| print("\n[1/10] Loading UNet model...") | |
| unet = UNet2DConditionModel.from_pretrained( | |
| base_path, | |
| subfolder="unet", | |
| torch_dtype=torch.float16, | |
| ) | |
| unet.requires_grad_(False) | |
| # torch.compile() μ μ© - μΆλ‘ μλ 20-40% ν₯μ (PyTorch 2.0+) | |
| # μ£Όμ: 첫 λ²μ§Έ μΆλ‘ μ μ»΄νμΌλ‘ μΈν΄ λ릴 μ μμ | |
| if is_compile_for_zeroGPU == True: | |
| print("β UNet model loaded successfully") | |
| else: | |
| if hasattr(torch, 'compile'): | |
| try: | |
| unet = torch.compile(unet, mode="reduce-overhead") | |
| print("β UNet model loaded and compiled successfully") | |
| except Exception as e: | |
| print(f"β UNet model loaded (compile skipped: {e})") | |
| else: | |
| print("β UNet model loaded successfully") | |
| print("\n[2/10] Loading tokenizers...") | |
| tokenizer_one = AutoTokenizer.from_pretrained( | |
| base_path, | |
| subfolder="tokenizer", | |
| revision=None, | |
| use_fast=False, | |
| ) | |
| tokenizer_two = AutoTokenizer.from_pretrained( | |
| base_path, | |
| subfolder="tokenizer_2", | |
| revision=None, | |
| use_fast=False, | |
| ) | |
| print("β Tokenizers loaded successfully") | |
| print("\n[3/10] Loading noise scheduler...") | |
| noise_scheduler = DDPMScheduler.from_pretrained(base_path, subfolder="scheduler") | |
| print("β Noise scheduler loaded successfully") | |
| print("\n[4/10] Loading text encoders...") | |
| text_encoder_one = CLIPTextModel.from_pretrained( | |
| base_path, | |
| subfolder="text_encoder", | |
| torch_dtype=torch.float16, | |
| ) | |
| text_encoder_two = CLIPTextModelWithProjection.from_pretrained( | |
| base_path, | |
| subfolder="text_encoder_2", | |
| torch_dtype=torch.float16, | |
| ) | |
| print("β Text encoders loaded successfully") | |
| print("\n[5/10] Loading image encoder...") | |
| image_encoder = CLIPVisionModelWithProjection.from_pretrained( | |
| base_path, | |
| subfolder="image_encoder", | |
| torch_dtype=torch.float16, | |
| ) | |
| print("β Image encoder loaded successfully") | |
| print("\n[6/10] Loading VAE...") | |
| vae = AutoencoderKL.from_pretrained(base_path, | |
| subfolder="vae", | |
| torch_dtype=torch.float16, | |
| ) | |
| # torch.compile() μ μ© - VAE μΈμ½λ©/λμ½λ© μλ ν₯μ | |
| if is_compile_for_zeroGPU == True: | |
| print("β VAE loaded successfully") | |
| else: | |
| if hasattr(torch, 'compile'): | |
| try: | |
| vae = torch.compile(vae, mode="reduce-overhead") | |
| print("β VAE loaded and compiled successfully") | |
| except Exception as e: | |
| print(f"β VAE loaded (compile skipped: {e})") | |
| else: | |
| print("β VAE loaded successfully") | |
| print("\n[7/10] Loading UNet Encoder...") | |
| UNet_Encoder = UNet2DConditionModel_ref.from_pretrained( | |
| base_path, | |
| subfolder="unet_encoder", | |
| torch_dtype=torch.float16, | |
| ) | |
| # torch.compile() μ μ© - UNet Encoder μλ ν₯μ | |
| if is_compile_for_zeroGPU == True: | |
| print("β UNet Encoder loaded successfully") | |
| else: | |
| if hasattr(torch, 'compile'): | |
| try: | |
| UNet_Encoder = torch.compile(UNet_Encoder, mode="reduce-overhead") | |
| print("β UNet Encoder loaded and compiled successfully") | |
| except Exception as e: | |
| print(f"β UNet Encoder loaded (compile skipped: {e})") | |
| else: | |
| print("β UNet Encoder loaded successfully") | |
| print("\n[8/10] Initializing parsing and openpose models...") | |
| parsing_model = Parsing(0) | |
| openpose_model = OpenPose(0) | |
| print("β Parsing and OpenPose models initialized") | |
| print("\n[9/10] Configuring model parameters...") | |
| UNet_Encoder.requires_grad_(False) | |
| image_encoder.requires_grad_(False) | |
| vae.requires_grad_(False) | |
| unet.requires_grad_(False) | |
| text_encoder_one.requires_grad_(False) | |
| text_encoder_two.requires_grad_(False) | |
| tensor_transfrom = transforms.Compose( | |
| [ | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.5], [0.5]), | |
| ] | |
| ) | |
| print("β Model parameters configured") | |
| print("\n[10/10] Initializing TryonPipeline...") | |
| pipe = TryonPipeline.from_pretrained( | |
| base_path, | |
| unet=unet, | |
| vae=vae, | |
| feature_extractor= CLIPImageProcessor(), | |
| text_encoder = text_encoder_one, | |
| text_encoder_2 = text_encoder_two, | |
| tokenizer = tokenizer_one, | |
| tokenizer_2 = tokenizer_two, | |
| scheduler = noise_scheduler, | |
| image_encoder=image_encoder, | |
| torch_dtype=torch.float16, | |
| ) | |
| pipe.unet_encoder = UNet_Encoder | |
| print("β TryonPipeline initialized successfully") | |
| # torch, diffusers λ± λ²μ μ 리 ν μ μ© κ°λ₯. | |
| # # xFormers λ©λͺ¨λ¦¬ ν¨μ¨μ μ΄ν μ νμ±ν (λ©λͺ¨λ¦¬ 20-30% κ°μ, μλ 10-20% ν₯μ) | |
| # print("\n[Optimization] Enabling xFormers memory efficient attention...") | |
| # try: | |
| # pipe.enable_xformers_memory_efficient_attention() | |
| # print("β xFormers memory efficient attention enabled") | |
| # except Exception as e: | |
| # print(f"β xFormers not available, using default attention: {e}") | |
| print("\n" + "=" * 60) | |
| print("All models loaded successfully!") | |
| print("=" * 60 + "\n") | |
| # Warm-up: 첫 λ²μ§Έ μΆλ‘ μ§μ° κ°μλ₯Ό μν λͺ¨λΈ μ΄κΈ°ν | |
| # JIT μ»΄νμΌ, CUDA 컀λ λ‘λ© λ±μ 미리 μν | |
| print("=" * 60) | |
| print("Warming up models (CPU)...") | |
| print("=" * 60) | |
| def warmup_models_cpu(): | |
| """μ± μμ μ CPU λͺ¨λΈ μ΄κΈ°νλ₯Ό μν Warm-up ν¨μ""" | |
| try: | |
| # CPUμμ ν μ€νΈ μλ² λ© Warm-up (Tokenizer + Text Encoder μ΄κΈ°ν) | |
| print("[CPU Warm-up 1/2] Text Encoder warm-up...") | |
| with torch.no_grad(): | |
| dummy_prompt = "a photo of clothing" | |
| dummy_tokens = tokenizer_one( | |
| dummy_prompt, | |
| padding="max_length", | |
| max_length=tokenizer_one.model_max_length, | |
| truncation=True, | |
| return_tensors="pt" | |
| ) | |
| # CPUμμ μ€ν κ°λ₯ν μ΄κΈ°ν | |
| _ = text_encoder_one(dummy_tokens.input_ids, output_hidden_states=True) | |
| print("β Text Encoder warmed up") | |
| # Tensor λ³ν Warm-up | |
| print("[CPU Warm-up 2/2] Tensor transform warm-up...") | |
| dummy_img = Image.new('RGB', (768, 1024), color='white') | |
| _ = tensor_transfrom(dummy_img) | |
| print("β Tensor transform warmed up") | |
| return True | |
| except Exception as e: | |
| print(f"β CPU Warm-up partially completed: {e}") | |
| return False | |
| # CPU Warm-up μ€ν | |
| warmup_success = warmup_models_cpu() | |
| if warmup_success: | |
| print("\nβ CPU warm-up completed successfully") | |
| else: | |
| print("\nβ CPU warm-up completed with warnings") | |
| print("=" * 60 + "\n") | |
| # torch.compile μ€λ₯ μ eager λͺ¨λλ‘ ν΄λ°± μ€μ | |
| # 컀μ€ν UNet forward λ©μλ νΈνμ± λ¬Έμ λμ | |
| if is_compile_for_zeroGPU == True: | |
| print("β torch.compile is disabled for ZeroGPU") | |
| else: | |
| try: | |
| import torch._dynamo | |
| torch._dynamo.config.suppress_errors = True | |
| print("β torch._dynamo.config.suppress_errors enabled (fallback to eager mode on error)") | |
| except Exception as e: | |
| print(f"β torch._dynamo config not available: {e}") | |
| # GPU Warm-up ν¨μ (μ± λ‘λ μ μλ μ€ν) | |
| # Text Encoder, VAE GPU λ‘λ© λ° CUDA 컀λ μ΄κΈ°ν | |
| def warmup_gpu(): | |
| """μ± λ‘λ μ GPU λͺ¨λΈ μ΄κΈ°νλ₯Ό μν Warm-up ν¨μ""" | |
| try: | |
| device = "cuda" | |
| print("=" * 60) | |
| print("GPU Warm-up: Loading models to GPU and initializing CUDA kernels...") | |
| print("=" * 60) | |
| # λͺ¨λΈμ GPUλ‘ μ΄λ | |
| print("[GPU Warm-up 1/4] Moving models to GPU...") | |
| pipe.to(device) | |
| pipe.unet_encoder.to(device) | |
| print("β Models moved to GPU") | |
| # λλ―Έ ν μ μμ± | |
| with torch.no_grad(): | |
| with torch.cuda.amp.autocast(): | |
| # 1. λλ―Έ ν둬ννΈ μλ² λ© μμ± (Text Encoder GPU warm-up) | |
| print("[GPU Warm-up 2/4] Text Encoder GPU warm-up...") | |
| dummy_prompt = "a photo of white t-shirt" | |
| _ = pipe.encode_prompt( | |
| dummy_prompt, | |
| num_images_per_prompt=1, | |
| do_classifier_free_guidance=True, | |
| negative_prompt="low quality", | |
| ) | |
| print("β Text Encoder GPU warmed up") | |
| # 2. λλ―Έ μ΄λ―Έμ§λ‘ VAE μΈμ½λ©/λμ½λ© (VAE GPU warm-up) | |
| print("[GPU Warm-up 3/4] VAE GPU warm-up...") | |
| dummy_img = torch.randn(1, 3, 1024, 768).to(device, torch.float16) | |
| latents = pipe.vae.encode(dummy_img).latent_dist.sample() | |
| _ = pipe.vae.decode(latents) | |
| print("β VAE GPU warmed up (encode + decode)") | |
| # 3. CUDA λκΈ°ν (컀λ λ‘λ© μλ£ λκΈ°) | |
| print("[GPU Warm-up 4/4] CUDA synchronization...") | |
| torch.cuda.synchronize() | |
| print("β CUDA kernels initialized") | |
| # GPU λ©λͺ¨λ¦¬ μ 리 | |
| torch.cuda.empty_cache() | |
| print("\n" + "=" * 60) | |
| print("β GPU Warm-up completed!") | |
| print(" Text Encoder, VAE ready. UNet will compile on first request.") | |
| print(" (torch.compile errors will fallback to eager mode)") | |
| print("=" * 60 + "\n") | |
| return "GPU Warm-up completed successfully!" | |
| except Exception as e: | |
| print(f"\nβ GPU Warm-up failed: {e}") | |
| print(" Models will be loaded on first user request.") | |
| return f"GPU Warm-up skipped: {e}" | |
| # μ΄λ―Έμ§ μ μ²λ¦¬ ν¨μ | |
| def preprocess_image(image): | |
| # HEIC μ΄λ―Έμ§ μ²λ¦¬ | |
| if isinstance(image, np.ndarray): | |
| image = Image.fromarray(image) | |
| # HEIC μ΄λ―Έμ§λ₯Ό JPEGλ‘ λ³ν - μ΄κ±° μ λ¨Ήνλ κ±° κ°μλ°.... | |
| try: | |
| output = io.BytesIO() | |
| image.convert("RGB").save(output, format="JPEG", quality=95) | |
| output.seek(0) | |
| image = Image.open(output) | |
| except Exception as e: | |
| print(f"Error converting image: {e}") | |
| # λ³ν μ€ν¨ μ μλ³Έ μ΄λ―Έμ§ μ¬μ© | |
| image = image.convert("RGB") | |
| # μ΄λ―Έμ§ ν¬κΈ° κ°μ Έμ€κΈ° | |
| width, height = image.size | |
| # 3:4 λΉμ¨λ‘ μ€μ μλ₯΄κΈ° | |
| target_width = int(min(width, height * (3 / 4))) | |
| target_height = int(min(height, width * (4 / 3))) | |
| left = (width - target_width) / 2 | |
| top = (height - target_height) / 2 | |
| right = (width + target_width) / 2 | |
| bottom = (height + target_height) / 2 | |
| # μ΄λ―Έμ§ μλ₯΄κΈ° | |
| cropped_img = image.crop((left, top, right, bottom)) | |
| # 768x1024λ‘ λ¦¬μ¬μ΄μ§ | |
| resized_img = cropped_img.resize((768, 1024), resample=Image.Resampling.LANCZOS) | |
| return resized_img | |
| # URLμμ μ΄λ―Έμ§ κ°μ Έμ€κΈ° ν¨μ | |
| def load_image_from_url(url): | |
| try: | |
| response = session.get(url, stream=True, timeout=10) | |
| response.raise_for_status() # HTTP μ€λ₯ νμΈ | |
| # μ΄λ―Έμ§ λ€μ΄λ‘λ | |
| img = Image.open(response.raw).convert("RGB") | |
| # JPEGλ‘ λ³ν | |
| output = io.BytesIO() | |
| img.save(output, format="JPEG", quality=95) | |
| output.seek(0) | |
| # λ³νλ JPEG μ΄λ―Έμ§ λ°ν | |
| jpeg_img = Image.open(output) | |
| return jpeg_img | |
| except requests.exceptions.RequestException as e: | |
| print(f"Error downloading image from URL: {e}") | |
| return None | |
| except Exception as e: | |
| print(f"Error processing image from URL: {e}") | |
| return None | |
| def process_url_image(url): | |
| """Process image from URL and return PIL Image""" | |
| if not url or not url.strip(): | |
| return None | |
| # URL μ ν¨μ± κ²μ¬ | |
| try: | |
| result = urlparse(url) | |
| if not all([result.scheme, result.netloc]): | |
| print("Invalid URL format") | |
| return None | |
| except Exception as e: | |
| print(f"Error parsing URL: {e}") | |
| return None | |
| img = load_image_from_url(url) | |
| if img is None: | |
| print("Failed to load image from URL") | |
| return None | |
| return preprocess_image(img) | |
| def load_example_for_editor(image_path): | |
| """Load example image for ImageEditor component""" | |
| if image_path is None: | |
| return None | |
| # ImageEditorλ νΉμ νμμ κΈ°λνλ―λ‘ λμ λ리 ννλ‘ λ°ν | |
| return { | |
| "background": image_path, | |
| "layers": None, | |
| "composite": None | |
| } | |
| def download_model_file(model_path, urls): | |
| """Download model file from multiple URLs if it doesn't exist""" | |
| if os.path.exists(model_path): | |
| print(f"Model file already exists: {model_path}") | |
| return True | |
| os.makedirs(os.path.dirname(model_path), exist_ok=True) | |
| for url in urls: | |
| try: | |
| print(f"Downloading from: {url}") | |
| response = requests.get(url, stream=True) | |
| response.raise_for_status() | |
| total_size = int(response.headers.get('content-length', 0)) | |
| block_size = 8192 | |
| with open(model_path, 'wb') as f: | |
| downloaded = 0 | |
| for chunk in response.iter_content(chunk_size=block_size): | |
| if chunk: | |
| f.write(chunk) | |
| downloaded += len(chunk) | |
| if total_size > 0: | |
| percent = (downloaded / total_size) * 100 | |
| if(percent % 10 == 0): | |
| print(f"\rDownload progress: {percent:.1f}%", end='', flush=True) | |
| print(f"\nSuccessfully downloaded: {model_path}") | |
| return True | |
| except Exception as e: | |
| print(f"Failed to download from {url}: {e}") | |
| continue | |
| print(f"Failed to download model file from all URLs: {model_path}") | |
| return False | |
| def download_densepose_model(): | |
| """Download DensePose model file""" | |
| model_path = "ckpt/densepose/model_final_162be9.pkl" | |
| urls = [ | |
| "https://dl.fbaipublicfiles.com/densepose/densepose_rcnn_R_50_FPN_s1x/165712039/model_final_162be9.pkl", | |
| "https://github.com/facebookresearch/densepose/releases/download/v1.0/model_final_162be9.pkl" | |
| ] | |
| return download_model_file(model_path, urls) | |
| def download_openpose_model(): | |
| """Download OpenPose model file""" | |
| model_path = "ckpt/openpose/ckpts/body_pose_model.pth" | |
| urls = [ | |
| "https://huggingface.co/lllyasviel/Annotators/resolve/main/body_pose_model.pth" | |
| ] | |
| return download_model_file(model_path, urls) | |
| def download_humanparsing_models(): | |
| """Download Human Parsing model files""" | |
| base_url = "https://huggingface.co/Longcat2957/humanparsing-onnx/resolve/main" | |
| models = [ | |
| ("ckpt/humanparsing/parsing_atr.onnx", f"{base_url}/parsing_atr.onnx"), | |
| ("ckpt/humanparsing/parsing_lip.onnx", f"{base_url}/parsing_lip.onnx") | |
| ] | |
| success = True | |
| for model_path, url in models: | |
| if os.path.exists(model_path): | |
| print(f"Human parsing model already exists: {model_path}") | |
| continue | |
| print(f"Downloading {model_path} from {url}") | |
| if download_model_file(model_path, [url]): | |
| print(f"Successfully downloaded: {model_path}") | |
| else: | |
| print(f"Failed to download: {model_path}") | |
| success = False | |
| return success | |
| def download_all_models(): | |
| """Download all required model files""" | |
| print("Checking and downloading required model files...") | |
| # Download DensePose model | |
| print("\n[1/3] Downloading DensePose model...") | |
| densepose_success = download_densepose_model() | |
| if densepose_success: | |
| print("β DensePose model ready") | |
| else: | |
| print("β DensePose model download failed (will download on demand)") | |
| # Download OpenPose model | |
| print("\n[2/3] Downloading OpenPose model...") | |
| openpose_success = download_openpose_model() | |
| if openpose_success: | |
| print("β OpenPose model ready") | |
| else: | |
| print("β OpenPose model download failed (will download on demand)") | |
| # Download Human Parsing models | |
| print("\n[3/3] Downloading Human Parsing models...") | |
| parsing_success = download_humanparsing_models() | |
| if parsing_success: | |
| print("β Human Parsing models ready") | |
| else: | |
| print("β Human Parsing models download failed (will download on demand)") | |
| return densepose_success and openpose_success and parsing_success | |
| def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop, denoise_steps,seed): | |
| device = "cuda" | |
| openpose_model.preprocessor.body_estimation.model.to(device) | |
| pipe.to(device) | |
| pipe.unet_encoder.to(device) | |
| garm_img= garm_img.convert("RGB").resize((768,1024)) | |
| human_img_orig = dict["background"].convert("RGB") | |
| if is_checked_crop: | |
| width, height = human_img_orig.size | |
| target_width = int(min(width, height * (3 / 4))) | |
| target_height = int(min(height, width * (4 / 3))) | |
| left = (width - target_width) / 2 | |
| top = (height - target_height) / 2 | |
| right = (width + target_width) / 2 | |
| bottom = (height + target_height) / 2 | |
| cropped_img = human_img_orig.crop((left, top, right, bottom)) | |
| crop_size = cropped_img.size | |
| human_img = cropped_img.resize((768,1024)) | |
| else: | |
| human_img = human_img_orig.resize((768,1024)) | |
| if is_checked: | |
| keypoints = openpose_model(human_img.resize((384,512))) | |
| model_parse, _ = parsing_model(human_img.resize((384,512))) | |
| mask, mask_gray = get_mask_location('hd', "upper_body", model_parse, keypoints) | |
| mask = mask.resize((768,1024)) | |
| else: | |
| mask = pil_to_binary_mask(dict['layers'][0].convert("RGB").resize((768, 1024))) | |
| # mask = transforms.ToTensor()(mask) | |
| # mask = mask.unsqueeze(0) | |
| mask_gray = (1-transforms.ToTensor()(mask)) * tensor_transfrom(human_img) | |
| mask_gray = to_pil_image((mask_gray+1.0)/2.0) | |
| human_img_arg = _apply_exif_orientation(human_img.resize((384,512))) | |
| human_img_arg = convert_PIL_to_numpy(human_img_arg, format="BGR") | |
| # DensePose λͺ¨λΈ λ€μ΄λ‘λ λ° κ²½λ‘ μ€μ | |
| densepose_model_path = './ckpt/densepose/model_final_162be9.pkl' | |
| # λͺ¨λΈ νμΌμ΄ μμΌλ©΄ λ€μ΄λ‘λ μλ | |
| if not os.path.exists(densepose_model_path): | |
| print("DensePose model not found, attempting to download...") | |
| download_success = download_densepose_model() | |
| if not download_success: | |
| print("Failed to download DensePose model") | |
| return None, None | |
| args = apply_net.create_argument_parser().parse_args(('show', './configs/densepose_rcnn_R_50_FPN_s1x.yaml', densepose_model_path, 'dp_segm', '-v', '--opts', 'MODEL.DEVICE', 'cuda')) | |
| # verbosity = getattr(args, "verbosity", None) | |
| pose_img = args.func(args,human_img_arg) | |
| pose_img = pose_img[:,:,::-1] | |
| pose_img = Image.fromarray(pose_img).resize((768,1024)) | |
| with torch.no_grad(): | |
| # Extract the images | |
| with torch.cuda.amp.autocast(): | |
| with torch.no_grad(): | |
| prompt = "model is wearing " + garment_des | |
| negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality" | |
| with torch.inference_mode(): | |
| ( | |
| prompt_embeds, | |
| negative_prompt_embeds, | |
| pooled_prompt_embeds, | |
| negative_pooled_prompt_embeds, | |
| ) = pipe.encode_prompt( | |
| prompt, | |
| num_images_per_prompt=1, | |
| do_classifier_free_guidance=True, | |
| negative_prompt=negative_prompt, | |
| ) | |
| prompt = "a photo of " + garment_des | |
| negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality" | |
| if not isinstance(prompt, List): | |
| prompt = [prompt] * 1 | |
| if not isinstance(negative_prompt, List): | |
| negative_prompt = [negative_prompt] * 1 | |
| with torch.inference_mode(): | |
| ( | |
| prompt_embeds_c, | |
| _, | |
| _, | |
| _, | |
| ) = pipe.encode_prompt( | |
| prompt, | |
| num_images_per_prompt=1, | |
| do_classifier_free_guidance=False, | |
| negative_prompt=negative_prompt, | |
| ) | |
| pose_img = tensor_transfrom(pose_img).unsqueeze(0).to(device,torch.float16) | |
| garm_tensor = tensor_transfrom(garm_img).unsqueeze(0).to(device,torch.float16) | |
| generator = torch.Generator(device).manual_seed(seed) if seed is not None else None | |
| images = pipe( | |
| prompt_embeds=prompt_embeds.to(device,torch.float16), | |
| negative_prompt_embeds=negative_prompt_embeds.to(device,torch.float16), | |
| pooled_prompt_embeds=pooled_prompt_embeds.to(device,torch.float16), | |
| negative_pooled_prompt_embeds=negative_pooled_prompt_embeds.to(device,torch.float16), | |
| num_inference_steps=denoise_steps, | |
| generator=generator, | |
| strength = 1.0, | |
| pose_img = pose_img.to(device,torch.float16), | |
| text_embeds_cloth=prompt_embeds_c.to(device,torch.float16), | |
| cloth = garm_tensor.to(device,torch.float16), | |
| mask_image=mask, | |
| image=human_img, | |
| height=1024, | |
| width=768, | |
| ip_adapter_image = garm_img.resize((768,1024)), | |
| guidance_scale=2.0, | |
| )[0] | |
| if is_checked_crop: | |
| out_img = images[0].resize(crop_size) | |
| human_img_orig.paste(out_img, (int(left), int(top))) | |
| return human_img_orig, mask_gray | |
| else: | |
| return images[0], mask_gray | |
| # return images[0], mask_gray | |
| print("\n" + "=" * 60) | |
| print("Loading Example Images...") | |
| print("=" * 60) | |
| garm_list = os.listdir(os.path.join(example_path,"cloth")) | |
| garm_list_path = [os.path.join(example_path,"cloth",garm) for garm in garm_list] | |
| print(f"β Found {len(garm_list_path)} garment example images") | |
| human_list = os.listdir(os.path.join(example_path,"human")) | |
| human_list_path = [os.path.join(example_path,"human",human) for human in human_list] | |
| print(f"β Found {len(human_list_path)} human example images") | |
| # human_ex_listλ₯Ό λ¨μν μ΄λ―Έμ§ κ²½λ‘ λ¦¬μ€νΈλ‘ λ³κ²½ (그리λ νμλ₯Ό μν΄) | |
| human_ex_list = human_list_path | |
| ##default human | |
| print("\n" + "=" * 60) | |
| print("Creating Gradio Application Interface...") | |
| print("=" * 60) | |
| image_blocks = gr.Blocks().queue() | |
| with image_blocks as demo: | |
| print("β Gradio Blocks created") | |
| gr.Markdown("## DXCO : GENAI-VTON") | |
| gr.Markdown("μμ±λ¨, μ€μ§μ, μ‘°λ―Όμ£Ό based on IDM-VTON") | |
| gr.Markdown("* 맨 μ²μ μΆλ‘ μ [5λΆ] κ±Έλ¦Ό - compileκ³Ό GPU warm-up *") | |
| gr.Markdown("κΆμ₯ μ΄λ―Έμ§ μ¬μ΄μ¦ - 3:4λΉμ¨(384x512,768x1024)") | |
| with gr.Row(): | |
| with gr.Column(): | |
| imgs = gr.ImageEditor(sources='upload', type="pil", label='λμ μ΄λ―Έμ§', interactive=True) | |
| with gr.Row(): | |
| img_url_input = gr.Textbox(label="λμ μ΄λ―Έμ§ URL", placeholder="μ) https://example.com/human_image.jpg") | |
| with gr.Row(): | |
| is_checked = gr.Checkbox(label="Yes", info="μλ λ§μ€νΉ",value=True) | |
| with gr.Row(): | |
| is_checked_crop = gr.Checkbox(label="Yes", info="μλ ν¬λ‘ λ° λ¦¬μ¬μ΄μ§",value=True) | |
| example = gr.Examples( | |
| inputs=imgs, | |
| examples_per_page=10, | |
| examples=human_ex_list | |
| ) | |
| with gr.Column(): | |
| garm_img = gr.Image(label="μμ μ΄λ―Έμ§", sources='upload', type="pil") | |
| with gr.Row(): | |
| garm_url_input = gr.Textbox(label="μμ μ΄λ―Έμ§ URL", placeholder="μ) https://example.com/garment.jpg") | |
| with gr.Row(elem_id="prompt-container"): | |
| with gr.Row(): | |
| prompt = gr.Textbox(placeholder="Description of garment ex) Short Sleeve Round Neck T-shirts", show_label=False, elem_id="prompt") | |
| example = gr.Examples( | |
| inputs=garm_img, | |
| examples_per_page=8, | |
| examples=garm_list_path) | |
| with gr.Column(): | |
| masked_img = gr.Image(label="Masked image output", elem_id="masked-img", show_share_button=False) | |
| with gr.Column(): | |
| image_out = gr.Image(label="Output", elem_id="output-img", show_share_button=False) | |
| with gr.Column(): | |
| try_button = gr.Button(value="Try-on") | |
| with gr.Accordion(label="Advanced Settings", open=False): | |
| with gr.Row(): | |
| denoise_steps = gr.Number(label="Denoising Steps", minimum=20, maximum=40, value=30, step=1) | |
| seed = gr.Number(label="Seed", minimum=-1, maximum=2147483647, step=1, value=42) | |
| # is_checked = gr.Number(value=True) | |
| # μ΄λ―Έμ§ μ λ‘λ μ μ μ²λ¦¬ | |
| # imgs.upload( | |
| # fn=preprocess_image, | |
| # inputs=imgs, | |
| # outputs=imgs, # μ μ²λ¦¬λ μ΄λ―Έμ§λ₯Ό ImageEditorμ λ€μ νμ | |
| # ) | |
| # λμ μ΄λ―Έμ§: URL μ λ ₯ μ²λ¦¬ | |
| img_url_input.change( | |
| fn=lambda url: process_url_image(url), | |
| inputs=img_url_input, | |
| outputs=imgs, | |
| ) | |
| # μμ μ΄λ―Έμ§: URL μ λ ₯ μ²λ¦¬ | |
| garm_url_input.change( | |
| fn=lambda url: process_url_image(url), | |
| inputs=garm_url_input, | |
| outputs=garm_img, | |
| ) | |
| try_button.click( | |
| fn=start_tryon, | |
| inputs=[imgs, garm_img, prompt, is_checked, is_checked_crop, denoise_steps, seed], | |
| outputs=[image_out, masked_img], | |
| api_name='tryon' | |
| ) | |
| # GPU Warm-up μν νμμ© (μ¨κΉ) | |
| warmup_status = gr.Textbox(visible=False) | |
| # μ± λ‘λ μ GPU Warm-up μλ μ€ν (torch.compile 첫 μ»΄νμΌ) | |
| if is_compile_for_zeroGPU == True: | |
| print("β GPU warm-up is disabled for ZeroGPU") | |
| else: | |
| demo.load( | |
| fn=warmup_gpu, | |
| inputs=None, | |
| outputs=warmup_status, | |
| ) | |
| print("β Gradio interface components created") | |
| print("β Event handlers configured") | |
| print("β GPU warm-up scheduled on app load") | |
| print("\n" + "=" * 60) | |
| print("Gradio Application Interface Created Successfully!") | |
| print("=" * 60) | |
| # DensePose λͺ¨λΈ λ€μ΄λ‘λ | |
| print("\n" + "=" * 60) | |
| print("Checking and Downloading Additional Models...") | |
| print("=" * 60) | |
| try: | |
| download_all_models() | |
| print("\nβ All model files downloaded successfully.") | |
| except Exception as e: | |
| print(f"\nβ Warning: Could not download all model files: {e}") | |
| print("The models will be downloaded when needed during inference.") | |
| # μ± μ€ν | |
| print("\n" + "=" * 60) | |
| print("Launching Application Server...") | |
| print("=" * 60) | |
| if __name__ == "__main__": | |
| try: | |
| print("Starting GENAI-VTON application on http://0.0.0.0:7860") | |
| print("Please wait while the server starts...") | |
| image_blocks.launch(server_name="0.0.0.0", server_port=7860, share=False) | |
| except Exception as e: | |
| print(f"\nβ Error starting the application: {e}") | |
| print("Please check if all required dependencies are installed.") | |
| import traceback | |
| traceback.print_exc() | |