check claude result

app.py

@@ -1,7 +1,6 @@
 import os
 # Prevent libgomp crashes in Spaces
 os.environ["OMP_NUM_THREADS"] = "1"
-
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
 os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
 os.environ["HTTP_PROXY"] = ""
@@ -17,6 +16,15 @@ import base64
 import io
 import json
 import uuid
+from fastapi import FastAPI, HTTPException, UploadFile, File
+from fastapi.responses import JSONResponse, FileResponse
+from pydantic import BaseModel, validator
+from typing import Optional, Union
+import tempfile
+import aiofiles
+import asyncio
+from urllib.parse import urlparse
+import httpx
 
 # Import pose transfer related modules
 try:
@@ -42,6 +50,9 @@ except ImportError as e:
 MODEL_REPO = "recky101/new_l_cfld_model"
 PERSISTENT_DIR = "/data/models"
 
+# Initialize FastAPI app
+app = FastAPI(title="Pose Transfer API", version="1.0.0", description="API for AI-powered pose transfer")
+
 # ==============================
 # GLOBALS
 # ==============================
@@ -57,7 +68,41 @@ test_pairs = None
 annotation_file = None
 
 # ==============================
-# UTILS
+# PYDANTIC MODELS
+# ==============================
+class PoseTransferRequest(BaseModel):
+    source_image: Optional[str] = None  # Base64 encoded image
+    source_image_url: Optional[str] = None  # URL to image
+    test_pair_index: int
+    output_format: Optional[str] = "base64"  # "base64" or "url"
+    
+    @validator('test_pair_index')
+    def validate_test_pair_index(cls, v):
+        if v < 0 or v >= 4040:  # Based on your maximum value
+            raise ValueError(f'Test pair index must be between 0 and 4039')
+        return v
+    
+    @validator('source_image', 'source_image_url')
+    def validate_image_input(cls, v, values, field):
+        # Ensure at least one image input is provided
+        if field.name == 'source_image_url' and not v and not values.get('source_image'):
+            raise ValueError('Either source_image or source_image_url must be provided')
+        return v
+
+class PoseTransferResponse(BaseModel):
+    success: bool
+    message: str
+    output_image: Optional[str] = None  # Base64 encoded or URL
+    output_format: str
+    processing_time: Optional[float] = None
+
+class ModelStatusResponse(BaseModel):
+    model_ready: bool
+    message: str
+    download_progress: Optional[str] = None
+
+# ==============================
+# UTILITY FUNCTIONS (Same as your original code)
 # ==============================
 def log(msg: str):
     """Append a log message and return the full log as a string."""
@@ -210,7 +255,6 @@ def download_models():
         log("💡 Trying alternative download method...")
         download_individual_folders_with_retry()
 
-
 def initialize_models():
     """Initialize the pose transfer models after download."""
     global vae, model, unet, noise_scheduler, test_pairs, annotation_file
@@ -229,9 +273,9 @@ def initialize_models():
                 size = os.path.getsize(os.path.join(root, f)) / (1024*1024)
                 log(f"{subindent}{f} ({size:.2f} MB)")
         
-        log("🔄 Initializing pose transfer models after tyeh checking ll thigns out...")
+        log("🔄 Initializing pose transfer models after checking all things out...")
         try:
-        # Initialize models
+            # Initialize models
             noise_scheduler = DDPMScheduler.from_pretrained(os.path.join(PERSISTENT_DIR, "pretrained_models/scheduler/scheduler_config.json"))
             log("🔄 noise done")
             vae = VariationalAutoencoder(pretrained_path=os.path.join(PERSISTENT_DIR, "pretrained_models/vae")).eval().requires_grad_(False).cuda()
@@ -244,15 +288,16 @@ def initialize_models():
         except Exception as e:
             log(f"❌ Error during model initialization: {str(e)}")
             log(e)
+        
         # Load model weights
         model.load_state_dict(torch.load(
             os.path.join(PERSISTENT_DIR, "checkpoints", "pytorch_model.bin"), map_location="cpu"
         ), strict=False)
-        log("🔄 checkoitns  done")
+        log("🔄 checkpoints done")
         unet.load_state_dict(torch.load(
             os.path.join(PERSISTENT_DIR, "checkpoints", "pytorch_model_1-001.bin"), map_location="cpu"
         ), strict=False)
-        log("🔄 checkoitns22  done ")
+        log("🔄 checkpoints22 done ")
         
         # Load test data
         test_pairs_path = os.path.join(PERSISTENT_DIR, "fashion", "fasion-resize-pairs-test.csv")
@@ -262,8 +307,6 @@ def initialize_models():
         annotation_file = pd.read_csv(annotation_file_path, sep=':')
         annotation_file = annotation_file.set_index('name')
         
-
-        
         log("✅ Models initialized successfully")
         
     except Exception as e:
@@ -276,7 +319,6 @@ def build_pose_img(annotation_file, img_path):
     log(f"📄 Index Sample: {annotation_file.index[:5]}")
     log(f"📄 Does key exist?: {os.path.basename(img_path) in annotation_file.index}")
 
-
     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)
@@ -285,7 +327,6 @@ def build_pose_img(annotation_file, img_path):
     return pose_img
 
 def pose_transfer(source_image, test_pair_index):
-    test_pair_index = int(test_pair_index)
     """Perform pose transfer from source image to target pose."""
     global vae, model, unet, noise_scheduler, test_pairs, annotation_file
     
@@ -298,6 +339,7 @@ def pose_transfer(source_image, test_pair_index):
     # Get target image path
     img_to_path = test_pairs.iloc[test_pair_index]["to"]
     log(f"🔄 img_to_path: {img_to_path}")
+    
     # Build pose image
     pose_img_tensor = build_pose_img(annotation_file, img_to_path).unsqueeze(0)
     
@@ -340,9 +382,7 @@ def pose_transfer(source_image, test_pair_index):
         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))
-    # Convert tensor to PIL image without resizing
+    # Convert tensor to PIL image
     output_img = Image.fromarray(
         (sampling_imgs[0] * 255.)
         .permute((1, 2, 0))
@@ -352,20 +392,227 @@ def pose_transfer(source_image, test_pair_index):
         .astype(np.uint8)
     )
     
-     # ✅ Save image in-memory as PNG, preserving original size
-    img_bytes = io.BytesIO()
-    output_img.save(img_bytes, format="PNG")
-    img_bytes.seek(0)
-
     log("✅ Pose transfer completed successfully")
     log(f"🔄 output_img size: {output_img.size}")
-    log(f"🔄 output_img butes: {img_bytes}")
+    
+    return output_img
 
+# ==============================
+# IMAGE PROCESSING UTILITIES
+# ==============================
+async def load_image_from_base64(base64_str: str) -> Image.Image:
+    """Load PIL Image from base64 string."""
+    try:
+        # Remove data URL prefix if present
+        if base64_str.startswith('data:image'):
+            base64_str = base64_str.split(',', 1)[1]
+        
+        # Decode base64
+        image_data = base64.b64decode(base64_str)
+        image = Image.open(io.BytesIO(image_data))
+        
+        # Convert to RGB if necessary
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+            
+        return image
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Invalid base64 image: {str(e)}")
 
+async def load_image_from_url(url: str) -> Image.Image:
+    """Load PIL Image from URL."""
+    try:
+        # Validate URL
+        parsed_url = urlparse(url)
+        if not parsed_url.scheme or not parsed_url.netloc:
+            raise HTTPException(status_code=400, detail="Invalid URL format")
+        
+        # Download image
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            response = await client.get(url)
+            response.raise_for_status()
+            
+            # Check content type
+            content_type = response.headers.get('content-type', '')
+            if not content_type.startswith('image/'):
+                raise HTTPException(status_code=400, detail="URL does not point to an image")
+            
+            # Load image
+            image = Image.open(io.BytesIO(response.content))
+            
+            # Convert to RGB if necessary
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+                
+            return image
+            
+    except httpx.HTTPError as e:
+        raise HTTPException(status_code=400, detail=f"Failed to download image: {str(e)}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing image from URL: {str(e)}")
 
+def pil_to_base64(image: Image.Image, format: str = "PNG") -> str:
+    """Convert PIL Image to base64 string."""
+    buffer = io.BytesIO()
+    image.save(buffer, format=format)
+    buffer.seek(0)
+    img_str = base64.b64encode(buffer.getvalue()).decode()
+    return f"data:image/{format.lower()};base64,{img_str}"
 
-    return output_img
+# ==============================
+# API ENDPOINTS
+# ==============================
+@app.get("/")
+async def root():
+    """Root endpoint with API information."""
+    return {
+        "message": "Pose Transfer API",
+        "version": "1.0.0",
+        "endpoints": {
+            "POST /pose-transfer": "Perform pose transfer",
+            "GET /model-status": "Check model status",
+            "POST /download-models": "Download models",
+            "GET /health": "Health check"
+        }
+    }
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint."""
+    return {
+        "status": "healthy",
+        "model_ready": model_ready,
+        "timestamp": pd.Timestamp.now().isoformat()
+    }
+
+@app.get("/model-status", response_model=ModelStatusResponse)
+async def get_model_status():
+    """Get current model status and download progress."""
+    global model_ready, download_log, download_thread
+    
+    is_downloading = download_thread and download_thread.is_alive()
+    progress = "\n".join(download_log) if download_log else None
+    
+    return ModelStatusResponse(
+        model_ready=model_ready,
+        message="Models are ready" if model_ready else ("Download in progress" if is_downloading else "Models not downloaded"),
+        download_progress=progress
+    )
+
+@app.post("/download-models")
+async def start_model_download():
+    """Start model download process."""
+    global download_thread, download_log, cancel_download
+    
+    if download_thread and download_thread.is_alive():
+        return JSONResponse(
+            status_code=409,
+            content={"message": "Download already in progress"}
+        )
+    
+    if model_ready:
+        return JSONResponse(
+            content={"message": "Models already downloaded and ready"}
+        )
+    
+    # Start download in background
+    download_log = []
+    cancel_download = False
+    download_thread = threading.Thread(target=download_models)
+    download_thread.start()
+    
+    return JSONResponse(
+        content={"message": "Model download started"}
+    )
+
+@app.post("/pose-transfer", response_model=PoseTransferResponse)
+async def api_pose_transfer(request: PoseTransferRequest):
+    """Perform pose transfer via API."""
+    import time
+    start_time = time.time()
+    
+    try:
+        # Check if models are ready
+        if not model_ready:
+            raise HTTPException(
+                status_code=503, 
+                detail="Models not ready. Please download models first using /download-models endpoint."
+            )
+        
+        # Load source image
+        source_image = None
+        if request.source_image:
+            source_image = await load_image_from_base64(request.source_image)
+        elif request.source_image_url:
+            source_image = await load_image_from_url(request.source_image_url)
+        else:
+            raise HTTPException(
+                status_code=400,
+                detail="Either source_image (base64) or source_image_url must be provided"
+            )
+        
+        # Perform pose transfer
+        try:
+            output_image = pose_transfer(source_image, request.test_pair_index)
+        except ValueError as e:
+            raise HTTPException(status_code=400, detail=str(e))
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Pose transfer failed: {str(e)}")
+        
+        # Format output based on requested format
+        processing_time = time.time() - start_time
+        
+        if request.output_format == "base64":
+            output_base64 = pil_to_base64(output_image)
+            return PoseTransferResponse(
+                success=True,
+                message="Pose transfer completed successfully",
+                output_image=output_base64,
+                output_format="base64",
+                processing_time=processing_time
+            )
+        else:
+            # For URL format, save to temporary file and return file path
+            # In production, you might want to save to cloud storage instead
+            temp_filename = f"pose_transfer_{uuid.uuid4()}.png"
+            temp_path = os.path.join(tempfile.gettempdir(), temp_filename)
+            output_image.save(temp_path, "PNG")
+            
+            return PoseTransferResponse(
+                success=True,
+                message="Pose transfer completed successfully",
+                output_image=f"/download/{temp_filename}",
+                output_format="url",
+                processing_time=processing_time
+            )
+            
+    except HTTPException:
+        raise
+    except Exception as e:
+        processing_time = time.time() - start_time
+        return PoseTransferResponse(
+            success=False,
+            message=f"An error occurred: {str(e)}",
+            processing_time=processing_time
+        )
 
+@app.get("/download/{filename}")
+async def download_file(filename: str):
+    """Download generated image files."""
+    file_path = os.path.join(tempfile.gettempdir(), filename)
+    
+    if not os.path.exists(file_path):
+        raise HTTPException(status_code=404, detail="File not found")
+    
+    return FileResponse(
+        path=file_path,
+        filename=filename,
+        media_type='image/png'
+    )
+
+# ==============================
+# GRADIO UI (Keep existing functionality)
+# ==============================
 def start_download():
     """Start model download in a separate thread."""
     global download_thread, download_log, cancel_download
@@ -385,8 +632,6 @@ def get_download_status():
     if download_thread and download_thread.is_alive():
         return status + "\n\n⏳ Download in progress..."
     elif is_model_ready():
-        # tree = get_directory_tree(PERSISTENT_DIR)
-        # log(f"\n{tree}")
         return status + "\n\n✅ Models are ready."
     return status
 
@@ -397,11 +642,7 @@ def cancel_download_fn():
     log("⛔ Download cancelled by user.")
     return "Download cancelled."
 
-# ==============================
-# GRADIO UI ONLY (NO API ENDPOINTS)
-# ==============================
 def gradio_pose_transfer(source_image, test_pair_index):
-    test_pair_index = int(test_pair_index)
     """Gradio interface for pose transfer."""
     try:
         if not model_ready:
@@ -411,24 +652,17 @@ def gradio_pose_transfer(source_image, test_pair_index):
             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)
+        output_image = pose_transfer(source_image, int(test_pair_index))
         
-        # Save image to in-memory bytes (PNG)
-        img_bytes = io.BytesIO()
-        output_image.save(img_bytes, format="PNG")
-        img_bytes.seek(0)
-
-
-        return output_image, img_bytes, "Pose transfer successful"
-
+        return output_image, "Pose transfer successful"
         
     except Exception as e:
         import traceback
         tb = traceback.format_exc()
         log(f"❌ Error during pose transfer:\n{tb}")
-        log(f"❌ the value recieved is as follow: \n{test_pair_index}")
         return None, f"Error during pose transfer: {str(e)}"
 
+# Create Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("## 🧩 Model Downloader & Pose Transfer")
     gr.Markdown(f"**Model Source:** [{MODEL_REPO}](https://huggingface.co/{MODEL_REPO})")
@@ -469,7 +703,6 @@ with gr.Blocks() as demo:
             
             with gr.Column():
                 output_image = gr.Image(label="Output Image", type="pil")
-                download_btn = gr.File(label="Download Output", file_types=[".png"])
                 status_message = gr.Textbox(label="Status", interactive=False)
         
         generate_btn.click(
@@ -478,42 +711,116 @@ with gr.Blocks() as demo:
             outputs=[output_image, status_message]
         )
 
-# ==============================
-# START APP
-# ==============================
-
-
-from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse
-import traceback
-
+    with gr.Tab("API Documentation"):
+        gr.Markdown("""
+        ## API Endpoints
+        
+        ### POST /pose-transfer
+        Perform pose transfer with JSON request:
+        ```json
+        {
+            "source_image": "base64_encoded_image_string",  // Optional
+            "source_image_url": "https://example.com/image.jpg",  // Optional
+            "test_pair_index": 42,
+            "output_format": "base64"  // "base64" or "url"
+        }
+        ```
+        
+        ### GET /model-status
+        Check if models are downloaded and ready
+        
+        ### POST /download-models
+        Start model download process
+        
+        ### GET /health
+        Health check endpoint
+        
+        ## Example Usage (Python)
+        ```python
+        import requests
+        import base64
+        
+        # Load and encode image
+        with open("source_image.jpg", "rb") as f:
+            image_base64 = base64.b64encode(f.read()).decode()
+        
+        # API request
+        response = requests.post("http://localhost:7860/pose-transfer", json={
+            "source_image": image_base64,
+            "test_pair_index": 100,
+            "output_format": "base64"
+        })
+        
+        result = response.json()
+        if result["success"]:
+            # Save output image
+            output_base64 = result["output_image"].split(",")[1]
+            with open("output_image.png", "wb") as f:
+                f.write(base64.b64decode(output_base64))
+            print(f"Processing time: {result['processing_time']:.2f}s")
+        ```
+        
+        ## Example Usage (cURL)
+        ```bash
+        # Using base64 image
+        curl -X POST "http://localhost:7860/pose-transfer" \
+        -H "Content-Type: application/json" \
+        -d '{
+            "source_image": "iVBORw0KGgoAAAANSUhEUgAA...",
+            "test_pair_index": 100,
+            "output_format": "base64"
+        }'
+        
+        # Using image URL
+        curl -X POST "http://localhost:7860/pose-transfer" \
+        -H "Content-Type: application/json" \
+        -d '{
+            "source_image_url": "https://example.com/image.jpg",
+            "test_pair_index": 100,
+            "output_format": "url"
+        }'
+# Mount FastAPI app with Gradio
+    """) 
 
-def read_image_from_input(image_input: str):
-    """
-    Accepts either:
-      - Base64 image string ("data:image/png;base64,...")
-      - Image URL ("https://...")
-    Returns: PIL.Image
-    """
-    try:
-        if image_input.startswith("http://") or image_input.startswith("https://"):
-            response = requests.get(image_input)
-            response.raise_for_status()
-            return Image.open(io.BytesIO(response.content)).convert("RGB")
-        else:
-            # Base64 input
-            if "," in image_input:
-                image_input = image_input.split(",", 1)[1]
-            image_bytes = base64.b64decode(image_input)
-            return Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    except Exception as e:
-        raise ValueError(f"Invalid image input: {str(e)}")
+        
+                       
+app.mount("/gradio", gr.mount_gradio_app(demo, app, path="/gradio"))
 
+# ==============================
+# STARTUP EVENT
+# ==============================
+@app.on_event("startup")
+async def startup_event():
+    """Initialize models on startup if available."""
+    global model_ready
+    if is_model_ready() and not model_ready:
+        log("🚀 Starting up API server...")
+        log("🔍 Models found on startup, initializing...")
+        # Run initialization in background thread to avoid blocking startup
+        init_thread = threading.Thread(target=initialize_models)
+        init_thread.start()
 
+# ==============================
+# MAIN EXECUTION
+# ==============================
 if __name__ == "__main__":
-    # Initialize models if they exist
-    if is_model_ready():
+    import uvicorn
+    
+    # Check if models are ready on startup
+    if is_model_ready() and not model_ready:
+        print("📁 Models found, initializing...")
         model_ready = True
         initialize_models()
     
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    # Run both FastAPI and Gradio
+    print("🚀 Starting Pose Transfer API server...")
+    print("📱 Gradio UI available at: http://localhost:7860/gradio")
+    print("🔌 API endpoints available at: http://localhost:7860")
+    print("📖 API documentation at: http://localhost:7860/docs")
+    
+    uvicorn.run(
+        app, 
+        host="0.0.0.0", 
+        port=7860,
+        log_level="info"
+    )