Refactor README and app.py for improved clarity; updated title, emoji, and SDK version, added HF OAuth support, enhanced model downloading logic, and improved error handling.

README.md

@@ -1,13 +1,14 @@
 ---
-title: Tryonnix 2D Engine
-emoji: 👕
-colorFrom: blue
-colorTo: purple
+title: IDM VTON
+emoji: 👕👔👚
+colorFrom: yellow
+colorTo: red
 sdk: gradio
-sdk_version: 4.26.0
+sdk_version: 4.24.0
 app_file: app.py
 pinned: false
-license: other
-# ADD THIS LINE BELOW:
-hf_oauth: true
----
+license: cc-by-nc-sa-4.0
+short_description: High-fidelity Virtual Try-on
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -2,21 +2,16 @@ import sys
 import os
 import gc
 import shutil
-import huggingface_hub
 
-# --- Force Login (Optional but recommended) ---
-# If you added HF_TOKEN to secrets, this ensures the container is logged in
-token = os.getenv("HF_TOKEN")
-if token:
-    print("🔐 Authenticating with Secret Token...")
-    huggingface_hub.login(token=token)
-
-# --- Install Detectron2 if missing ---
+# --- 1. System Setup & Error Handling ---
+# Force install detectron2 if missing
 try:
     import detectron2
 except ImportError:
+    print("⚠️ Detectron2 missing. Installing...")
     os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
 
+import requests
 import gradio as gr
 import spaces
 from PIL import Image
@@ -28,7 +23,7 @@ from huggingface_hub import hf_hub_download
 
 sys.path.append('./')
 
-# --- Import Local Modules ---
+# Import Local Modules
 try:
     from utils_mask import get_mask_location
     from src.tryon_pipeline import StableDiffusionXLInpaintPipeline as TryonPipeline
@@ -39,7 +34,7 @@ try:
     from detectron2.data.detection_utils import convert_PIL_to_numpy, _apply_exif_orientation
     import apply_net
 except ImportError as e:
-    raise ImportError(f"CRITICAL ERROR: Missing core modules. {e}")
+    raise ImportError(f"CRITICAL ERROR: Missing core modules. Error: {e}")
 
 from transformers import (
     CLIPImageProcessor, CLIPVisionModelWithProjection, CLIPTextModel,
@@ -47,31 +42,72 @@ from transformers import (
 )
 from diffusers import DDPMScheduler, AutoencoderKL
 
-# --- Robust Downloader ---
+# ---------------------------------------------------------
+# 2. ROBUST MODEL DOWNLOADER (The Fix)
+# ---------------------------------------------------------
 def download_model_robust(repo_id, filename, local_path):
-    if os.path.exists(local_path) and os.path.getsize(local_path) > 1000:
-        return
+    if os.path.exists(local_path):
+        # Quick size check to ensure it's not an empty corrupt file
+        if os.path.getsize(local_path) > 1000:
+            print(f"✅ Found {local_path}")
+            return
+        else:
+            print(f"⚠️ Corrupt file found at {local_path}, redownloading...")
+            os.remove(local_path)
+    
+    print(f"⬇️ Downloading {filename} to {local_path}...")
     try:
+        # Create directory
         os.makedirs(os.path.dirname(local_path), exist_ok=True)
-        hf_hub_download(repo_id=repo_id, filename=filename, local_dir=os.path.dirname(local_path), local_dir_use_symlinks=False)
-        actual_path = os.path.join(os.path.dirname(local_path), filename)
-        if actual_path != local_path and os.path.exists(actual_path):
-            shutil.move(actual_path, local_path)
+        
+        # Download using Hugging Face Hub (Fast & Cached)
+        downloaded_file = hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            local_dir=os.path.dirname(local_path),
+            local_dir_use_symlinks=False
+        )
+        
+        # If the filename in repo is different from target, rename it
+        # (hf_hub_download saves to local_dir/filename)
+        actual_download_path = os.path.join(os.path.dirname(local_path), filename)
+        if actual_download_path != local_path:
+            # Move it to the exact expected path if different
+            if os.path.exists(actual_download_path):
+                 shutil.move(actual_download_path, local_path)
+                 
+        print(f"✅ Successfully downloaded {local_path}")
+        
     except Exception as e:
-        print(f"Failed download {filename}: {e}")
+        print(f"❌ Failed to download {filename}: {e}")
+        # Manual Fallback for complex paths
+        try:
+            url = f"https://huggingface.co/{repo_id}/resolve/main/{filename}"
+            print(f"🔄 Trying direct URL fallback: {url}")
+            os.system(f"wget -O {local_path} {url}")
+        except:
+            pass
 
 def check_and_download_models():
+    print("⏳ VALIDATING MODELS...")
+    
+    # 1. Parsing & OpenPose (From Camenduru)
     download_model_robust("camenduru/IDM-VTON", "humanparsing/parsing_atr.onnx", "ckpt/humanparsing/parsing_atr.onnx")
     download_model_robust("camenduru/IDM-VTON", "humanparsing/parsing_lip.onnx", "ckpt/humanparsing/parsing_lip.onnx")
     download_model_robust("camenduru/IDM-VTON", "densepose/model_final_162be9.pkl", "ckpt/densepose/model_final_162be9.pkl")
     download_model_robust("camenduru/IDM-VTON", "openpose/ckpts/body_pose_model.pth", "ckpt/openpose/ckpts/body_pose_model.pth")
+    
+    # 2. IP Adapter (From h94)
     download_model_robust("h94/IP-Adapter", "sdxl_models/ip-adapter-plus_sdxl_vit-h.bin", "ckpt/ip_adapter/ip-adapter-plus_sdxl_vit-h.bin")
     download_model_robust("h94/IP-Adapter", "models/image_encoder/config.json", "ckpt/image_encoder/config.json")
     download_model_robust("h94/IP-Adapter", "models/image_encoder/pytorch_model.bin", "ckpt/image_encoder/pytorch_model.bin")
 
+# EXECUTE DOWNLOAD BEFORE LOADING ANYTHING
 check_and_download_models()
 
-# --- Load Models ---
+# ---------------------------------------------------------
+# 3. LOAD MODELS
+# ---------------------------------------------------------
 base_path = 'yisol/IDM-VTON'
 def load_models():
     unet = UNet2DConditionModel.from_pretrained(base_path, subfolder="unet", torch_dtype=torch.float16)
@@ -84,9 +120,11 @@ def load_models():
     vae = AutoencoderKL.from_pretrained(base_path, subfolder="vae", torch_dtype=torch.float16)
     UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(base_path, subfolder="unet_encoder", torch_dtype=torch.float16)
     
+    # Initialize Preprocessors
     parsing_model = Parsing(0)
     openpose_model = OpenPose(0)
     
+    # Freeze Weights
     UNet_Encoder.requires_grad_(False)
     image_encoder.requires_grad_(False)
     vae.requires_grad_(False)
@@ -106,23 +144,20 @@ def load_models():
 pipe, openpose_model, parsing_model = load_models()
 tensor_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
 
-# --- Inference ---
+# ---------------------------------------------------------
+# 4. INFERENCE
+# ---------------------------------------------------------
 @spaces.GPU(duration=120)
-def start_tryon(human_img, garm_img, garment_des, is_checked, is_checked_crop, denoise_steps, seed, profile: gr.OAuthProfile | None = None):
-    # CRITICAL: If OAuth is enabled, we verify the user here
-    if profile is None:
-         # It will still run, but this log helps debug
-         print("⚠️ User running without OAuth Profile (might hit limits)")
-    else:
-         print(f"✅ Pro User Detected: {profile.name}")
-
+def start_tryon(human_img, garm_img, garment_des, is_checked, is_checked_crop, denoise_steps, seed):
     device = "cuda"
+    
     try:
         openpose_model.preprocessor.body_estimation.model.to(device)
         pipe.to(device)
         pipe.unet_encoder.to(device)
 
-        if not human_img or not garm_img: raise gr.Error("Missing images")
+        if not human_img or not garm_img:
+            raise gr.Error("Please upload both Human and Garment images.")
 
         garm_img = garm_img.convert("RGB").resize((768, 1024))
         human_img_orig = human_img.convert("RGB")
@@ -146,7 +181,7 @@ def start_tryon(human_img, garm_img, garment_des, is_checked, is_checked_crop, d
             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))
-            
+
             mask_gray = (1 - transforms.ToTensor()(mask)) * tensor_transform(human_img)
             mask_gray = to_pil_image((mask_gray + 1.0) / 2.0)
 
@@ -161,9 +196,13 @@ def start_tryon(human_img, garm_img, garment_des, is_checked, is_checked_crop, d
             negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"
             
             with torch.cuda.amp.autocast():
-                (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_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_c = "a photo of " + garment_des
-                (prompt_embeds_c, _, _, _) = pipe.encode_prompt(prompt_c, num_images_per_prompt=1, do_classifier_free_guidance=False, negative_prompt=negative_prompt)
+                (prompt_embeds_c, _, _, _) = pipe.encode_prompt(
+                    prompt_c, num_images_per_prompt=1, do_classifier_free_guidance=False, negative_prompt=negative_prompt
+                )
 
                 pose_img = tensor_transform(pose_img).unsqueeze(0).to(device, torch.float16)
                 garm_tensor = tensor_transform(garm_img).unsqueeze(0).to(device, torch.float16)
@@ -188,23 +227,26 @@ def start_tryon(human_img, garm_img, garment_des, is_checked, is_checked_crop, d
             final_result = human_img_orig
         else:
             final_result = images[0]
-            
+
         return final_result, mask_gray
 
     except Exception as e:
         raise gr.Error(f"Error: {e}")
+    
     finally:
-        try: del keypoints, model_parse, mask, pose_img, prompt_embeds, garm_tensor
-        except: pass
+        # Memory Cleanup
+        try:
+            del keypoints, model_parse, mask, pose_img, prompt_embeds, garm_tensor
+        except:
+            pass
         gc.collect()
         torch.cuda.empty_cache()
 
-# --- UI ---
+# ---------------------------------------------------------
+# 5. UI
+# ---------------------------------------------------------
 with gr.Blocks(theme=gr.themes.Soft(), title="Tryonnix Engine") as demo:
-    gr.Markdown("# ✨ Tryonnix 2D Engine")
-    # LOGIN BUTTON: This is the fix. It forces the Pro Token to be sent.
-    gr.LoginButton()
-    
+    gr.Markdown("# ✨ Tryonnix 2D Engine (Stable)")
     with gr.Row():
         with gr.Column():
             img_human = gr.Image(label="Human", type="pil", height=400)
@@ -219,7 +261,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Tryonnix Engine") as demo:
             out = gr.Image(label="Result", type="pil", height=600)
             mask_out = gr.Image(label="Mask", type="pil", visible=False)
 
-    btn.click(fn=start_tryon, inputs=[img_human, img_garm, desc, chk1, chk2, steps, seed], outputs=[out, mask_out])
+    btn.click(fn=start_tryon, inputs=[img_human, img_garm, desc, chk1, chk2, steps, seed], outputs=[out, mask_out], api_name="tryon")
 
 if __name__ == "__main__":
     demo.queue(max_size=10).launch()