Update app.py

app.py

@@ -1,20 +1,42 @@
-from fastapi import FastAPI, Query
-from gradio_client import Client, handle_file
+from fastapi import FastAPI, Query, HTTPException
+from gradio_client import Client
 import requests
+import os
 
 app = FastAPI()
 
+HF_MODEL = "yanze/PuLID-FLUX"
+BASE_URL = "https://yanze-pulid-flux.hf.space/file="  # Ensure this is correct
+TEMP_DIR = "/tmp"  # Folder to store temp image downloads
+
+def download_image(image_url: str) -> str:
+    """Downloads image from URL and saves it locally."""
+    try:
+        response = requests.get(image_url, stream=True)
+        if response.status_code != 200:
+            raise HTTPException(status_code=400, detail="Failed to download image")
+
+        filename = os.path.join(TEMP_DIR, "input_image.jpg")
+        with open(filename, "wb") as f:
+            for chunk in response.iter_content(1024):
+                f.write(chunk)
+        return filename
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error downloading image: {str(e)}")
+
 @app.get("/generate")
 def generate_image(prompt: str, image_url: str):
-    client = Client("yanze/PuLID-FLUX")
-    id_image = handle_file(image_url)
+    client = Client(HF_MODEL)
+    
+    # Download image from URL first
+    image_path = download_image(image_url)
+
     width, height = 1080, 1080  # Instagram format (fixed size)
-    base_url = "https://yanze-pulid-flux.hf.space/file="
 
     try:
         result = client.predict(
             prompt=prompt,
-            id_image=id_image,
+            id_image=image_path,  # File path instead of URL
             start_step=0,
             guidance=4,
             seed="-1",
@@ -29,8 +51,12 @@ def generate_image(prompt: str, image_url: str):
             api_name="/generate_image"
         )
 
+        if not result or not isinstance(result, list) or len(result) == 0:
+            raise HTTPException(status_code=500, detail="Model did not return a valid response")
+
         file_path = result[0]
-        full_url = f"{base_url}{file_path}"
+        full_url = f"{BASE_URL}{file_path}"
         return {"image_url": full_url}
+    
     except Exception as e:
-        return {"error": str(e)}
+        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")