Update app/main.py

app/main.py

@@ -1,8 +1,19 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Query
+from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from gradio_client import Client
 
 app = FastAPI()
+
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
 client = Client("Efficient-Large-Model/SanaSprint")
 
 class GenerationRequest(BaseModel):
@@ -15,6 +26,33 @@ class GenerationRequest(BaseModel):
     guidance_scale: float = 4.5
     num_inference_steps: int = 2
 
+@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
+):
+    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"
+        )
+        return {"result": result}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
 @app.post("/generate")
 async def generate_image(request: GenerationRequest):
     try: