Update app.py

app.py

@@ -1,10 +1,12 @@
 from fastapi import FastAPI, File, UploadFile, Request, HTTPException
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
 import httpx
 import os
 from sse_starlette.sse import EventSourceResponse
 import asyncio
+import urllib.request
+import io
 
 app = FastAPI(title="Virtual Try-On API", description="API to forward images and handle virtual try-on requests")
 
@@ -118,4 +120,30 @@ async def event_generator(session_hash: str):
 @app.get("/sse")
 async def sse_proxy(session_hash: str, request: Request):
     """Proxy Server-Sent Events for the virtual try-on process."""
-    return EventSourceResponse(event_generator(session_hash))
+    return EventSourceResponse(event_generator(session_hash))
+
+@app.get("/download-image")
+async def download_image(path: str):
+    """Download and stream an image from the external API."""
+    try:
+        async with httpx.AsyncClient() as client:
+            response = await client.get("https://kwai-kolors-kolors-virtual-try-on.hf.space/file="+path)
+            response.raise_for_status()
+            
+            # Create a BytesIO object for streaming
+            image_stream = io.BytesIO(response.content)
+            
+            # Return StreamingResponse with appropriate media type
+            return StreamingResponse(
+                image_stream,
+                media_type="image/webp",
+                headers={
+                    "Content-Disposition": "inline; filename=image.webp"
+                }
+            )
+    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)}")