Create main.py

main.py

@@ -0,0 +1,75 @@
+from fastapi import FastAPI, Request, Form
+from fastapi.responses import HTMLResponse
+from fastapi.staticfiles import StaticFiles
+from fastapi.templating import Jinja2Templates
+import requests, json
+
+# FastAPI setup
+app = FastAPI()
+# Serve static files (if needed)
+app.mount("/static", StaticFiles(directory="static"), name="static")
+# Templates directory
+templates = Jinja2Templates(directory="templates")
+
+# External API config
+BASE_URL = "https://black-forest-labs-flux-1-dev.hf.space/gradio_api/call/infer"
+HEADERS = {"Content-Type": "application/json"}
+DEFAULT_PAYLOAD = {
+    "data": ["", 0, True, 1024, 1024, 4, 60]
+}
+
+# Helper: request an event_id for given prompt
+def get_event_id(prompt: str) -> str:
+    payload = DEFAULT_PAYLOAD.copy()
+    # insert prompt
+    payload["data"][0] = prompt
+    resp = requests.post(BASE_URL, headers=HEADERS, json=payload)
+    resp.raise_for_status()
+    data = resp.json()
+    event_id = data.get("event_id")
+    if not event_id:
+        raise RuntimeError(f"No event_id returned: {data}")
+    return event_id
+
+# Helper: stream SSE until 'complete' event, then extract URL
+def stream_until_complete(event_id: str) -> str:
+    url = f"{BASE_URL}/{event_id}"
+    with requests.get(url, headers=HEADERS, stream=True) as resp:
+        resp.raise_for_status()
+        buffer = ""
+        for chunk in resp.iter_content(chunk_size=None, decode_unicode=True):
+            buffer += chunk
+            while "\n\n" in buffer:
+                message, buffer = buffer.split("\n\n", 1)
+                evt = None
+                payload = None
+                for line in message.splitlines():
+                    if line.startswith("event:"):
+                        evt = line.split("event:",1)[1].strip()
+                    elif line.startswith("data:"):
+                        payload = line.split("data:",1)[1].strip()
+                if evt == "complete" and payload:
+                    parsed = json.loads(payload)
+                    file_info = parsed[0]
+                    return file_info.get("url")
+    raise RuntimeError("Stream ended without complete event")
+
+# Home page: form for prompt
+@app.get("/", response_class=HTMLResponse)
+async def home(request: Request):
+    # no-cache headers
+    return templates.TemplateResponse("index.html", {"request": request})
+
+# Generate endpoint: processes form and renders result
+@app.post("/generate", response_class=HTMLResponse)
+async def generate(request: Request, prompt: str = Form(...)):
+    try:
+        event_id = get_event_id(prompt)
+        image_url = stream_until_complete(event_id)
+        return templates.TemplateResponse(
+            "index.html", {"request": request, "image_url": image_url, "prompt": prompt}
+        )
+    except Exception as e:
+        return templates.TemplateResponse(
+            "index.html", {"request": request, "error": str(e), "prompt": prompt}
+        )