Update app/main.py

app/main.py

@@ -2,12 +2,13 @@ from fastapi import FastAPI, HTTPException, Query
 from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 from gradio_client import Client
-import requests
+from PIL import Image
 import io
+import os
 
 app = FastAPI()
 
-# Allow CORS for all origins (you can restrict this in production)
+# Allow CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -16,14 +17,14 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Initialize Gradio client
+# Connect to the model
 client = Client("K00B404/flux_666")
 
 @app.get("/")
 def root():
-    return {"message": "Welcome to the Image Generator API!"}
+    return {"message": "Welcome to the Flux 666 Image Generator API!"}
 
-@app.get("/generate_image")
+@app.get("/gen")
 def generate_image(
     prompt: str = Query(..., description="Prompt for image generation"),
     basemodel: str = "black-forest-labs/FLUX.1-schnell",
@@ -38,8 +39,7 @@ def generate_image(
     process_lora: bool = False
 ):
     try:
-        # Call the Gradio prediction API
-        image_url = client.predict(
+        result = client.predict(
             prompt=prompt,
             basemodel=basemodel,
             width=width,
@@ -54,12 +54,17 @@ def generate_image(
             api_name="/gen"
         )
 
-        # Download the image
-        response = requests.get(image_url)
-        response.raise_for_status()
+        # result is a list of local file paths
+        if isinstance(result, list) and result:
+            image_path = result[0]
+        else:
+            raise ValueError("No image returned from model.")
 
-        # Return the image stream to the browser
-        return StreamingResponse(io.BytesIO(response.content), media_type="image/png")
+        # Read the image from disk
+        with open(image_path, "rb") as img_file:
+            img_bytes = img_file.read()
+
+        return StreamingResponse(io.BytesIO(img_bytes), media_type="image/png")
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))