checking the server

app.py

@@ -12,6 +12,7 @@ os.environ["https_proxy"] = ""
 import threading
 import requests
 from huggingface_hub import snapshot_download
+import gradio as gr
 import base64
 import io
 import json
@@ -261,6 +262,8 @@ 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:
@@ -273,6 +276,7 @@ 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)
@@ -316,7 +320,7 @@ def pose_transfer(source_image, test_pair_index):
         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_lake(v), v]).to(dtype=weight_dtype)
+        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)
@@ -337,6 +341,8 @@ def pose_transfer(source_image, test_pair_index):
         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
     output_img = Image.fromarray(
         (sampling_imgs[0] * 255.)
         .permute((1, 2, 0))
@@ -346,14 +352,169 @@ 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
 
+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():
+        # tree = get_directory_tree(PERSISTENT_DIR)
+        # log(f"\n{tree}")
+        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."
+
 # ==============================
-# FASTAPI SERVER
+# 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:
+            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)
+        
+        # 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"
+
+        
+    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)}"
+
+with gr.Blocks() as demo:
+    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.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..."
+        )
+
+        # 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)
 
-from fastapi import FastAPI, Form, UploadFile, HTTPException
+        # 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")
+        
+        with gr.Row():
+            with gr.Column():
+                source_image = gr.Image(label="Source Image", type="pil")
+                test_pair_index = gr.Number(
+                    label="Test Pair Index", 
+                    value=0, 
+                    minimum=0, 
+                    maximum=4039
+                )
+                generate_btn = gr.Button("🚀 Generate Pose Transfer")
+            
+            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(
+            fn=gradio_pose_transfer,
+            inputs=[source_image, test_pair_index],
+            outputs=[output_image, status_message]
+        )
+
+# ==============================
+# START APP
+# ==============================
+
+
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+import traceback
+
+
+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)}")
+
+
+
+
+
+
+from fastapi import FastAPI, Form, UploadFile
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
@@ -372,15 +533,6 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-@app.get("/health")
-async def health_check():
-    """Health check endpoint to verify server status"""
-    return {
-        "status": "healthy",
-        "model_ready": model_ready,
-        "message": "Server is running" if model_ready else "Server running but models not ready"
-    }
-
 @app.post("/pose-transfer")
 async def pose_transfer_api(
     test_pair_index: int = Form(...),
@@ -388,9 +540,6 @@ async def pose_transfer_api(
     image_base64: str = Form(None)
 ):
     try:
-        if not model_ready:
-            raise HTTPException(status_code=503, detail="Models not ready. Please wait for models to be downloaded and initialized.")
-        
         # Get input image
         if source_image:
             image_bytes = await source_image.read()
@@ -401,13 +550,12 @@ async def pose_transfer_api(
             image_bytes = base64.b64decode(image_base64)
             image = Image.open(BytesIO(image_bytes)).convert("RGB")
         else:
-            raise HTTPException(status_code=400, detail="No image provided")
-        
-        # Validate test pair index
-        if test_pair_index < 0 or test_pair_index >= len(test_pairs):
-            raise HTTPException(status_code=400, detail=f"Test pair index must be between 0 and {len(test_pairs)-1}")
+            return JSONResponse(
+                content={"status": "error", "message": "No image provided"},
+                status_code=400
+            )
 
-        # Perform pose transfer
+        # 🔹 Call your pose transfer model
         output_image = pose_transfer(image, test_pair_index)
 
         # Convert result to base64
@@ -422,35 +570,28 @@ async def pose_transfer_api(
                 "output_image_base64": f"data:image/png;base64,{img_base64}"
             }
         )
-    except HTTPException:
-        raise
     except Exception as e:
         import traceback
         tb = traceback.format_exc()
-        print(f"Error during pose transfer: {str(e)}\n{tb}")
-        raise HTTPException(status_code=500, detail=f"Error during pose transfer: {str(e)}")
+        return JSONResponse(
+            content={"status": "error", "message": str(e), "traceback": tb},
+            status_code=500
+        )
 
-def initialize_app():
-    """Initialize models before starting the server"""
-    global model_ready
-    
-    print("Checking if models are available...")
+if __name__ == "__main__":
     if is_model_ready():
-        print("✅ Models found. Initializing...")
         model_ready = True
         initialize_models()
-        print("✅ Models initialized successfully")
-    else:
-        print("⚠️ Models not found. Starting download process...")
-        # Download models synchronously (blocking)
-        download_models()
     
-    return app
+    uvicorn.run(app, host="0.0.0.0", port=7860)
 
-if __name__ == "__main__":
-    # Initialize models before starting server
-    initialize_app()
+
+
+
+# if __name__ == "__main__":
+#     # Initialize models if they exist
+#     if is_model_ready():
+#         model_ready = True
+#         initialize_models()
     
-    # Start FastAPI server only after models are ready
-    print("Starting server on http://0.0.0.0:8000")
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+#     demo.launch(server_name="0.0.0.0", server_port=7860)