new app.py with pose transfer jugaar

app.py

@@ -10,6 +10,29 @@ import threading
 import requests
 from huggingface_hub import snapshot_download
 import gradio as gr
+import base64
+import io
+import json
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
+import uuid
+
+# Import pose transfer related modules
+from diffusers import DDPMScheduler
+from defaults import pose_transfer_C as cfg
+from pose_transfer_train import build_model
+from models import UNet, VariationalAutoencoder
+import torch
+import numpy as np
+import pandas as pd
+from pose_utils import (cords_to_map, draw_pose_from_cords,
+                        load_pose_cords_from_strings)
+import random
+from PIL import Image
+from torchvision import transforms
+import copy
 
 # ==============================
 # CONFIGURATION
@@ -23,6 +46,13 @@ PERSISTENT_DIR = "/data/models"
 download_thread = None
 download_log = []
 cancel_download = False
+model_ready = False
+vae = None
+model = None
+unet = None
+noise_scheduler = None
+test_pairs = None
+annotation_file = None
 
 # ==============================
 # UTILS
@@ -113,7 +143,7 @@ def download_individual_folders_with_retry():
 
 def download_models():
     """Download models with logging and show directory tree after completion."""
-    global cancel_download
+    global cancel_download, model_ready, vae, model, unet, noise_scheduler, test_pairs, annotation_file
     
     try:
         # Test connectivity first
@@ -128,6 +158,8 @@ def download_models():
             log("📂 Directory structure:")
             tree = get_directory_tree(PERSISTENT_DIR)
             log(f"\n{tree}")
+            model_ready = True
+            initialize_models()
             return
 
         os.makedirs(PERSISTENT_DIR, exist_ok=True)
@@ -167,6 +199,8 @@ def download_models():
         log("📂 Listing downloaded directory structure...")
         tree = get_directory_tree(PERSISTENT_DIR)
         log(f"\n{tree}")
+        model_ready = True
+        initialize_models()
 
     except Exception as e:
         error_msg = str(e)
@@ -174,8 +208,108 @@ def download_models():
         log("💡 Trying alternative download method...")
         download_individual_folders_with_retry()
 
-# ... rest of your code remains the same ...
+def initialize_models():
+    """Initialize the pose transfer models after download."""
+    global vae, model, unet, noise_scheduler, test_pairs, annotation_file
+    
+    try:
+        log("🔄 Initializing pose transfer models...")
+        
+        # Initialize models
+        noise_scheduler = DDPMScheduler.from_pretrained(os.path.join(PERSISTENT_DIR, "pretrained_models/scheduler/scheduler_config.json"))
+        vae = VariationalAutoencoder(pretrained_path=os.path.join(PERSISTENT_DIR, "pretrained_models/vae")).eval().requires_grad_(False).cuda()
+        model = build_model(cfg).eval().requires_grad_(False).cuda()
+        unet = UNet(cfg).eval().requires_grad_(False).cuda()
+        
+        # Load model weights
+        model.load_state_dict(torch.load(
+            os.path.join(PERSISTENT_DIR, "checkpoints", "pytorch_model.bin"), map_location="cpu"
+        ), strict=False)
+        unet.load_state_dict(torch.load(
+            os.path.join(PERSISTENT_DIR, "checkpoints", "pytorch_model_1.bin"), map_location="cpu"
+        ), strict=False)
+        
+        # Load test data
+        test_pairs_path = os.path.join(PERSISTENT_DIR, "fashion", "fasion-resize-pairs-test.csv")
+        test_pairs = pd.read_csv(test_pairs_path)
+        
+        annotation_file_path = os.path.join(PERSISTENT_DIR, "fashion", "fasion-resize-annotation-test.csv")
+        annotation_file = pd.read_csv(annotation_file_path, sep=':')
+        annotation_file = annotation_file.set_index('name')
+        
+        log("✅ Models initialized successfully")
+        
+    except Exception as e:
+        log(f"❌ Error initializing models: {str(e)}")
+
+def build_pose_img(annotation_file, img_path):
+    """Build pose image from annotation file."""
+    string = annotation_file.loc[os.path.basename(img_path)]
+    array = load_pose_cords_from_strings(string['keypoints_y'], string['keypoints_x'])
+    pose_map = torch.tensor(cords_to_map(array, (256, 256), (256, 176)).transpose(2, 0, 1), dtype=torch.float32)
+    pose_img = torch.tensor(draw_pose_from_cords(array, (256, 256), (256, 176)).transpose(2, 0, 1) / 255., dtype=torch.float32)
+    pose_img = torch.cat([pose_img, pose_map], dim=0)
+    return pose_img
+
+def pose_transfer(source_image, test_pair_index):
+    """Perform pose transfer from source image to target pose."""
+    global vae, model, unet, noise_scheduler, test_pairs, annotation_file
+    
+    if not model_ready:
+        raise ValueError("Models not ready. Please download models first.")
+    
+    if test_pair_index < 0 or test_pair_index >= len(test_pairs):
+        raise ValueError(f"Test pair index must be between 0 and {len(test_pairs)-1}")
+    
+    # Get target image path
+    img_to_path = test_pairs.iloc[test_pair_index]["to"]
+    
+    # Build pose image
+    pose_img_tensor = build_pose_img(annotation_file, img_to_path).unsqueeze(0)
+    
+    # Transform source image
+    trans = transforms.Compose([
+        transforms.Resize([256, 256], interpolation=transforms.InterpolationMode.BICUBIC, antialias=True),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+    ])
+    img_from_tensor = trans(source_image).unsqueeze(0)
+    
+    # Perform pose transfer
+    with torch.no_grad():
+        c_new, down_block_additional_residuals, up_block_additional_residuals = model({
+            "img_cond": img_from_tensor.cuda(), "pose_img": pose_img_tensor.cuda()})
+        noisy_latents = torch.randn((1, 4, 64, 64)).cuda()
+        weight_dtype = torch.float32
+        bsz = 1
+
+        c_new = torch.cat([c_new[:bsz], c_new[:bsz], c_new[bsz:]])
+        down_block_additional_residuals = [torch.cat([torch.zeros_like(sample), sample, sample]).to(dtype=weight_dtype) \
+                                            for sample in down_block_additional_residuals]
+        up_block_additional_residuals = {k: torch.cat([torch.zeros_like(v), torch.zeros_like(v), v]).to(dtype=weight_dtype) \
+                                            for k, v in up_block_additional_residuals.items()}
+
+        noise_scheduler.set_timesteps(cfg.TEST.NUM_INFERENCE_STEPS)
+        for t in noise_scheduler.timesteps:
+            inputs = torch.cat([noisy_latents, noisy_latents, noisy_latents], dim=0)
+            inputs = noise_scheduler.scale_model_input(inputs, timestep=t)
+            noise_pred = unet(sample=inputs, timestep=t, encoder_hidden_states=c_new,
+                down_block_additional_residuals=copy.deepcopy(down_block_additional_residuals),
+                up_block_additional_residuals=copy.deepcopy(up_block_additional_residuals))
 
+            noise_pred_uc, noise_pred_down, noise_pred_full = noise_pred.chunk(3)
+            noise_pred = noise_pred_uc + \
+                            cfg.TEST.DOWN_BLOCK_GUIDANCE_SCALE * (noise_pred_down - noise_pred_uc) + \
+                            cfg.TEST.FULL_GUIDANCE_SCALE * (noise_pred_full - noise_pred_down)
+            noisy_latents = noise_scheduler.step(noise_pred, t, noisy_latents)[0]
+
+        sampling_imgs = vae.decode(noisy_latents) * 0.5 + 0.5  # denormalize
+        sampling_imgs = sampling_imgs.clamp(0, 1)
+    
+    # Convert to PIL image
+    output_img = Image.fromarray((sampling_imgs[0] * 255.).permute((1, 2, 0)).long().cpu().numpy().astype(np.uint8)).resize((256, 256))
+    
+    return output_img
 
 def start_download():
     """Start model download in a separate thread."""
@@ -206,34 +340,156 @@ def cancel_download_fn():
     log("⛔ Download cancelled by user.")
     return "Download cancelled."
 
+# ==============================
+# API Models
+# ==============================
+class PoseTransferRequest(BaseModel):
+    source_image: str  # base64 encoded image
+    test_pair_index: int  # between 0 and 4039
+
+class PoseTransferResponse(BaseModel):
+    output_image: str  # base64 encoded output image
+    success: bool
+    message: str
+
+# ==============================
+# FastAPI App
+# ==============================
+api_app = FastAPI(title="Pose Transfer API")
+
+# Add CORS middleware
+api_app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+@api_app.post("/api/pose-transfer", response_model=PoseTransferResponse)
+async def api_pose_transfer(request: PoseTransferRequest):
+    """API endpoint for pose transfer."""
+    try:
+        if not model_ready:
+            return PoseTransferResponse(
+                output_image="",
+                success=False,
+                message="Models not ready. Please download models first."
+            )
+        
+        # Decode base64 image
+        image_data = base64.b64decode(request.source_image)
+        source_image = Image.open(io.BytesIO(image_data)).convert("RGB")
+        
+        # Perform pose transfer
+        output_image = pose_transfer(source_image, request.test_pair_index)
+        
+        # Convert output to base64
+        buffered = io.BytesIO()
+        output_image.save(buffered, format="PNG")
+        output_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        
+        return PoseTransferResponse(
+            output_image=output_base64,
+            success=True,
+            message="Pose transfer successful"
+        )
+        
+    except Exception as e:
+        return PoseTransferResponse(
+            output_image="",
+            success=False,
+            message=f"Error during pose transfer: {str(e)}"
+        )
+
+@api_app.get("/api/status")
+async def api_status():
+    """API endpoint to check model status."""
+    return {
+        "model_ready": model_ready,
+        "test_pairs_count": len(test_pairs) if test_pairs is not None else 0
+    }
+
 # ==============================
 # GRADIO UI
 # ==============================
+def gradio_pose_transfer(source_image, test_pair_index):
+    """Gradio interface for pose transfer."""
+    try:
+        if not model_ready:
+            return None, "Models not ready. Please download models first."
+        
+        if test_pair_index < 0 or test_pair_index >= len(test_pairs):
+            return None, f"Test pair index must be between 0 and {len(test_pairs)-1}"
+        
+        # Perform pose transfer
+        output_image = pose_transfer(source_image, test_pair_index)
+        
+        return output_image, "Pose transfer successful"
+        
+    except Exception as e:
+        return None, f"Error during pose transfer: {str(e)}"
+
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧩 Model Downloader with Live Status & Persistent Storage")
+    gr.Markdown("## 🧩 Model Downloader & Pose Transfer")
     gr.Markdown(f"**Model Source:** [{MODEL_REPO}](https://huggingface.co/{MODEL_REPO})")
     gr.Markdown("**Required folders:** checkpoints/, pretrained_models/, fashion/")
 
-    with gr.Row():
-        start_btn = gr.Button("📥 Download Models")
-        cancel_btn = gr.Button("❌ Cancel Download")
+    with gr.Tab("Model Download"):
+        with gr.Row():
+            start_btn = gr.Button("📥 Download Models")
+            cancel_btn = gr.Button("❌ Cancel Download")
 
-    status_box = gr.Textbox(
-        label="Download Logs",
-        lines=25,
-        interactive=False,
-        placeholder="Click 'Download Models' to start downloading..."
-    )
+        status_box = gr.Textbox(
+            label="Download Logs",
+            lines=25,
+            interactive=False,
+            placeholder="Click 'Download Models' to start downloading..."
+        )
 
-    # Button bindings
-    start_btn.click(fn=start_download, inputs=None, outputs=status_box)
-    cancel_btn.click(fn=cancel_download_fn, inputs=None, outputs=status_box)
+        # Button bindings
+        start_btn.click(fn=start_download, inputs=None, outputs=status_box)
+        cancel_btn.click(fn=cancel_download_fn, inputs=None, outputs=status_box)
 
-    # Periodic refresh of logs
-    demo.load(fn=get_download_status, inputs=None, outputs=status_box, every=2)
+        # Periodic refresh of logs
+        demo.load(fn=get_download_status, inputs=None, outputs=status_box, every=2)
+
+    with gr.Tab("Pose Transfer"):
+        gr.Markdown("## Pose Transfer")
+        gr.Markdown(f"Available test pairs: {len(test_pairs) if test_pairs is not None else 'Loading...'}")
+        
+        with gr.Row():
+            with gr.Column():
+                source_image = gr.Image(label="Source Image", type="pil")
+                test_pair_index = gr.Number(
+                    label="Test Pair Index (0-4039)", 
+                    value=0, 
+                    minimum=0, 
+                    maximum=4039 if test_pairs is not None else 0
+                )
+                generate_btn = gr.Button("🚀 Generate Pose Transfer")
+            
+            with gr.Column():
+                output_image = gr.Image(label="Output Image")
+                status_message = gr.Textbox(label="Status", interactive=False)
+        
+        generate_btn.click(
+            fn=gradio_pose_transfer,
+            inputs=[source_image, test_pair_index],
+            outputs=[output_image, status_message]
+        )
 
 # ==============================
-# START APP
+# START APPS
 # ==============================
 if __name__ == "__main__":
-    demo.launch()
+    # Initialize models if they exist
+    if is_model_ready():
+        model_ready = True
+        initialize_models()
+    
+    # Mount the Gradio app
+    api_app.mount("/", gr.routes.App.create_app(demo))
+    
+    import uvicorn
+    uvicorn.run(api_app, host="0.0.0.0", port=7860)

download.py

@@ -0,0 +1,239 @@
+import os
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
+os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
+os.environ["HTTP_PROXY"] = ""
+os.environ["HTTPS_PROXY"] = ""
+os.environ["http_proxy"] = ""
+os.environ["https_proxy"] = ""
+
+import threading
+import requests
+from huggingface_hub import snapshot_download
+import gradio as gr
+
+# ==============================
+# CONFIGURATION
+# ==============================
+MODEL_REPO = "recky101/new_l_cfld_model"
+PERSISTENT_DIR = "/data/models"
+
+# ==============================
+# GLOBALS
+# ==============================
+download_thread = None
+download_log = []
+cancel_download = False
+
+# ==============================
+# UTILS
+# ==============================
+def log(msg: str):
+    """Append a log message and return the full log as a string."""
+    global download_log
+    download_log.append(msg)
+    print(msg)
+    return "\n".join(download_log)
+
+def test_huggingface_connectivity():
+    """Test if we can reach Hugging Face servers."""
+    try:
+        response = requests.get("https://huggingface.co/", timeout=10)
+        if response.status_code == 200:
+            return True, "✅ Can reach Hugging Face"
+        else:
+            return False, f"❌ Hugging Face returned status {response.status_code}"
+    except Exception as e:
+        return False, f"❌ Cannot reach Hugging Face: {str(e)}"
+
+def is_model_ready():
+    """Check if required folders already exist."""
+    checkpoints = os.path.join(PERSISTENT_DIR, "checkpoints")
+    pretrained = os.path.join(PERSISTENT_DIR, "pretrained_models")
+    fashion = os.path.join(PERSISTENT_DIR, "fashion")
+    return all(os.path.exists(path) for path in [checkpoints, pretrained, fashion])
+
+def get_directory_tree(root_dir, indent=""):
+    """Recursively build a directory tree string for logs."""
+    tree_str = ""
+    try:
+        items = sorted(os.listdir(root_dir))
+    except Exception as e:
+        return f"{indent}❌ [Error accessing {root_dir}]: {e}\n"
+
+    for i, item in enumerate(items):
+        path = os.path.join(root_dir, item)
+        connector = "└── " if i == len(items) - 1 else "├── "
+        tree_str += f"{indent}{connector}{item}\n"
+        if os.path.isdir(path):
+            tree_str += get_directory_tree(path, indent + ("    " if i == len(items) - 1 else "│   "))
+    return tree_str
+
+def download_individual_folders_with_retry():
+    """Alternative method with robust retry logic."""
+    try:
+        folders_to_download = ["checkpoints", "pretrained_models", "fashion"]
+        max_retries = 3
+        
+        for folder in folders_to_download:
+            if cancel_download:
+                log("⛔ Download cancelled during individual folder download.")
+                return
+                
+            folder_path = os.path.join(PERSISTENT_DIR, folder)
+            if os.path.exists(folder_path):
+                log(f"✅ Folder {folder} already exists, skipping...")
+                continue
+                
+            for attempt in range(max_retries):
+                try:
+                    log(f"⬇️ Downloading folder: {folder} (Attempt {attempt + 1}/{max_retries})")
+                    
+                    snapshot_download(
+                        repo_id=MODEL_REPO,
+                        local_dir=folder_path,
+                        resume_download=True,
+                        local_dir_use_symlinks=False,
+                        allow_patterns=f"{folder}/*"
+                    )
+                    log(f"✅ Successfully downloaded {folder}")
+                    break
+                    
+                except Exception as e:
+                    if attempt == max_retries - 1:
+                        log(f"❌ Failed to download {folder} after {max_retries} attempts: {str(e)}")
+                    else:
+                        log(f"⚠️ Attempt {attempt + 1} failed for {folder}: {str(e)}")
+                        import time
+                        time.sleep(5)
+            
+        log("✅ All folders processed.")
+        
+    except Exception as e:
+        log(f"❌ Individual folder download failed: {str(e)}")
+
+def download_models():
+    """Download models with logging and show directory tree after completion."""
+    global cancel_download
+    
+    try:
+        # Test connectivity first
+        success, message = test_huggingface_connectivity()
+        log(message)
+        if not success:
+            log("🌐 Please check your internet connection and try again")
+            return
+
+        if is_model_ready():
+            log("✅ Models already downloaded. Skipping...")
+            log("📂 Directory structure:")
+            tree = get_directory_tree(PERSISTENT_DIR)
+            log(f"\n{tree}")
+            return
+
+        os.makedirs(PERSISTENT_DIR, exist_ok=True)
+
+        log("⬇️ Starting model download from Hugging Face Hub...")
+        
+        # Add retry logic
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                if cancel_download:
+                    log("⛔ Download cancelled during process.")
+                    return
+                    
+                log(f"🔄 Attempt {attempt + 1}/{max_retries}")
+                
+                snapshot_download(
+                    repo_id=MODEL_REPO,
+                    local_dir=PERSISTENT_DIR,
+                    resume_download=True,
+                    local_dir_use_symlinks=False,
+                )
+                break
+                
+            except Exception as e:
+                if attempt == max_retries - 1:
+                    raise e
+                log(f"⚠️ Attempt {attempt + 1} failed: {str(e)}")
+                import time
+                time.sleep(10)
+        
+        if cancel_download:
+            log("⛔ Download cancelled during process.")
+            return
+            
+        log("✅ Download completed successfully.")
+        log("📂 Listing downloaded directory structure...")
+        tree = get_directory_tree(PERSISTENT_DIR)
+        log(f"\n{tree}")
+
+    except Exception as e:
+        error_msg = str(e)
+        log(f"❌ Download failed after {max_retries} attempts: {error_msg}")
+        log("💡 Trying alternative download method...")
+        download_individual_folders_with_retry()
+
+# ... rest of your code remains the same ...
+
+
+def start_download():
+    """Start model download in a separate thread."""
+    global download_thread, download_log, cancel_download
+    if download_thread and download_thread.is_alive():
+        return "⚠️ Download already running..."
+
+    download_log = []
+    cancel_download = False
+    download_thread = threading.Thread(target=download_models)
+    download_thread.start()
+    return "📥 Download started..."
+
+def get_download_status():
+    """Get the latest log status."""
+    global download_thread
+    status = "\n".join(download_log) if download_log else "Preparing download..."
+    if download_thread and download_thread.is_alive():
+        return status + "\n\n⏳ Download in progress..."
+    elif is_model_ready():
+        return status + "\n\n✅ Models are ready."
+    return status
+
+def cancel_download_fn():
+    """Cancel request handler."""
+    global cancel_download
+    cancel_download = True
+    log("⛔ Download cancelled by user.")
+    return "Download cancelled."
+
+# ==============================
+# GRADIO UI
+# ==============================
+with gr.Blocks() as demo:
+    gr.Markdown("## 🧩 Model Downloader with Live Status & Persistent Storage")
+    gr.Markdown(f"**Model Source:** [{MODEL_REPO}](https://huggingface.co/{MODEL_REPO})")
+    gr.Markdown("**Required folders:** checkpoints/, pretrained_models/, fashion/")
+
+    with gr.Row():
+        start_btn = gr.Button("📥 Download Models")
+        cancel_btn = gr.Button("❌ Cancel Download")
+
+    status_box = gr.Textbox(
+        label="Download Logs",
+        lines=25,
+        interactive=False,
+        placeholder="Click 'Download Models' to start downloading..."
+    )
+
+    # Button bindings
+    start_btn.click(fn=start_download, inputs=None, outputs=status_box)
+    cancel_btn.click(fn=cancel_download_fn, inputs=None, outputs=status_box)
+
+    # Periodic refresh of logs
+    demo.load(fn=get_download_status, inputs=None, outputs=status_box, every=2)
+
+# ==============================
+# START APP
+# ==============================
+if __name__ == "__main__":
+    demo.launch()