Async InvSR: POST returns job_id instantly, poll for result (fixes Swagger timeout)

app.py

@@ -73,6 +73,8 @@ _INVSR_PATH = Path("/app/InvSR")
 _sampler_invsr = None
 _invsr_status = "not_loaded"
 _invsr_error = None
+_invsr_jobs: dict[str, dict] = {}
+_job_counter = 0
 
 def _patch_invsr_source():
     p = _INVSR_PATH / "sampler_invsr.py"
@@ -261,10 +263,33 @@ async def route_metrics(file: UploadFile = File(...)):
 
 @app.post("/upscale/invsr")
 async def route_invsr(file: UploadFile = File(...), num_steps: int = Query(1, ge=1, le=5)):
+    if _invsr_status == "error":
+        raise HTTPException(500, f"InvSR not loaded: {_invsr_error}")
+    if _sampler_invsr is None:
+        raise HTTPException(503, f"InvSR is {_invsr_status}. Check /health for status.")
+    global _job_counter
+    _job_counter += 1
+    job_id = str(_job_counter)
+    _invsr_jobs[job_id] = {"status": "queued", "image_bytes": await file.read(), "num_steps": num_steps}
+    threading.Thread(target=_run_invsr_job, args=(job_id,), daemon=True).start()
+    return JSONResponse({"job_id": job_id, "status": "queued", "check": f"/upscale/invsr/{job_id}"})
+
+def _run_invsr_job(job_id: str):
+    job = _invsr_jobs.get(job_id)
+    if not job: return
     try:
-        result = upscale_invsr(await file.read(), num_steps=num_steps)
-    except HTTPException:
-        raise
+        job["status"] = "processing"
+        job["result"] = upscale_invsr(job["image_bytes"], job["num_steps"])
+        job["status"] = "done"
     except Exception as e:
-        raise HTTPException(500, f"InvSR failed: {e}")
-    return StreamingResponse(BytesIO(result), media_type="image/png")
+        job["status"] = "error"
+        job["error"] = str(e)
+
+@app.get("/upscale/invsr/{job_id}")
+async def route_invsr_status(job_id: str):
+    job = _invsr_jobs.get(job_id)
+    if not job:
+        raise HTTPException(404, "Job not found")
+    if job["status"] == "done":
+        return StreamingResponse(BytesIO(job["result"]), media_type="image/png")
+    return JSONResponse({"job_id": job_id, "status": job["status"], "error": job.get("error")})