877 lines
29 KiB
Python
877 lines
29 KiB
Python
"""
|
|
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 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, 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")
|
|
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",
|
|
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_NEGATIVE = "6"
|
|
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")
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_methods=["GET", "POST", "DELETE"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
# --- 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 _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(
|
|
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
|
|
|
|
|
|
# --- 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,
|
|
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
|
|
|
|
# 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"]
|
|
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_bytes = _comfy_fetch_image(outputs)
|
|
|
|
if is_transparent:
|
|
png_bytes = _apply_transparency(png_bytes)
|
|
|
|
return png_bytes
|
|
|
|
|
|
# --- 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 _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")
|
|
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:
|
|
print(f"Error opening {fname}: {e}")
|
|
jobs[job_id]["failed"] += len(prompts)
|
|
|
|
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 | 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": total_tasks, "done": 0, "failed": 0}
|
|
t = threading.Thread(
|
|
target=_batch_worker,
|
|
args=(job_id, req.filenames, prompts, req.seed, req.max_area, req.group_id),
|
|
daemon=True,
|
|
)
|
|
t.start()
|
|
return {"job_id": job_id, "total": total_tasks}
|
|
|
|
|
|
@app.get("/poses")
|
|
def get_poses():
|
|
return _load_poses()
|
|
|
|
|
|
@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("/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:
|
|
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}")
|
|
|
|
|
|
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(...),
|
|
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}")
|
|
|
|
png = _run_pipeline(pil, prompt, seed, max_area, steps, cfg, sampler_name, scheduler)
|
|
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"),
|
|
port=int(os.environ.get("PORT", "8500")))
|