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() | |