some changes

app.py

@@ -1,9 +1,3 @@
-# common
-
-
-# ====================================>GPT
-
-
 import os
 import sys
 import torch
@@ -21,23 +15,25 @@ import io
 from fastapi import FastAPI, HTTPException
 import uvicorn
 from pydantic import BaseModel
+import threading
+import logging
+import requests
+from pathlib import Path
 
 # ===========================
-# ENVIRONMENT SETUP
+# LOGGING SETUP
 # ===========================
-# Disable telemetry and set environment
-
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
+# ===========================
+# ENVIRONMENT SETUP
+# ===========================
 os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'
 os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
 os.environ['HF_HOME'] = '/tmp/huggingface_cache'
-
-# max_threads = min(multiprocessing.cpu_count(), 2)
-# os.environ["OMP_NUM_THREADS"] = str(max_threads)
-# os.environ["MKL_NUM_THREADS"] = str(max_threads)
-# os.environ["OPENBLAS_NUM_THREADS"] = str(max_threads)
-
-
+# Disable hf_transfer for more reliable downloads
+os.environ['HF_HUB_ENABLE_HF_TRANSFER'] = '0'
 
 print("🚀 Starting CFLD Pose Transfer Application...")
 print(f"Python version: {sys.version}")
@@ -48,7 +44,7 @@ print(f"CUDA available: {torch.cuda.is_available()}")
 # IMPORTS (with error handling)
 # ===========================
 try:
-    from huggingface_hub import snapshot_download
+    from huggingface_hub import snapshot_download, hf_hub_download
     from diffusers import DDPMScheduler
     print("✅ Core dependencies imported successfully")
 except ImportError as e:
@@ -80,6 +76,8 @@ class ModelState:
         self.model_dir = None
         self.is_loaded = False
         self.is_downloaded = False
+        self.download_progress = ""
+        self.load_progress = ""
         
     def reset(self):
         """Reset model state for memory management"""
@@ -138,123 +136,245 @@ def build_pose_img(annotation_file, img_path, device='cuda'):
         print(f"Error building pose image: {e}")
         raise
 
+def download_swin_pretrained():
+    """Download Swin transformer pretrained weights"""
+    swin_dir = "pretrained_models/swin"
+    os.makedirs(swin_dir, exist_ok=True)
+    
+    swin_path = os.path.join(swin_dir, "swin_base_patch4_window12_384_22kto1k.pth")
+    
+    if os.path.exists(swin_path):
+        print("✅ Swin pretrained model already exists")
+        return True
+    
+    try:
+        print("📥 Downloading Swin pretrained model...")
+        # Alternative download URLs for Swin transformer
+        swin_urls = [
+            "https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window12_384_22kto1k.pth",
+            "https://download.pytorch.org/models/swin_b-68c6b09e.pth"
+        ]
+        
+        for url in swin_urls:
+            try:
+                response = requests.get(url, stream=True)
+                response.raise_for_status()
+                
+                with open(swin_path, 'wb') as f:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                
+                print(f"✅ Downloaded Swin model from {url}")
+                return True
+                
+            except Exception as e:
+                print(f"❌ Failed to download from {url}: {e}")
+                continue
+        
+        print("❌ All Swin download URLs failed")
+        return False
+        
+    except Exception as e:
+        print(f"❌ Error downloading Swin model: {e}")
+        return False
+
 # ===========================
-# MODEL DOWNLOADING (No GPU required)
+# MODEL DOWNLOADING (Enhanced)
 # ===========================
 def download_models():
-    """Download models from Hugging Face Hub (no GPU required)"""
+    """Download models from Hugging Face Hub with better error handling"""
     global model_state
     
     if model_state.is_downloaded:
-        print("✅ Models already downloaded")
-        return True
+        model_state.download_progress = "✅ Models already downloaded"
+        return "✅ Models already downloaded"
         
     try:
+        model_state.download_progress = "⏳ Starting model download..."
         print("⏳ Downloading models & data from repository...")
+        
         repo_id = "recky101/new_l_cfld_model"
-        model_state.model_dir = snapshot_download(
-            repo_id=repo_id,
-            cache_dir="/tmp/model_cache",
-            resume_download=True
-        )
+        cache_dir = "/tmp/model_cache"
+        
+        # Download with retries and better error handling
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                model_state.download_progress = f"⏳ Download attempt {attempt + 1}/{max_retries}..."
+                model_state.model_dir = snapshot_download(
+                    repo_id=repo_id,
+                    cache_dir=cache_dir,
+                    resume_download=True,
+                    local_files_only=False,
+                    force_download=False
+                )
+                break
+            except Exception as e:
+                print(f"❌ Download attempt {attempt + 1} failed: {e}")
+                if attempt == max_retries - 1:
+                    raise e
+                continue
+        
+        model_state.download_progress = f"📁 Downloaded to: {model_state.model_dir}"
         print(f"📁 Downloaded to: {model_state.model_dir}")
         
+        # Download additional Swin pretrained model
+        model_state.download_progress = "📥 Downloading Swin pretrained model..."
+        swin_success = download_swin_pretrained()
+        if not swin_success:
+            print("⚠️ Warning: Swin model download failed, but continuing...")
+        
         # Load dataset (doesn't require GPU)
+        model_state.download_progress = "📊 Loading fashion dataset..."
         print("📊 Loading fashion dataset...")
-        model_state.test_pairs = pd.read_csv(
-            os.path.join(model_state.model_dir, "fashion", "fasion-resize-pairs-test.csv")
-        )
-        model_state.annotation_file = pd.read_csv(
-            os.path.join(model_state.model_dir, "fashion", "fasion-resize-annotation-test.csv"), 
-            sep=":"
-        )
-        model_state.annotation_file = model_state.annotation_file.set_index("name")
+        
+        try:
+            model_state.test_pairs = pd.read_csv(
+                os.path.join(model_state.model_dir, "fashion", "fasion-resize-pairs-test.csv")
+            )
+            model_state.annotation_file = pd.read_csv(
+                os.path.join(model_state.model_dir, "fashion", "fasion-resize-annotation-test.csv"), 
+                sep=":"
+            )
+            model_state.annotation_file = model_state.annotation_file.set_index("name")
+        except Exception as e:
+            print(f"❌ Error loading dataset: {e}")
+            raise e
         
         model_state.is_downloaded = True
+        model_state.download_progress = f"✅ All models downloaded successfully! Loaded {len(model_state.test_pairs)} test pairs"
+        
         print("✅ All models downloaded successfully!")
         print(f"📈 Loaded {len(model_state.test_pairs)} test pairs")
         
-        return True
+        # Print directory structure for debugging
+        print_directory_structure()
+        
+        return model_state.download_progress
             
     except Exception as e:
-        print(f"❌ Error downloading models: {e}")
+        error_msg = f"❌ Error downloading models: {str(e)}"
+        model_state.download_progress = error_msg
+        print(error_msg)
         traceback.print_exc()
-        return False
+        return error_msg
+
+def print_directory_structure():
+    """Print directory structure for debugging"""
+    if not model_state.model_dir:
+        return
+    
+    print("\n📂 Downloaded directory structure:")
+    try:
+        for root, dirs, files in os.walk(model_state.model_dir):
+            level = root.replace(model_state.model_dir, '').count(os.sep)
+            indent = ' ' * 2 * level
+            print(f"{indent}{os.path.basename(root)}/")
+            subindent = ' ' * 2 * (level + 1)
+            for file in files[:5]:  # Limit to first 5 files per directory
+                print(f"{subindent}{file}")
+            if len(files) > 5:
+                print(f"{subindent}... and {len(files) - 5} more files")
+    except Exception as e:
+        print(f"Error listing directory: {e}")
 
 # ===========================
-# MODEL LOADING (GPU required)
+# MODEL LOADING (Enhanced)
 # ===========================
 def load_models_gpu():
     """Load models on GPU with proper memory management"""
     global model_state
     
     if model_state.is_loaded:
-        return True
+        model_state.load_progress = "✅ Models already loaded"
+        return "✅ Models already loaded"
         
     # Ensure models are downloaded first
     if not model_state.is_downloaded:
-        success = download_models()
-        if not success:
-            return False
+        model_state.load_progress = "📥 Downloading models first..."
+        download_result = download_models()
+        if "❌" in download_result:
+            model_state.load_progress = "❌ Failed to download models"
+            return "❌ Failed to download models"
     
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    model_state.load_progress = f"🔧 Loading models on device: {device}..."
     print(f"🔧 Loading models on device: {device}")
     
     try:
         with gpu_memory_guard():
             # Load scheduler
+            model_state.load_progress = "🔧 Loading scheduler..."
             print("🔧 Loading scheduler...")
             model_state.noise_scheduler = DDPMScheduler.from_pretrained(
                 os.path.join(model_state.model_dir, "pretrained_models/scheduler")
             )
 
             # Load VAE
+            model_state.load_progress = "🔧 Loading VAE..."
             print("🔧 Loading VAE...")
             model_state.vae = VariationalAutoencoder(
                 pretrained_path=os.path.join(model_state.model_dir, "pretrained_models/vae")
             ).eval().requires_grad_(False).to(device)
 
             # Load main model
+            model_state.load_progress = "🔧 Loading main model..."
             print("🔧 Loading main model...")
-            model_state.model = build_model(cfg).eval().requires_grad_(False).to(device)
+            try:
+                model_state.model = build_model(cfg).eval().requires_grad_(False).to(device)
+            except FileNotFoundError as e:
+                if "swin_base_patch4_window12_384_22kto1k.pth" in str(e):
+                    print("⚠️ Swin pretrained model not found, attempting to download...")
+                    swin_success = download_swin_pretrained()
+                    if swin_success:
+                        model_state.model = build_model(cfg).eval().requires_grad_(False).to(device)
+                    else:
+                        raise e
+                else:
+                    raise e
 
             # Load UNet
+            model_state.load_progress = "🔧 Loading UNet..."
             print("🔧 Loading UNet...")
             model_state.unet = UNet(cfg).eval().requires_grad_(False).to(device)
 
             # Load weights
+            model_state.load_progress = "📦 Loading model weights..."
             print("📦 Loading model weights...")
-            model_weights = torch.load(
-                os.path.join(model_state.model_dir, "checkpoints/pytorch_model.bin"), 
-                map_location=device
-            )
+            
+            model_weights_path = os.path.join(model_state.model_dir, "checkpoints/pytorch_model.bin")
+            unet_weights_path = os.path.join(model_state.model_dir, "checkpoints/pytorch_model_1.bin")
+            
+            if not os.path.exists(model_weights_path):
+                raise FileNotFoundError(f"Model weights not found: {model_weights_path}")
+            if not os.path.exists(unet_weights_path):
+                raise FileNotFoundError(f"UNet weights not found: {unet_weights_path}")
+            
+            model_weights = torch.load(model_weights_path, map_location=device)
             model_state.model.load_state_dict(model_weights, strict=False)
             del model_weights  # Free memory
             
-            unet_weights = torch.load(
-                os.path.join(model_state.model_dir, "checkpoints/pytorch_model_1.bin"), 
-                map_location=device
-            )
+            unet_weights = torch.load(unet_weights_path, map_location=device)
             model_state.unet.load_state_dict(unet_weights, strict=False)
             del unet_weights  # Free memory
 
             model_state.is_loaded = True
+            model_state.load_progress = "✅ All models loaded successfully!"
             print("✅ All models loaded successfully!")
             
-            return True
+            return model_state.load_progress
             
     except Exception as e:
-        print(f"❌ Error loading models: {e}")
+        error_msg = f"❌ Error loading models: {str(e)}"
+        model_state.load_progress = error_msg
+        print(error_msg)
         traceback.print_exc()
         model_state.reset()
-        return False
+        return error_msg
 
 # ===========================
 # INFERENCE FUNCTION
 # ===========================
-# ===========================
-# INFERENCE FUNCTION (Modified)
-# ===========================
 def perform_inference(img_from_array, pair_index=None):
     """Perform pose transfer inference using test pair reference"""
     global model_state
@@ -265,9 +385,9 @@ def perform_inference(img_from_array, pair_index=None):
         with gpu_memory_guard():
             # Ensure models are loaded
             if not model_state.is_loaded:
-                success = load_models_gpu()
-                if not success:
-                    return None, "Failed to load models", None, None
+                load_result = load_models_gpu()
+                if "❌" in load_result:
+                    return None, f"Failed to load models: {load_result}", None, None
 
             # Convert numpy array back to PIL Image
             img_from = Image.fromarray(img_from_array.astype(np.uint8))
@@ -306,7 +426,7 @@ def perform_inference(img_from_array, pair_index=None):
 
             print("🚀 Running inference...")
             
-            # Main inference (same as before)
+            # Main inference
             with torch.no_grad():
                 c_new, down_block_additional_residuals, up_block_additional_residuals = model_state.model({
                     "img_cond": img_from_tensor, 
@@ -358,20 +478,20 @@ def perform_inference(img_from_array, pair_index=None):
                 output_array = (sampling_imgs[0] * 255.).permute((1, 2, 0)).cpu().numpy().astype(np.uint8)
                 
                 print("✅ Inference completed successfully!")
-                return output_array, f"Success! Used test pair: {pair_index}", ref_image, img_to_path
+                return output_array, f"✅ Success! Used test pair: {pair_index}", ref_image, img_to_path
                 
     except Exception as e:
         print(f"❌ Error in inference: {e}")
         traceback.print_exc()
-        return None, f"Error: {str(e)}", None, None
+        return None, f"❌ Error: {str(e)}", None, None
 
 # ===========================
-# GRADIO INTERFACE FUNCTIONS (Modified)
+# GRADIO INTERFACE FUNCTIONS
 # ===========================
 def gradio_inference(img_from, pair_index):
     """Gradio-compatible inference function"""
     if img_from is None:
-        return None, "❌ Please upload an image first!", None, None
+        return None, "❌ Please upload an image first!", None, ""
     
     try:
         # Convert PIL to numpy for GPU function
@@ -379,7 +499,7 @@ def gradio_inference(img_from, pair_index):
         result_array, message, ref_image, ref_path = perform_inference(img_array, pair_index)
         
         if result_array is None:
-            return None, message, None, None
+            return None, message, None, ""
         
         # Convert back to PIL for gradio display
         result_image = Image.fromarray(result_array)
@@ -390,99 +510,67 @@ def gradio_inference(img_from, pair_index):
         error_msg = f"❌ Inference failed: {str(e)}"
         print(error_msg)
         traceback.print_exc()
-        return None, error_msg, None, None
-
-# ===========================
-# GRADIO INTERFACE (Modified)
-# ===========================
-
+        return None, error_msg, None, ""
 
 def check_model_status():
-    """Check if models are loaded"""
+    """Check model status and return appropriate message"""
     if model_state.is_loaded:
-        return "✅ **Status:** Models loaded and ready!"
+        return "✅ **Status:** Models loaded and ready for inference!"
     elif model_state.is_downloaded:
         return "🔄 **Status:** Models downloaded, ready to load on first inference"
+    elif model_state.download_progress:
+        return f"📥 **Status:** {model_state.download_progress}"
     else:
-        return "📥 **Status:** Models not downloaded. Click 'Download Models' first!"
+        return "📥 **Status:** Ready to download models"
 
-# ===========================
-# API FUNCTIONS
-# ===========================
-class PoseTransferRequest(BaseModel):
-    image_base64: str
-    pose_index: int = -1  # -1 for random
-
-class PoseTransferResponse(BaseModel):
-    success: bool
-    message: str
-    image_base64: str = None
-    pose_index: int = -1
-
-def base64_to_pil(image_base64):
-    """Convert base64 image to PIL Image"""
+def preview_test_pair(pair_index):
+    """Preview the test pair without running inference"""
     try:
-        if image_base64.startswith('data:image'):
-            # Remove data URL prefix if present
-            image_base64 = image_base64.split(',', 1)[1]
+        if not model_state.is_downloaded:
+            return None, "❌ Download models first!", ""
+        
+        if pair_index is None:
+            pair_index = 0
+            
+        pair_index = min(int(pair_index), len(model_state.test_pairs) - 1)
+        pair = model_state.test_pairs.iloc[pair_index]
+        img_to_path = pair["to"]
         
-        image_data = base64.b64decode(image_base64)
-        image = Image.open(io.BytesIO(image_data))
-        return image.convert('RGB')
+        ref_image_path = os.path.join(model_state.model_dir, "fashion", "test_highres", img_to_path)
+        if os.path.exists(ref_image_path):
+            ref_image = Image.open(ref_image_path).convert("RGB")
+            return ref_image, f"✅ Preview: Test pair {pair_index}", img_to_path
+        else:
+            return None, f"❌ Reference image not found: {img_to_path}", ""
+            
     except Exception as e:
-        raise ValueError(f"Invalid image data: {str(e)}")
-
-def pil_to_base64(image):
-    """Convert PIL Image to base64"""
-    buffered = io.BytesIO()
-    image.save(buffered, format="JPEG")
-    img_str = base64.b64encode(buffered.getvalue()).decode()
-    return f"data:image/jpeg;base64,{img_str}"
+        return None, f"❌ Preview error: {str(e)}", ""
 
-def api_inference(image_base64, pose_index=-1):
-    """API inference function"""
+# ===========================
+# AUTO-DOWNLOAD FUNCTION
+# ===========================
+def auto_download_models():
+    """Automatically download models on startup"""
     try:
-        # Convert base64 to PIL image
-        img_from = base64_to_pil(image_base64)
-        
-        # Convert PIL to numpy for inference function
-        img_array = np.array(img_from)
-        result_array, message, _ = perform_inference(img_array, pose_index)
-        
-        if result_array is None:
-            return PoseTransferResponse(
-                success=False,
-                message=message,
-                pose_index=pose_index
-            )
-        
-        # Convert result to base64
-        result_image = Image.fromarray(result_array)
-        result_base64 = pil_to_base64(result_image)
-        
-        return PoseTransferResponse(
-            success=True,
-            message=message,
-            image_base64=result_base64,
-            pose_index=pose_index
-        )
+        print("🚀 Auto-downloading models on startup...")
+        download_models()
         
+        # Auto-load models if GPU is available
+        if torch.cuda.is_available():
+            print("🚀 Auto-loading models (GPU detected)...")
+            load_models_gpu()
+        else:
+            print("⚠️ No GPU detected, models will be loaded on first inference")
+            
     except Exception as e:
-        error_msg = f"Inference failed: {str(e)}"
-        print(error_msg)
-        traceback.print_exc()
-        return PoseTransferResponse(
-            success=False,
-            message=error_msg,
-            pose_index=pose_index
-        )
+        print(f"❌ Auto-download failed: {e}")
 
 # ===========================
 # GRADIO INTERFACE
 # ===========================
 def create_interface():
     with gr.Blocks(
-        title="🎭 CFLD Pose Transfer - Professional Demo",
+        title="🎭 CFLD Pose Transfer - Auto-Loading Demo",
         theme=gr.themes.Soft(),
         css="""
         .gradio-container {
@@ -503,9 +591,9 @@ def create_interface():
     ) as demo:
         
         gr.Markdown("""
-        # 🎭 CFLD Pose Transfer - Professional Demo
+        # 🎭 CFLD Pose Transfer - Auto-Loading Demo
         
-        **Upload source image + select reference pose from test pairs**
+        **Models download and load automatically! Just upload an image and generate.**
         
         ---
         """)
@@ -514,21 +602,10 @@ def create_interface():
         with gr.Row():
             with gr.Column():
                 status_display = gr.Markdown(
-                    "📥 **Status:** Ready to download models...", 
+                    "🚀 Starting up... Models will download automatically", 
                     elem_classes=["status-box"]
                 )
         
-        # Model management buttons
-        with gr.Row():
-            download_btn = gr.Button(
-                "📥 Download Models (No GPU)", 
-                variant="secondary"
-            )
-            load_btn = gr.Button(
-                "🔧 Load Models (GPU Required)", 
-                variant="secondary"
-            )
-        
         with gr.Row(equal_height=True):
             # Input column
             with gr.Column(scale=1):
@@ -589,7 +666,8 @@ def create_interface():
                 
                 ref_path_display = gr.Textbox(
                     label="Reference Image Path",
-                    interactive=False
+                    interactive=False,
+                    value=""
                 )
                 
                 gr.Markdown("""
@@ -598,40 +676,7 @@ def create_interface():
                 The pose is selected based on the test pair index you provide.
                 """)
 
-        # Event handlers
-        download_btn.click(
-            fn=download_models,
-            outputs=[status_display],
-            show_progress=True
-        )
-        
-        load_btn.click(
-            fn=load_models_gpu,
-            outputs=[status_display],
-            show_progress=True
-        )
-        
-        # Preview test pair
-        def preview_test_pair(pair_index):
-            """Preview the test pair without running inference"""
-            try:
-                if not model_state.is_downloaded:
-                    return None, "Download models first!", None
-                
-                pair_index = min(int(pair_index), len(model_state.test_pairs) - 1)
-                pair = model_state.test_pairs.iloc[pair_index]
-                img_to_path = pair["to"]
-                
-                ref_image_path = os.path.join(model_state.model_dir, "fashion", "test_highres", img_to_path)
-                if os.path.exists(ref_image_path):
-                    ref_image = Image.open(ref_image_path).convert("RGB")
-                    return ref_image, f"Preview: Test pair {pair_index}", img_to_path
-                else:
-                    return None, f"Reference image not found: {img_to_path}", None
-                    
-            except Exception as e:
-                return None, f"Preview error: {str(e)}", None
-        
+        # Event handlers        
         preview_btn.click(
             fn=preview_test_pair,
             inputs=[pair_index],
@@ -650,41 +695,27 @@ def create_interface():
         demo.load(
             fn=check_model_status,
             outputs=[status_display],
-            every=5
+            every=3
         )
         
         # Footer
         gr.Markdown("""
         ---
-        **How it works:** 
-        1. Upload your source image
-        2. Select a test pair index (0-4499) for reference pose
-        3. Click generate to transfer the pose
-        4. Use preview to see the reference pose first
+        **Fully Automatic Setup:** 
+        - ✅ Models download automatically on startup
+        - ✅ Models load automatically on first use
+        - ✅ No manual intervention required
+        - ✅ Full directory structure printed to console for debugging
+        
+        **How to use:** 
+        1. Wait for models to download (automatic)
+        2. Upload your source image
+        3. Select a test pair index (0-4499) for reference pose
+        4. Click generate to transfer the pose
         """)
 
     return demo
 
-# ===========================
-# FASTAPI SETUP
-# ===========================
-app = FastAPI(title="CFLD Pose Transfer API")
-
-@app.post("/api/pose-transfer", response_model=PoseTransferResponse)
-async def api_pose_transfer(request: PoseTransferRequest):
-    """API endpoint for pose transfer"""
-    return api_inference(request.image_base64, request.pose_index)
-
-@app.get("/health")
-async def health_check():
-    """Health check endpoint"""
-    return {
-        "status": "healthy",
-        "models_downloaded": model_state.is_downloaded,
-        "models_loaded": model_state.is_loaded,
-        "cuda_available": torch.cuda.is_available()
-    }
-
 # ===========================
 # MAIN EXECUTION
 # ===========================
@@ -692,33 +723,15 @@ if __name__ == "__main__":
     print("🌟 Creating Gradio interface...")
     demo = create_interface()
     
-    # Download models on startup (no GPU required)
-    print("📥 Downloading models on startup...")
-    download_models()
+    # Start auto-download in background thread
+    download_thread = threading.Thread(target=auto_download_models, daemon=True)
+    download_thread.start()
     
     print("🚀 Launching application...")
-    # Run both Gradio and FastAPI
-    # Note: In production, you might want to run these separately
-    import threading
-    
-    def run_gradio():
-        demo.launch(
-            server_name="0.0.0.0",
-            server_port=7860,
-            show_error=True,
-            share=False,
-            debug=False
-        )
-    
-    def run_api():
-        uvicorn.run(app, host="0.0.0.0", port=8000)
-    
-    # Start both servers in separate threads
-    gradio_thread = threading.Thread(target=run_gradio)
-    api_thread = threading.Thread(target=run_api)
-    
-    gradio_thread.start()
-    api_thread.start()
-    
-    gradio_thread.join()
-    api_thread.join()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+        share=False,
+        debug=False
+    )