feat: add async job API with concurrency control (/get-hairswap-async, /job/{id})

Hair_stable_new_fresh/server.py

@@ -8,6 +8,8 @@ from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, Header
 from fastapi.responses import FileResponse, JSONResponse
 from pydantic import BaseModel
 import torch
+import asyncio
+import time
 
 import numpy as np
 from PIL import Image
@@ -254,3 +256,82 @@ def logs(_=Depends(verify_bearer)):
     return JSONResponse({"logs": ["service running"], "db": "not_configured"})
 
 
+# -------------------- Async job API --------------------
+MAX_CONCURRENCY = int(os.environ.get("MAX_CONCURRENCY", "1"))
+_sem = asyncio.Semaphore(MAX_CONCURRENCY)
+_jobs = {}
+
+
+@app.post("/get-hairswap-async")
+async def get_hairswap_async(req: HairSwapRequest, _=Depends(verify_bearer)):
+    job_id = str(uuid.uuid4())
+    _jobs[job_id] = {"status": "queued", "result": None, "error": None, "started_at": None, "ended_at": None}
+
+    async def _run_job():
+        _jobs[job_id]["status"] = "running"
+        _jobs[job_id]["started_at"] = time.time()
+        try:
+            async with _sem:
+                # reuse the same core flow as sync endpoint
+                def find_file(image_id: str) -> str:
+                    for name in os.listdir(UPLOAD_DIR):
+                        if name.startswith(image_id):
+                            return os.path.join(UPLOAD_DIR, name)
+                    raise HTTPException(status_code=404, detail=f"Image id not found: {image_id}")
+
+                source_path = find_file(req.source_id)
+                reference_path = find_file(req.reference_id)
+                LOGGER.info(f"[job {job_id}] Found source: {source_path}, reference: {reference_path}")
+
+                model = get_model()
+                LOGGER.info(f"[job {job_id}] Model loaded successfully")
+
+                LOGGER.info(f"[job {job_id}] Starting hair transfer...")
+                try:
+                    try:
+                        sched_main = type(model.pipeline.scheduler).__name__ if hasattr(model, "pipeline") else None
+                        sched_bald = type(model.remove_hair_pipeline.scheduler).__name__ if hasattr(model, "remove_hair_pipeline") else None
+                        LOGGER.info(f"[job {job_id}] Schedulers -> main: {sched_main}, remove_hair: {sched_bald}")
+                    except Exception:
+                        pass
+
+                    id_np, out_np, bald_np, ref_np = model.Hair_Transfer(
+                        source_image=source_path,
+                        reference_image=reference_path,
+                        random_seed=-1,
+                        step=20,
+                        guidance_scale=req.guidance_scale,
+                        scale=req.scale,
+                        controlnet_conditioning_scale=req.controlnet_conditioning_scale,
+                        size=448,
+                    )
+                    result_id = str(uuid.uuid4())
+                    out_img = Image.fromarray((out_np * 255.).astype(np.uint8))
+                    filename = f"{result_id}.png"
+                    out_path = os.path.join(RESULTS_DIR, filename)
+                    out_img.save(out_path)
+                    _jobs[job_id]["result"] = {"filename": filename}
+                    _jobs[job_id]["status"] = "completed"
+                    LOGGER.info(f"[job {job_id}] Completed -> {out_path}")
+                except Exception as e:
+                    LOGGER.error(f"[job {job_id}] Hair transfer failed: {str(e)}")
+                    _jobs[job_id]["error"] = str(e)
+                    _jobs[job_id]["status"] = "failed"
+        except Exception as e:
+            _jobs[job_id]["error"] = str(e)
+            _jobs[job_id]["status"] = "failed"
+        finally:
+            _jobs[job_id]["ended_at"] = time.time()
+
+    asyncio.create_task(_run_job())
+    return {"job_id": job_id, "status": "queued"}
+
+
+@app.get("/job/{job_id}")
+def job_status(job_id: str, _=Depends(verify_bearer)):
+    data = _jobs.get(job_id)
+    if not data:
+        raise HTTPException(status_code=404, detail="job not found")
+    return data
+
+