aa

tour-comfy/edit_api.py

@@ -0,0 +1,191 @@
+"""
+edit_api.py — headless throughput API for Qwen-Image-Edit Rapid-AIO (v23 Q8 GGUF)
+running on top of a local ComfyUI server.
+
+Flow per request:  image + prompt  ->  upload to ComfyUI  ->  inject into the
+workflow graph  ->  queue  ->  poll until done  ->  return the edited PNG.
+
+Run ComfyUI first (run_comfyui.sh), then this service (start_api.sh).
+"""
+
+import io
+import os
+import json
+import time
+import uuid
+import random
+import copy
+
+import requests
+from PIL import Image
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi.responses import Response
+
+# --- config -----------------------------------------------------------------
+COMFY = os.environ.get("COMFY_URL", "http://127.0.0.1:8188").rstrip("/")
+WORKFLOW_PATH = os.environ.get(
+    "WORKFLOW_PATH",
+    os.path.join(os.path.dirname(os.path.abspath(__file__)), "workflow_qwen_edit.json"),
+)
+# Default target pixel area for the output latent. The MI50 is not fast, so we
+# cap at ~1MP by default; raise via MAX_AREA env if you want bigger output.
+MAX_AREA = int(os.environ.get("MAX_AREA", str(1024 * 1024)))
+GEN_TIMEOUT = int(os.environ.get("GEN_TIMEOUT", "600"))  # seconds per request
+
+# Node ids in workflow_qwen_edit.json (kept stable on purpose).
+NODE_LOADIMAGE = "4"
+NODE_POSITIVE = "5"
+NODE_LATENT = "7"
+NODE_KSAMPLER = "8"
+NODE_SAVE = "10"
+
+MAX_SEED = 2**32 - 1
+
+with open(WORKFLOW_PATH, "r", encoding="utf-8") as f:
+    BASE_WORKFLOW = json.load(f)
+
+app = FastAPI(title="Qwen-Image-Edit Rapid-AIO API", version="1.0")
+
+
+# --- helpers ----------------------------------------------------------------
+def _round16(x: int) -> int:
+    return max(16, int(round(x / 16.0)) * 16)
+
+
+def _target_size(w: int, h: int, max_area: int) -> tuple[int, int]:
+    """Scale (w, h) to ~max_area preserving aspect, rounded to /16."""
+    scale = (max_area / float(w * h)) ** 0.5
+    return _round16(w * scale), _round16(h * scale)
+
+
+def _comfy_upload(img_bytes: bytes, filename: str) -> str:
+    """Upload an image to ComfyUI's input dir; return the stored name."""
+    r = requests.post(
+        f"{COMFY}/upload/image",
+        files={"image": (filename, img_bytes, "image/png")},
+        data={"overwrite": "true", "type": "input"},
+        timeout=60,
+    )
+    r.raise_for_status()
+    j = r.json()
+    name = j["name"]
+    sub = j.get("subfolder", "")
+    return f"{sub}/{name}" if sub else name
+
+
+def _comfy_queue(graph: dict, client_id: str) -> str:
+    r = requests.post(
+        f"{COMFY}/prompt",
+        json={"prompt": graph, "client_id": client_id},
+        timeout=60,
+    )
+    if r.status_code != 200:
+        raise HTTPException(502, f"ComfyUI rejected workflow: {r.text}")
+    return r.json()["prompt_id"]
+
+
+def _comfy_wait(prompt_id: str, deadline: float) -> dict:
+    """Poll /history until the prompt finishes; return its outputs dict."""
+    while time.time() < deadline:
+        r = requests.get(f"{COMFY}/history/{prompt_id}", timeout=30)
+        if r.status_code == 200:
+            hist = r.json()
+            if prompt_id in hist:
+                entry = hist[prompt_id]
+                status = entry.get("status", {})
+                if status.get("status_str") == "error":
+                    raise HTTPException(500, f"ComfyUI execution error: {json.dumps(status)}")
+                outputs = entry.get("outputs", {})
+                if outputs:
+                    return outputs
+        time.sleep(0.5)
+    raise HTTPException(504, f"Generation timed out after {GEN_TIMEOUT}s")
+
+
+def _comfy_fetch_image(outputs: dict) -> bytes:
+    node_out = outputs.get(NODE_SAVE) or next(
+        (v for v in outputs.values() if "images" in v), None
+    )
+    if not node_out or not node_out.get("images"):
+        raise HTTPException(500, "No output image produced")
+    img = node_out["images"][0]
+    r = requests.get(
+        f"{COMFY}/view",
+        params={
+            "filename": img["filename"],
+            "subfolder": img.get("subfolder", ""),
+            "type": img.get("type", "output"),
+        },
+        timeout=60,
+    )
+    r.raise_for_status()
+    return r.content
+
+
+# --- routes -----------------------------------------------------------------
+@app.get("/health")
+def health():
+    try:
+        requests.get(f"{COMFY}/system_stats", timeout=5).raise_for_status()
+        return {"status": "ok", "comfy": COMFY}
+    except Exception as e:
+        raise HTTPException(503, f"ComfyUI unreachable at {COMFY}: {e}")
+
+
+@app.post("/edit")
+async def edit(
+    image: UploadFile = File(...),
+    prompt: str = Form(...),
+    seed: int = Form(-1),
+    steps: int = Form(4),
+    cfg: float = Form(1.0),
+    sampler_name: str = Form("euler_ancestral"),
+    scheduler: str = Form("beta"),
+    max_area: int = Form(0),
+):
+    raw = await image.read()
+    try:
+        pil = Image.open(io.BytesIO(raw)).convert("RGB")
+    except Exception as e:
+        raise HTTPException(400, f"Invalid image: {e}")
+
+    area = max_area if max_area > 0 else MAX_AREA
+    w, h = _target_size(pil.width, pil.height, area)
+
+    buf = io.BytesIO()
+    pil.save(buf, format="PNG")
+    stored = _comfy_upload(buf.getvalue(), f"in_{uuid.uuid4().hex[:8]}.png")
+
+    if seed is None or seed < 0:
+        seed = random.randint(0, MAX_SEED)
+
+    graph = copy.deepcopy(BASE_WORKFLOW)
+    graph[NODE_LOADIMAGE]["inputs"]["image"] = stored
+    graph[NODE_POSITIVE]["inputs"]["prompt"] = prompt
+    graph[NODE_LATENT]["inputs"]["width"] = w
+    graph[NODE_LATENT]["inputs"]["height"] = h
+    ks = graph[NODE_KSAMPLER]["inputs"]
+    ks.update(seed=seed, steps=steps, cfg=cfg,
+              sampler_name=sampler_name, scheduler=scheduler)
+
+    client_id = uuid.uuid4().hex
+    prompt_id = _comfy_queue(graph, client_id)
+    outputs = _comfy_wait(prompt_id, time.time() + GEN_TIMEOUT)
+    png = _comfy_fetch_image(outputs)
+
+    return Response(
+        content=png,
+        media_type="image/png",
+        headers={
+            "X-Seed": str(seed),
+            "X-Width": str(w),
+            "X-Height": str(h),
+            "X-Prompt-Id": prompt_id,
+        },
+    )
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host=os.environ.get("HOST", "0.0.0.0"),
+                port=int(os.environ.get("PORT", "8500")))