Update app/main.py

app/main.py

@@ -5,67 +5,82 @@ from gradio_client import Client
 
 app = FastAPI()
 
-
+# Enable CORS for browser access
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=["*"],  # Adjust as needed for security
+    allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
+# Initialize the Gradio Client
+client = Client("K00B404/flux_666")
 
-client = Client("Efficient-Large-Model/SanaSprint")
-
+# Request body schema for POST requests
 class GenerationRequest(BaseModel):
     prompt: str
-    model_size: str = "1.6B"
-    seed: int = 0
-    randomize_seed: bool = True
-    width: int = 1024
-    height: int = 1024
-    guidance_scale: float = 4.5
-    num_inference_steps: int = 2
+    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.get("/gen")
-def generate_image_get(
-    prompt: str = Query(..., description="Prompt for image generation"),
-    model_size: str = "1.6B",
-    seed: int = 0,
-    randomize_seed: bool = True,
-    width: int = 1024,
-    height: int = 1024,
-    guidance_scale: float = 4.5,
-    num_inference_steps: int = 2
-):
+@app.post("/generate")
+async def generate_image(request: GenerationRequest):
     try:
         result = client.predict(
-            prompt=prompt,
-            model_size=model_size,
-            seed=seed,
-            randomize_seed=randomize_seed,
-            width=width,
-            height=height,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_inference_steps,
-            api_name="/infer"
+            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.post("/generate")
-async def generate_image(request: GenerationRequest):
+# Optional: GET endpoint for browser address bar access
+@app.get("/gen")
+async def generate_image_get(
+    prompt: str = Query(..., description="Prompt for image generation"),
+    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
+):
     try:
         result = client.predict(
-            prompt=request.prompt,
-            model_size=request.model_size,
-            seed=request.seed,
-            randomize_seed=request.randomize_seed,
-            width=request.width,
-            height=request.height,
-            guidance_scale=request.guidance_scale,
-            num_inference_steps=request.num_inference_steps,
-            api_name="/infer"
+            prompt=prompt,
+            basemodel=basemodel,
+            width=width,
+            height=height,
+            scales=scales,
+            steps=steps,
+            seed=seed,
+            upscale_factor=upscale_factor,
+            process_upscale=process_upscale,
+            lora_model=lora_model,
+            process_lora=process_lora,
+            api_name="/gen"
         )
         return {"result": result}
     except Exception as e: