Enhance FastAPI application to support image generation with Inference API. Updated API title, added model configuration, and implemented a new endpoint for generating images from text prompts. Retained echo functionality for testing.

app.py

@@ -1,16 +1,50 @@
+import io, os
+from typing import Optional
 from fastapi import FastAPI
+from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
+from huggingface_hub import InferenceClient
 
-app = FastAPI(title="Hello API")
+app = FastAPI(title="ImageGen API")
 
-class Prompt(BaseModel):
-    text: str
+# Choose a model that supports text->image via Inference API.
+# FLUX.1-schnell is fast; you can swap for SDXL or others later.
+MODEL_ID = os.getenv("MODEL_ID", "black-forest-labs/FLUX.1-schnell")
+HF_TOKEN = os.getenv("HF_TOKEN")  # set in Space Secrets
+client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)
+
+class GenReq(BaseModel):
+    prompt: str
+    width: int = 768
+    height: int = 768
+    seed: Optional[int] = None
+    num_inference_steps: Optional[int] = None  # some models ignore this
 
 @app.get("/")
 def root():
-    return {"ok": True, "message": "API is live"}
+    return {"ok": True, "message": "API is live", "model": MODEL_ID}
+
+@app.post("/generate")
+def generate(req: GenReq):
+    try:
+        # InferenceClient returns a PIL.Image for text_to_image
+        img = client.text_to_image(
+            prompt=req.prompt,
+            size=f"{req.width}x{req.height}",
+            seed=req.seed,
+            num_inference_steps=req.num_inference_steps
+        )
+        buf = io.BytesIO()
+        img.save(buf, format="PNG")
+        buf.seek(0)
+        return StreamingResponse(buf, media_type="image/png")
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)
+
+# Keep your original echo for testing
+class Prompt(BaseModel):
+    text: str
 
 @app.post("/echo")
 def echo(p: Prompt):
-    # do your logic here
     return {"length": len(p.text), "upper": p.text.upper()}

requirements.txt

@@ -1,2 +1,4 @@
 fastapi
 uvicorn
+huggingface_hub
+python-multipart