Update app.py

app.py

@@ -1,81 +1,66 @@
-from fastapi import FastAPI, File, UploadFile, Request
+from fastapi import FastAPI, File, UploadFile, Request, HTTPException
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 import httpx
-import io
-import uuid
 import os
 from sse_starlette.sse import EventSourceResponse
 import asyncio
 
-app = FastAPI()
+app = FastAPI(title="Virtual Try-On API", description="API to forward images and handle virtual try-on requests")
+
+# Configuration
+ALLOWED_IMAGE_TYPES = {"image/png", "image/jpeg"}
+BASE_URL = "https://kwai-kolors-kolors-virtual-try-on.hf.space"
+
+class TryOnRequest(BaseModel):
+    person_image_path: str
+    garment_image_path: str
+    session_hash: str
 
 @app.get("/")
 async def root():
-    return {"message": "FastAPI to forward image to external API"}
-
+    return {"message": "Virtual Try-On API"}
 
-@app.post("/forward-upload")
+@app.post("/upload")
 async def forward_upload(file: UploadFile = File(...)):
-    if file.content_type not in ["image/png", "image/jpeg"]:
-        return JSONResponse(status_code=400, content={"error": "Invalid file type. Only PNG or JPEG allowed."})
+    """Upload an image file to the external API."""
+    if file.content_type not in ALLOWED_IMAGE_TYPES:
+        raise HTTPException(status_code=400, detail="Invalid file type. Only PNG or JPEG allowed.")
 
-    target_url = "https://kwai-kolors-kolors-virtual-try-on.hf.space/upload?upload_id=r9ywzifblvk"
+    target_url = f"{BASE_URL}/upload"
+    headers = {"accept": "*/*",}
 
     try:
         file_content = await file.read()
         files = {"files": (file.filename, file_content, file.content_type)}
-
+        
         async with httpx.AsyncClient() as client:
-            response = await client.post(
-                target_url,
-                files=files,
-                headers={
-                    "accept": "*/*",
-                    "accept-language": "en-GB,en-US;q=0.9,en;q=0.8,gu;q=0.7",
-                    "origin": "https://kwai-kolors-kolors-virtual-try-on.hf.space",
-                    "referer": "https://kwai-kolors-kolors-virtual-try-on.hf.space/?__theme=system",
-                    "user-agent": "Mozilla/5.0 (iPhone; CPU iPhone OS 16_6 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.6 Mobile/15E148 Safari/604.1",
-                },
-            )
-
-        if response.status_code == 200:
-            return JSONResponse(content={"status": "Success", "response": response.json()})
-        else:
-            return JSONResponse(status_code=response.status_code, content={"error": response.text})
-
+            response = await client.post(target_url, files=files, headers=headers)
+            response.raise_for_status()
+            return {"status": "success", "response": response.json()}
+            
+    except httpx.HTTPStatusError as e:
+        raise HTTPException(status_code=e.response.status_code, detail=str(e.response.text))
     except httpx.HTTPError as e:
-        return JSONResponse(status_code=500, content={"error": f"HTTP error occurred: {str(e)}"})
+        raise HTTPException(status_code=500, detail=f"HTTP error occurred: {str(e)}")
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": f"An error occurred: {str(e)}"})
-
-
-class TryOnRequest(BaseModel):
-    person_image_path: str
-    garment_image_path: str
-    session_hash: str 
-
-@app.post("/virtual-try-on")
-async def join(request: TryOnRequest):
-    url = "https://kwai-kolors-kolors-virtual-try-on.hf.space/queue/join?__theme=system"
+        raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}")
 
+@app.post("join")
+async def virtual_try_on(request: TryOnRequest):
+    """Process virtual try-on request with person and garment images."""
+    url = f"{BASE_URL}/queue/join"
     headers = {
         "accept": "*/*",
-        "content-type": "application/json",
+        "content-type": "application/json"
     }
 
-    # Extract orig_name from file path
-    person_orig_name = os.path.basename(request.person_image_path)
-    garment_orig_name = os.path.basename(request.garment_image_path)
-
-    #session_hash = str(uuid.uuid4())[:11]
-
     payload = {
         "data": [
             {
                 "path": request.person_image_path,
-                "url": f"https://kwai-kolors-kolors-virtual-try-on.hf.space/file={request.person_image_path}",
-                "orig_name": person_orig_name,
+                "url": f"{BASE_URL}/file={request.person_image_path}",
+                "orig_name": os.path.basename(request.person_image_path),
                 "size": None,
                 "mime_type": None,
                 "is_stream": False,
@@ -83,8 +68,8 @@ async def join(request: TryOnRequest):
             },
             {
                 "path": request.garment_image_path,
-                "url": f"https://kwai-kolors-kolors-virtual-try-on.hf.space/file={request.garment_image_path}",
-                "orig_name": garment_orig_name,
+                "url": f"{BASE_URL}/file={request.garment_image_path}",
+                "orig_name": os.path.basename(request.garment_image_path),
                 "size": None,
                 "mime_type": None,
                 "is_stream": False,
@@ -102,31 +87,35 @@ async def join(request: TryOnRequest):
     try:
         async with httpx.AsyncClient() as client:
             response = await client.post(url, headers=headers, json=payload)
-
-        if response.status_code == 200:
+            response.raise_for_status()
             return response.json()
-        else:
-            raise HTTPException(status_code=response.status_code, detail=response.text)
-
-    except httpx.RequestError as e:
-        raise HTTPException(status_code=500, detail=f"Request failed: {str(e)}")
-
-
-async def event_generator(session_hash: str, request: Request):
-    SSE_URL = f"https://kwai-kolors-kolors-virtual-try-on.hf.space/queue/data?session_hash={session_hash}"
+            
+    except httpx.HTTPStatusError as e:
+        raise HTTPException(status_code=e.response.status_code, detail=str(e.response.text))
+    except httpx.HTTPError as e:
+        raise HTTPException(status_code=500, detail=f"HTTP error occurred: {str(e)}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}")
 
+async def event_generator(session_hash: str):
+    """Generate Server-Sent Events for the given session hash."""
+    sse_url = f"{BASE_URL}/queue/data?session_hash={session_hash}"
     headers = {
         "accept": "text/event-stream",
-        "content-type": "application/json",
+        "content-type": "application/json"
     }
 
-    async with httpx.AsyncClient(timeout=None) as client:
-        async with client.stream("GET", SSE_URL, headers=headers) as response:
-            async for line in response.aiter_lines():
-                if line.startswith("data:"):
-                    yield line.removeprefix("data:").strip()
-                await asyncio.sleep(0.01)
+    try:
+        async with httpx.AsyncClient(timeout=None) as client:
+            async with client.stream("GET", sse_url, headers=headers) as response:
+                async for line in response.aiter_lines():
+                    if line.startswith("data:"):
+                        yield line.removeprefix("data:").strip()
+                    await asyncio.sleep(0.01)
+    except Exception as e:
+        yield f"data: {{'error': 'SSE stream error: {str(e)}'}}"
 
-@app.get("/sse-proxy")
+@app.get("/sse")
 async def sse_proxy(session_hash: str, request: Request):
-    return EventSourceResponse(event_generator(session_hash, request))
+    """Proxy Server-Sent Events for the virtual try-on process."""
+    return EventSourceResponse(event_generator(session_hash))