aa

tour-comfy/edit_api.py

@@ -15,13 +15,17 @@ import time
 import uuid
 import random
 import copy
+import threading
 
 import requests
 from PIL import Image
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import Response
+from pydantic import BaseModel
 
 # --- config -----------------------------------------------------------------
+CONFIG_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.json")
 COMFY = os.environ.get("COMFY_URL", "http://127.0.0.1:8188").rstrip("/")
 WORKFLOW_PATH = os.environ.get(
     "WORKFLOW_PATH",
@@ -45,6 +49,12 @@ 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")
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["GET", "POST"],
+    allow_headers=["*"],
+)
 
 
 # --- helpers ----------------------------------------------------------------
@@ -58,6 +68,41 @@ def _target_size(w: int, h: int, max_area: int) -> tuple[int, int]:
     return _round16(w * scale), _round16(h * scale)
 
 
+def _prep_image(pil: Image.Image, max_area: int) -> tuple[Image.Image, int, int]:
+    """
+    Prepare image for ComfyUI:
+    1. If area > max_area, crop from bottom if height remains >= 256.
+    2. Otherwise scale (up or down) to fit area while preserving aspect.
+    3. Ensure dimensions are rounded to 16.
+    """
+    w, h = pil.width, pil.height
+    if w * h > max_area:
+        # Try to keep width and crop height from bottom
+        rw = _round16(w)
+        th = max_area // rw
+        if th >= 256:
+            rh = (th // 16) * 16
+            if rh < 16: rh = 16
+            
+            # To avoid black bars from .crop((0,0,rw,rh)) when rw > w,
+            # we crop to original w first, then resize to rw.
+            pil = pil.crop((0, 0, w, min(h, (rh * w) // rw)))
+            pil = pil.resize((rw, rh), resample=Image.LANCZOS)
+            return pil, rw, rh
+        else:
+            # Too wide to keep width and have decent height, scale both down
+            rw, rh = _target_size(w, h, max_area)
+            pil = pil.resize((rw, rh), resample=Image.LANCZOS)
+            return pil, rw, rh
+    else:
+        # Fits or is too small: scale UP to match the max_area budget
+        # (Legacy behavior that gives better model performance)
+        rw, rh = _target_size(w, h, max_area)
+        if rw != w or rh != h:
+            pil = pil.resize((rw, rh), resample=Image.LANCZOS)
+        return pil, rw, rh
+
+
 def _comfy_upload(img_bytes: bytes, filename: str) -> str:
     """Upload an image to ComfyUI's input dir; return the stored name."""
     r = requests.post(
@@ -122,7 +167,122 @@ def _comfy_fetch_image(outputs: dict) -> bytes:
     return r.content
 
 
+# --- pipeline helper ---------------------------------------------------------
+
+def _run_pipeline(
+    pil: Image.Image,
+    prompt: str,
+    seed: int = -1,
+    max_area: int = 0,
+    steps: int = 4,
+    cfg: float = 1.0,
+    sampler_name: str = "euler_ancestral",
+    scheduler: str = "beta",
+) -> bytes:
+    area = max_area if max_area > 0 else MAX_AREA
+    pil, w, h = _prep_image(pil, 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)
+    return _comfy_fetch_image(outputs)
+
+
+# --- batch state -------------------------------------------------------------
+
+jobs: dict[str, dict] = {}
+
+
+def _load_output_dir() -> str:
+    with open(CONFIG_PATH, "r") as f:
+        conf = json.load(f)
+    d = conf["output_dir"]
+    if not os.path.isabs(d):
+        d = os.path.normpath(os.path.join(os.path.dirname(CONFIG_PATH), "..", d))
+    return d
+
+
+def _batch_worker(job_id: str, filenames: list, prompt: str, seed: int, max_area: int):
+    output_dir = _load_output_dir()
+    for fname in filenames:
+        fpath = os.path.join(output_dir, fname)
+        try:
+            pil = Image.open(fpath).convert("RGB")
+            png = _run_pipeline(pil, prompt, seed, max_area)
+            ts = time.strftime("%Y%m%d_%H%M%S")
+            out_name = f"{ts}_{fname}"
+            with open(os.path.join(output_dir, out_name), "wb") as f:
+                f.write(png)
+            jobs[job_id]["done"] += 1
+        except Exception as e:
+            jobs[job_id]["failed"] += 1
+    jobs[job_id]["status"] = "done"
+
+
 # --- routes -----------------------------------------------------------------
+
+class ConfigUpdate(BaseModel):
+    prompt: str | None = None
+    seed: int | None = None
+
+
+@app.get("/config")
+def get_config():
+    with open(CONFIG_PATH, "r") as f:
+        return json.load(f)
+
+
+@app.post("/config")
+def update_config(update: ConfigUpdate):
+    with open(CONFIG_PATH, "r") as f:
+        conf = json.load(f)
+    if update.prompt is not None:
+        conf["prompt"] = update.prompt
+    if update.seed is not None:
+        conf["seed"] = update.seed
+    with open(CONFIG_PATH, "w") as f:
+        json.dump(conf, f, indent=2)
+    return {"prompt": conf["prompt"], "seed": conf["seed"]}
+
+
+class BatchRequest(BaseModel):
+    filenames: list[str]
+    prompt: str
+    seed: int = -1
+    max_area: int = 0
+
+
+@app.post("/batch")
+def start_batch(req: BatchRequest):
+    job_id = uuid.uuid4().hex[:8]
+    jobs[job_id] = {"status": "running", "total": len(req.filenames), "done": 0, "failed": 0}
+    t = threading.Thread(
+        target=_batch_worker,
+        args=(job_id, req.filenames, req.prompt, req.seed, req.max_area),
+        daemon=True,
+    )
+    t.start()
+    return {"job_id": job_id, "total": len(req.filenames)}
+
+
+@app.get("/batch/{job_id}")
+def get_batch(job_id: str):
+    if job_id not in jobs:
+        raise HTTPException(404, "Job not found")
+    return jobs[job_id]
+
+
 @app.get("/health")
 def health():
     try:
@@ -149,40 +309,8 @@ async def edit(
     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,
-        },
-    )
+    png = _run_pipeline(pil, prompt, seed, max_area, steps, cfg, sampler_name, scheduler)
+    return Response(content=png, media_type="image/png")
 
 
 if __name__ == "__main__":