Update app/main.py

app/main.py

@@ -1,60 +1,30 @@
 from fastapi import FastAPI, HTTPException, Query
+from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
 from gradio_client import Client
+import requests
+import io
 
 app = FastAPI()
 
-# Enable CORS for browser access
+# Allow CORS for all origins (you can restrict this in production)
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Adjust as needed for security
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Initialize the Gradio Client
+# Initialize Gradio client
 client = Client("K00B404/flux_666")
 
-# Request body schema for POST requests
-class GenerationRequest(BaseModel):
-    prompt: str
-    basemodel: str = "black-forest-labs/FLUX.1-schnell"
-    width: int = 1280
-    height: int = 768
-    scales: int = 8
-    steps: int = 8
-    seed: int = -1
-    upscale_factor: str = "2"
-    process_upscale: bool = False
-    lora_model: str = "XLabs-AI/flux-RealismLora"
-    process_lora: bool = False
-
-@app.post("/generate")
-async def generate_image(request: GenerationRequest):
-    try:
-        result = client.predict(
-            prompt=request.prompt,
-            basemodel=request.basemodel,
-            width=request.width,
-            height=request.height,
-            scales=request.scales,
-            steps=request.steps,
-            seed=request.seed,
-            upscale_factor=request.upscale_factor,
-            process_upscale=request.process_upscale,
-            lora_model=request.lora_model,
-            process_lora=request.process_lora,
-            api_name="/gen"
-        )
-        return {"result": result}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+@app.get("/")
+def root():
+    return {"message": "Welcome to the Flux 666 Image Generator API!"}
 
-# Optional: GET endpoint for browser address bar access
-@app.get("/gen")
-async def generate_image_get(
+@app.get("/generate_image")
+def generate_image(
     prompt: str = Query(..., description="Prompt for image generation"),
     basemodel: str = "black-forest-labs/FLUX.1-schnell",
     width: int = 1280,
@@ -68,7 +38,8 @@ async def generate_image_get(
     process_lora: bool = False
 ):
     try:
-        result = client.predict(
+        # Call the Gradio prediction API
+        image_url = client.predict(
             prompt=prompt,
             basemodel=basemodel,
             width=width,
@@ -82,6 +53,13 @@ async def generate_image_get(
             process_lora=process_lora,
             api_name="/gen"
         )
-        return {"result": result}
+
+        # Download the image
+        response = requests.get(image_url)
+        response.raise_for_status()
+
+        # Return the image stream to the browser
+        return StreamingResponse(io.BytesIO(response.content), media_type="image/png")
+
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))