Improve model downloading and error handling; added robust download logic and enhanced memory management

app.py

@@ -1,16 +1,17 @@
 import sys
 import os
-import gc  # <--- ADDED: Garbage Collection
+import gc
+import shutil
 
 # --- 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
-from requests.adapters import HTTPAdapter
-from urllib3.util.retry import Retry
 import gradio as gr
 import spaces
 from PIL import Image
@@ -18,6 +19,7 @@ import numpy as np
 import torch
 from torchvision import transforms
 from torchvision.transforms.functional import to_pil_image
+from huggingface_hub import hf_hub_download
 
 sys.path.append('./')
 
@@ -32,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,
@@ -41,20 +43,67 @@ from transformers import (
 from diffusers import DDPMScheduler, AutoencoderKL
 
 # ---------------------------------------------------------
-# 2. DOWNLOADER
+# 2. ROBUST MODEL DOWNLOADER (The Fix)
 # ---------------------------------------------------------
-def download_file(url, path):
-    if os.path.exists(path): return
-    # ... (Keep existing downloader logic if you wish, or use the robust one from before)
-    # For brevity, assuming files exist or you use the previous robust downloader code here.
-    # If not, paste the 'download_file' function from the previous response here.
-    print(f"Checking {path}...")
-    if not os.path.exists(path):
-        os.system(f"wget -O {path} {url}")
+def download_model_robust(repo_id, filename, local_path):
+    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)
+        
+        # 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 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():
-    # ... (Same file list as before)
-    pass # Call your download logic here
+    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()
 
 # ---------------------------------------------------------
 # 3. LOAD MODELS
@@ -71,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)
@@ -94,15 +145,13 @@ pipe, openpose_model, parsing_model = load_models()
 tensor_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
 
 # ---------------------------------------------------------
-# 4. INFERENCE (Fixed Memory Leak)
+# 4. INFERENCE
 # ---------------------------------------------------------
-# Increase duration to 120s to prevent timeouts
 @spaces.GPU(duration=120)
 def start_tryon(human_img, garm_img, garment_des, is_checked, is_checked_crop, denoise_steps, seed):
     device = "cuda"
     
     try:
-        # Move models to GPU
         openpose_model.preprocessor.body_estimation.model.to(device)
         pipe.to(device)
         pipe.unet_encoder.to(device)
@@ -185,9 +234,7 @@ def start_tryon(human_img, garm_img, garment_des, is_checked, is_checked_crop, d
         raise gr.Error(f"Error: {e}")
     
     finally:
-        # --- CRITICAL MEMORY CLEANUP ---
-        # This code runs no matter what, preventing the "3-4 run crash"
-        print("Cleaning GPU memory...")
+        # Memory Cleanup
         try:
             del keypoints, model_parse, mask, pose_img, prompt_embeds, garm_tensor
         except: