aa

tour-comfy/edit_api.py

@@ -16,16 +16,31 @@ import uuid
 import random
 import copy
 import threading
+import csv
+
+try:
+    from . import database
+    from . import embeddings
+    from . import naming
+except ImportError:
+    import database
+    import embeddings
+    import naming
 
 import requests
 from PIL import Image
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException, BackgroundTasks
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import Response
 from pydantic import BaseModel
+import shutil
+import re
 
 # --- config -----------------------------------------------------------------
-CONFIG_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.json")
+CONFIG_PATH  = os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.json")
+NAMES_PATH   = os.path.join(os.path.dirname(os.path.abspath(__file__)), "names.json")
+GROUPS_PATH  = os.path.join(os.path.dirname(os.path.abspath(__file__)), "groups.json")
+WD_MODEL     = os.environ.get("WD_MODEL", "SmilingWolf/wd-vit-tagger-v3")
 COMFY = os.environ.get("COMFY_URL", "http://127.0.0.1:8188").rstrip("/")
 WORKFLOW_PATH = os.environ.get(
     "WORKFLOW_PATH",
@@ -39,6 +54,7 @@ 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_NEGATIVE = "6"
 NODE_LATENT = "7"
 NODE_KSAMPLER = "8"
 NODE_SAVE = "10"
@@ -52,7 +68,7 @@ app = FastAPI(title="Qwen-Image-Edit Rapid-AIO API", version="1.0")
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
-    allow_methods=["GET", "POST"],
+    allow_methods=["GET", "POST", "DELETE"],
     allow_headers=["*"],
 )
 
@@ -167,8 +183,122 @@ def _comfy_fetch_image(outputs: dict) -> bytes:
     return r.content
 
 
+# --- WD tagger (lazy) -------------------------------------------------------
+
+_tagger      = None   # (model, transform, labels) once loaded
+_tagger_lock = threading.Lock()
+
+
+def _load_tagger():
+    global _tagger
+    if _tagger is not None:
+        return _tagger
+    with _tagger_lock:
+        if _tagger is not None:
+            return _tagger
+        import torch
+        import timm
+        from timm.data import create_transform, resolve_data_config
+        import huggingface_hub
+
+        model = timm.create_model(f"hf_hub:{WD_MODEL}", pretrained=True).eval()
+        if torch.cuda.is_available():
+            model = model.cuda()
+
+        cfg       = resolve_data_config(model.pretrained_cfg, model=model)
+        transform = create_transform(**cfg)
+
+        lpath  = huggingface_hub.hf_hub_download(WD_MODEL, "selected_tags.csv")
+        with open(lpath, newline="") as f:
+            rows   = list(csv.DictReader(f))
+        # category 0=general 4=character 9=rating
+        labels = [(r["name"], int(r.get("category", 9))) for r in rows]
+
+        _tagger = (model, transform, labels)
+    return _tagger
+
+
+def _run_tagger(pil_img: Image.Image, threshold: float = 0.35):
+    import torch
+    model, transform, labels = _load_tagger()
+    tensor = transform(pil_img.convert("RGB")).unsqueeze(0)
+    if torch.cuda.is_available():
+        tensor = tensor.cuda()
+    with torch.no_grad():
+        scores = torch.sigmoid(model(tensor))[0].cpu().tolist()
+    tags = [
+        {"tag": name, "score": round(score, 3), "cat": cat}
+        for (name, cat), score in zip(labels, scores)
+        if score >= threshold
+    ]
+    tags.sort(key=lambda x: -x["score"])
+    return tags
+
+
+def _tags_to_name(tags: list, max_tags: int = 8) -> str:
+    content = [t["tag"] for t in tags if t["cat"] in (0, 4)][:max_tags]
+    return " ".join(content).replace("_", " ")
+
+
+def _load_json(path: str) -> dict:
+    if os.path.exists(path):
+        with open(path) as f:
+            return json.load(f)
+    return {}
+
+
+def _save_json(path: str, data: dict):
+    with open(path, "w") as f:
+        json.dump(data, f, indent=2)
+
+
+def _apply_transparency(png_bytes: bytes) -> bytes:
+    """Use rembg to remove background and return PNG bytes with Alpha channel."""
+    try:
+        from rembg import remove
+        import io
+        from PIL import Image
+        img = Image.open(io.BytesIO(png_bytes))
+        # rembg works best on RGB
+        if img.mode != "RGB":
+            img = img.convert("RGB")
+        out = remove(img)
+        buf = io.BytesIO()
+        out.save(buf, format="PNG")
+        return buf.getvalue()
+    except Exception as e:
+        print(f"Error in transparency post-processing: {e}")
+        return png_bytes
+
+
 # --- pipeline helper ---------------------------------------------------------
 
+def _load_poses():
+    poses_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "poses.md")
+    if not os.path.exists(poses_path):
+        return {}
+    
+    poses = {}
+    current_pose = None
+    current_desc = []
+    
+    with open(poses_path, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line.startswith("# "):
+                if current_pose:
+                    poses[current_pose] = " ".join(current_desc).strip()
+                current_pose = line[2:].rstrip(":").strip()
+                current_desc = []
+            elif line and current_pose:
+                current_desc.append(line)
+        
+        if current_pose:
+            poses[current_pose] = " ".join(current_desc).strip()
+            
+    return poses
+
+
 def _run_pipeline(
     pil: Image.Image,
     prompt: str,
@@ -189,6 +319,12 @@ def _run_pipeline(
     graph = copy.deepcopy(BASE_WORKFLOW)
     graph[NODE_LOADIMAGE]["inputs"]["image"] = stored
     graph[NODE_POSITIVE]["inputs"]["prompt"] = prompt
+    
+    # Transparency detection
+    is_transparent = any(kw in prompt.lower() for kw in ["transparent", "no background", "remove background", "alpha channel"])
+    if is_transparent:
+        graph[NODE_NEGATIVE]["inputs"]["prompt"] = "checkerboard, grid, pattern, texture, background details"
+        
     graph[NODE_LATENT]["inputs"]["width"] = w
     graph[NODE_LATENT]["inputs"]["height"] = h
     ks = graph[NODE_KSAMPLER]["inputs"]
@@ -196,7 +332,12 @@ def _run_pipeline(
     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)
+    png_bytes = _comfy_fetch_image(outputs)
+    
+    if is_transparent:
+        png_bytes = _apply_transparency(png_bytes)
+        
+    return png_bytes
 
 
 # --- batch state -------------------------------------------------------------
@@ -213,20 +354,76 @@ def _load_output_dir() -> str:
     return d
 
 
-def _batch_worker(job_id: str, filenames: list, prompt: str, seed: int, max_area: int):
+def _move_to_trash(filepath: str):
+    if not filepath or not os.path.exists(filepath):
+        return
+    output_dir = _load_output_dir()
+    trash_dir = os.path.join(output_dir, ".trash")
+    os.makedirs(trash_dir, exist_ok=True)
+    
+    filename = os.path.basename(filepath)
+    ts = time.strftime("%Y%m%d_%H%M%S")
+    trash_path = os.path.join(trash_dir, f"{ts}_{filename}")
+    
+    try:
+        shutil.move(filepath, trash_path)
+    except Exception as e:
+        print(f"Error moving {filepath} to trash: {e}")
+
+
+def _batch_worker(job_id: str, filenames: list, prompts: list[str], seed: int, max_area: int, group_id: str | None = None):
     output_dir = _load_output_dir()
     for fname in filenames:
+        # If no group_id provided, try to inherit from source
+        actual_gid = group_id
+        if not actual_gid:
+            try:
+                person = database.get_person(fname)
+                if person and person[1]: # group_id is at index 1
+                    actual_gid = person[1]
+                else:
+                    # Create a new group for this standalone image
+                    actual_gid = naming.get_base_name(fname)
+                    # Update source image to join this group
+                    database.upsert_person(fname, group_id=actual_gid)
+            except Exception as e:
+                print(f"Error determining group for {fname}: {e}")
+
         fpath = os.path.join(output_dir, fname)
+        if not os.path.exists(fpath):
+            jobs[job_id]["failed"] += len(prompts)
+            continue
+
         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
+            for prompt in prompts:
+                try:
+                    png = _run_pipeline(pil, prompt, seed, max_area)
+                    ts = time.strftime("%Y%m%d_%H%M%S")
+                    
+                    # Clean filename to avoid nested timestamps
+                    clean_fname = naming.get_base_name(fname)
+                    out_name = f"{ts}_{clean_fname}"
+                    out_path = os.path.join(output_dir, out_name)
+                    
+                    with open(out_path, "wb") as f:
+                        f.write(png)
+                    
+                    # Register in DB
+                    try:
+                        embedding = embeddings.generate_embedding(out_path)
+                        database.upsert_person(out_name, filepath=out_path, embedding=embedding, group_id=actual_gid)
+                    except Exception as db_err:
+                        print(f"Database error in batch worker: {db_err}")
+                        
+                    jobs[job_id]["done"] += 1
+                except Exception as e:
+                    print(f"Error in batch for {fname} with prompt '{prompt}': {e}")
+                    jobs[job_id]["failed"] += 1
         except Exception as e:
-            jobs[job_id]["failed"] += 1
+            print(f"Error opening {fname}: {e}")
+            jobs[job_id]["failed"] += len(prompts)
+            
     jobs[job_id]["status"] = "done"
 
 
@@ -258,22 +455,31 @@ def update_config(update: ConfigUpdate):
 
 class BatchRequest(BaseModel):
     filenames: list[str]
-    prompt: str
+    prompt: str | list[str]
     seed: int = -1
     max_area: int = 0
+    group_id: str | None = None
 
 
 @app.post("/batch")
 def start_batch(req: BatchRequest):
+    prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt
+    total_tasks = len(req.filenames) * len(prompts)
+    
     job_id = uuid.uuid4().hex[:8]
-    jobs[job_id] = {"status": "running", "total": len(req.filenames), "done": 0, "failed": 0}
+    jobs[job_id] = {"status": "running", "total": total_tasks, "done": 0, "failed": 0}
     t = threading.Thread(
         target=_batch_worker,
-        args=(job_id, req.filenames, req.prompt, req.seed, req.max_area),
+        args=(job_id, req.filenames, prompts, req.seed, req.max_area, req.group_id),
         daemon=True,
     )
     t.start()
-    return {"job_id": job_id, "total": len(req.filenames)}
+    return {"job_id": job_id, "total": total_tasks}
+
+
+@app.get("/poses")
+def get_poses():
+    return _load_poses()
 
 
 @app.get("/batch/{job_id}")
@@ -283,6 +489,186 @@ def get_batch(job_id: str):
     return jobs[job_id]
 
 
+@app.get("/images")
+def list_images():
+    output_dir = _load_output_dir()
+    extensions = ('.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp', '.svg')
+    try:
+        # Try to get from DB first
+        try:
+            persons = database.list_persons()
+            # persons is (filename, name, group_id, clip_description)
+            db_images = []
+            for p in persons:
+                db_images.append({
+                    "filename": p[0],
+                    "name": p[1],
+                    "group_id": p[2],
+                    "clip_description": p[3]
+                })
+            # Still sort by mtime for consistency with filesystem
+            db_images.sort(key=lambda x: os.path.getmtime(os.path.join(output_dir, x["filename"])) if os.path.exists(os.path.join(output_dir, x["filename"])) else 0, reverse=True)
+            return {"images": db_images}
+        except Exception as db_err:
+            print(f"DB error in list_images: {db_err}")
+            # Fallback to filesystem
+            files = [f for f in os.listdir(output_dir) if f.lower().endswith(extensions)]
+            files.sort(key=lambda x: os.path.getmtime(os.path.join(output_dir, x)), reverse=True)
+            return {"images": [{"filename": f} for f in files]}
+    except Exception as e:
+        raise HTTPException(500, str(e))
+
+
+# --- tagging routes ----------------------------------------------------------
+
+class TagRequest(BaseModel):
+    filename: str
+    threshold: float = 0.35
+    max_tags: int = 8
+    group_id: str | None = None
+
+
+@app.post("/tag")
+def tag_image(req: TagRequest):
+    output_dir = _load_output_dir()
+    fpath = os.path.join(output_dir, req.filename)
+    if not os.path.exists(fpath):
+        raise HTTPException(404, "File not found in output dir")
+    try:
+        pil  = Image.open(fpath)
+        tags = _run_tagger(pil, req.threshold)
+        clip_desc = _tags_to_name(tags, req.max_tags)
+        auto_name = naming.generate_associative_name(tags)
+        
+        # Save to DB
+        try:
+            embedding = embeddings.generate_embedding(fpath)
+            database.upsert_person(req.filename, filepath=fpath, name=auto_name, clip_description=clip_desc, tags=tags, embedding=embedding, group_id=req.group_id)
+        except Exception as db_err:
+            print(f"Database error during tag: {db_err}")
+            
+        # Legacy fallback
+        try:
+            names = _load_json(NAMES_PATH)
+            names[req.filename] = clip_desc
+            _save_json(NAMES_PATH, names)
+        except: pass
+        
+        return {"filename": req.filename, "clip_description": clip_desc, "tags": tags[:30]}
+    except Exception as e:
+        raise HTTPException(500, str(e))
+
+
+@app.get("/names")
+def get_names():
+    try:
+        persons = database.list_persons()
+        return {p[0]: p[1] for p in persons if p[1]}
+    except:
+        return _load_json(NAMES_PATH)
+
+
+@app.post("/names/{filename}")
+def set_name(filename: str, body: dict):
+    name = body.get("name", "")
+    try:
+        database.upsert_person(filename, name=name)
+    except Exception as db_err:
+        print(f"Database error in set_name: {db_err}")
+    
+    # Legacy fallback
+    try:
+        names = _load_json(NAMES_PATH)
+        names[filename] = name
+        _save_json(NAMES_PATH, names)
+    except: pass
+    
+    return {"filename": filename, "name": name}
+
+
+# --- group routes ------------------------------------------------------------
+
+@app.get("/groups")
+def get_groups():
+    try:
+        persons = database.list_persons()
+        return {p[0]: p[2] for p in persons if p[2]}
+    except:
+        return _load_json(GROUPS_PATH)
+
+
+class MergeRequest(BaseModel):
+    filenames: list[str]
+    group_id: str | None = None
+
+
+@app.post("/groups/merge")
+def merge_groups(req: MergeRequest):
+    gid = req.group_id or f"cg_{uuid.uuid4().hex[:8]}"
+    for fname in req.filenames:
+        try:
+            database.upsert_person(fname, group_id=gid)
+        except Exception as db_err:
+            print(f"Database error in merge: {db_err}")
+    
+    # Legacy fallback
+    try:
+        groups = _load_json(GROUPS_PATH)
+        for fname in req.filenames:
+            groups[fname] = gid
+        _save_json(GROUPS_PATH, groups)
+    except: pass
+    
+    return {"group_id": gid, "files": req.filenames}
+
+
+class ExtractRequest(BaseModel):
+    filename: str
+
+
+@app.post("/groups/extract")
+def extract_from_group(req: ExtractRequest):
+    gid = f"solo:{req.filename}"
+    try:
+        database.upsert_person(req.filename, group_id=gid)
+    except Exception as db_err:
+        print(f"Database error in extract: {db_err}")
+    
+    # Legacy fallback
+    try:
+        groups = _load_json(GROUPS_PATH)
+        groups[req.filename] = gid
+        _save_json(GROUPS_PATH, groups)
+    except: pass
+    
+    return {"filename": req.filename}
+
+
+@app.get("/similar/{filename}")
+def get_similar(filename: str, limit: int = 10):
+    person = database.get_person(filename)
+    if not person or person[3] is None:
+        raise HTTPException(404, "Image or embedding not found")
+    
+    embedding = person[3]
+    results = database.search_similar(embedding, limit=limit)
+    
+    similar = []
+    for r in results:
+        # Avoid returning the same image as the most similar
+        if r[0] == filename:
+            continue
+        similar.append({
+            "filename": r[0],
+            "name": r[1],
+            "group_id": r[2],
+            "clip_description": r[3],
+            "distance": float(r[4])
+        })
+    
+    return {"filename": filename, "similar": similar}
+
+
 @app.get("/health")
 def health():
     try:
@@ -292,6 +678,119 @@ def health():
         raise HTTPException(503, f"ComfyUI unreachable at {COMFY}: {e}")
 
 
+def _crop_to_bbox(pil_img: Image.Image, margin: int = 20, top_margin: int = 20, headroom: float = 0.05) -> Image.Image:
+    if pil_img.mode != 'RGBA':
+        return pil_img
+    
+    alpha = pil_img.split()[-1]
+    bbox = alpha.getbbox()
+    if not bbox:
+        return pil_img
+    
+    left, upper, right, lower = bbox
+    left = max(0, left - margin)
+    upper = max(0, upper - top_margin)
+    right = min(pil_img.width, right + margin)
+    lower = min(pil_img.height, lower + margin)
+    
+    cropped = pil_img.crop((left, upper, right, lower))
+    if headroom > 0:
+        h_px = int(cropped.height * headroom)
+        if h_px > 0:
+            new_img = Image.new("RGBA", (cropped.width, cropped.height + h_px), (0, 0, 0, 0))
+            new_img.paste(cropped, (0, h_px))
+            return new_img
+    return cropped
+
+
+def _process_upload(file_path: str, filename: str, prompts: list[str], name: str | None = None, group_id: str | None = None):
+    output_dir = _load_output_dir()
+    try:
+        pil = Image.open(file_path)
+        
+        # 1. CLIP tag the source
+        tags = _run_tagger(pil.convert("RGB"))
+        clip_desc = _tags_to_name(tags)
+        auto_name = name or naming.generate_associative_name(tags)
+        
+        # 2. Embedding for source
+        embedding = embeddings.generate_embedding(file_path)
+        
+        # 3. Register source in DB (optional, but good for tracking)
+        # We'll use the original filename or a timestamped one
+        database.upsert_person(filename, filepath=file_path, name=auto_name, clip_description=clip_desc, tags=tags, embedding=embedding, group_id=group_id)
+        
+        # 4. Default behavior: Crop if needed
+        # We'll use default values from watcher
+        cropped_pil = _crop_to_bbox(pil)
+        
+        # 5. Run prompts
+        for i, prompt in enumerate(prompts):
+            try:
+                # Use RGB for pipeline
+                png = _run_pipeline(cropped_pil.convert("RGB"), prompt)
+                
+                ts = time.strftime("%Y%m%d_%H%M%S")
+                out_name = f"{ts}_{i}_{filename}"
+                if not out_name.lower().endswith(".png"):
+                    out_name += ".png"
+                out_path = os.path.join(output_dir, out_name)
+                
+                with open(out_path, "wb") as f:
+                    f.write(png)
+                
+                # Register output in DB
+                out_embedding = embeddings.generate_embedding(out_path)
+                database.upsert_person(out_name, filepath=out_path, name=auto_name, clip_description=clip_desc, embedding=out_embedding, group_id=group_id)
+            except Exception as e:
+                print(f"Error processing prompt '{prompt}' for {filename}: {e}")
+                
+    except Exception as e:
+        print(f"Error in _process_upload for {filename}: {e}")
+
+
+@app.post("/upload")
+def upload_image(
+    background_tasks: BackgroundTasks,
+    image: UploadFile = File(...),
+    prompts: str = Form(""),
+    name: str = Form(None),
+):
+    # Load config to get output_dir (we use output_dir for UI uploads to avoid watcher conflict)
+    with open(CONFIG_PATH, "r") as f:
+        conf = json.load(f)
+    output_dir = _load_output_dir()
+    os.makedirs(output_dir, exist_ok=True)
+    
+    ts = time.strftime("%Y%m%d_%H%M%S")
+    safe_filename = re.sub(r'[^a-zA-Z0-9_.-]', '_', image.filename)
+    # Ensure extension
+    if not safe_filename.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
+        safe_filename += ".png"
+    
+    filename = f"{ts}_{safe_filename}"
+    file_path = os.path.join(output_dir, filename)
+    
+    with open(file_path, "wb") as f:
+        shutil.copyfileobj(image.file, f)
+    
+    prompt_list = [p.strip() for p in prompts.split(",") if p.strip()]
+    
+    # Add base-set prompts if defined in config
+    base_prompts = conf.get("base_prompts", [])
+    if isinstance(base_prompts, list):
+        prompt_list.extend(base_prompts)
+    
+    if not prompt_list:
+        # Use default prompt from config
+        prompt_list = [conf.get("prompt", "high quality, masterpiece")]
+    
+    group_id = naming.get_base_name(filename)
+    background_tasks.add_task(_process_upload, file_path, filename, prompt_list, name, group_id)
+    
+    return {"status": "processing", "filename": filename, "group_id": group_id, "prompts": prompt_list}
+
+
 @app.post("/edit")
 async def edit(
     image: UploadFile = File(...),
@@ -313,6 +812,64 @@ async def edit(
     return Response(content=png, media_type="image/png")
 
 
+@app.delete("/images/{filename}")
+def delete_image(filename: str):
+    person = database.get_person(filename)
+    if person and person[5] and os.path.exists(person[5]):
+        _move_to_trash(person[5])
+
+    database.delete_person(filename)
+    return {"status": "deleted", "filename": filename}
+
+
+@app.delete("/groups/{group_id}")
+def delete_group(group_id: str):
+    files = database.get_group_files(group_id)
+    for filename, filepath in files:
+        if filepath and os.path.exists(filepath):
+            _move_to_trash(filepath)
+    
+    database.delete_group(group_id)
+    return {"status": "deleted", "group_id": group_id}
+
+
+@app.post("/remove-background/{filename}")
+def remove_background(filename: str):
+    person = database.get_person(filename)
+    if not person or not person[5] or not os.path.exists(person[5]):
+        raise HTTPException(404, "Image file not found")
+    
+    path = person[5]
+    with open(path, "rb") as f:
+        png_bytes = f.read()
+    
+    transparent_png = _apply_transparency(png_bytes)
+    
+    with open(path, "wb") as f:
+        f.write(transparent_png)
+        
+    return {"status": "success", "filename": filename}
+
+
+@app.post("/remove-background/group/{group_id}")
+def remove_background_group(group_id: str, background_tasks: BackgroundTasks):
+    def _bg_task():
+        files = database.get_group_files(group_id)
+        for filename, filepath in files:
+            if filepath and os.path.exists(filepath):
+                try:
+                    with open(filepath, "rb") as f:
+                        png_bytes = f.read()
+                    transparent_png = _apply_transparency(png_bytes)
+                    with open(filepath, "wb") as f:
+                        f.write(transparent_png)
+                except Exception as e:
+                    print(f"Error removing background for {filename}: {e}")
+    
+    background_tasks.add_task(_bg_task)
+    return {"status": "processing", "group_id": group_id}
+
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host=os.environ.get("HOST", "0.0.0.0"),