This commit is contained in:
mike
2026-06-19 20:44:36 +02:00
parent 1056f1d460
commit f1f693523b
5 changed files with 2704 additions and 114 deletions

View File

@@ -38,8 +38,6 @@ 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(
@@ -61,6 +59,9 @@ NODE_SAVE = "10"
MAX_SEED = 2**32 - 1
# Poses where the source image should be rotated 180° before pipeline for better results
ROTATE_180_POSES = {"the dragon", "dragon", "the draak", "draak"}
with open(WORKFLOW_PATH, "r", encoding="utf-8") as f:
BASE_WORKFLOW = json.load(f)
@@ -72,6 +73,43 @@ app.add_middleware(
allow_headers=["*"],
)
def _sync_car_html():
src = os.path.join(os.path.dirname(os.path.abspath(__file__)), "car.html")
if not os.path.exists(src):
return
try:
dest = os.path.join(_load_output_dir(), "car.html")
shutil.copy2(src, dest)
print(f"[car.html] synced → {dest}")
except Exception as e:
print(f"[car.html] sync warning: {e}")
def _watch_car_html():
src = os.path.join(os.path.dirname(os.path.abspath(__file__)), "car.html")
last_mtime = os.path.getmtime(src) if os.path.exists(src) else 0
while True:
time.sleep(1)
try:
mtime = os.path.getmtime(src)
if mtime != last_mtime:
last_mtime = mtime
dest = os.path.join(_load_output_dir(), "car.html")
shutil.copy2(src, dest)
print(f"[car.html] change detected → synced to {dest}")
except Exception:
pass
@app.on_event("startup")
def on_startup():
try:
database.migrate_schema()
except Exception as e:
print(f"DB migration warning: {e}")
_sync_car_html()
threading.Thread(target=_watch_car_html, daemon=True).start()
# --- helpers ----------------------------------------------------------------
def _round16(x: int) -> int:
@@ -240,18 +278,6 @@ def _tags_to_name(tags: list, max_tags: int = 8) -> str:
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:
@@ -277,28 +303,41 @@ 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_beta = False
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()
poses[current_pose] = {"text": " ".join(current_desc).strip(), "beta": current_beta}
raw = line[2:].rstrip(":").strip()
current_beta = bool(re.search(r'\(beta\)', raw, re.IGNORECASE))
current_pose = re.sub(r'\s*\(beta\)\s*', '', raw, flags=re.IGNORECASE).strip()
current_desc = []
elif line and current_pose:
current_desc.append(line)
if current_pose:
poses[current_pose] = " ".join(current_desc).strip()
if current_pose:
poses[current_pose] = {"text": " ".join(current_desc).strip(), "beta": current_beta}
return poses
def _detect_has_background(pil: Image.Image) -> bool:
"""Return False when the image has significant alpha transparency (background removed)."""
if pil.mode != 'RGBA':
return True
alpha = pil.split()[3]
hist = alpha.histogram()
transparent_px = sum(hist[:128])
return transparent_px / (pil.width * pil.height) < 0.1
def _run_pipeline(
pil: Image.Image,
prompt: str,
@@ -308,6 +347,7 @@ def _run_pipeline(
cfg: float = 1.0,
sampler_name: str = "euler_ancestral",
scheduler: str = "beta",
extra_images: list = None, # additional PIL images wired to image2, image3
) -> bytes:
area = max_area if max_area > 0 else MAX_AREA
pil, w, h = _prep_image(pil, area)
@@ -319,12 +359,27 @@ def _run_pipeline(
graph = copy.deepcopy(BASE_WORKFLOW)
graph[NODE_LOADIMAGE]["inputs"]["image"] = stored
graph[NODE_POSITIVE]["inputs"]["prompt"] = prompt
# Inject extra reference images as image2 / image3 on the positive encoder
if extra_images:
for i, extra_pil in enumerate(extra_images[:2]):
extra_buf = io.BytesIO()
extra_pil.convert("RGB").save(extra_buf, format="PNG")
extra_stored = _comfy_upload(extra_buf.getvalue(), f"in_{uuid.uuid4().hex[:8]}.png")
node_id = str(11 + i) # "11" → image2, "12" → image3
img_key = f"image{i + 2}"
graph[node_id] = {
"class_type": "LoadImage",
"inputs": {"image": extra_stored},
"_meta": {"title": f"ref image {i + 2}"},
}
graph[NODE_POSITIVE]["inputs"][img_key] = [node_id, 0]
# 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_NEGATIVE]["inputs"]["prompt"] = "checkerboard, grid, pattern, texture, background details, watermark, deformed anatomy"
graph[NODE_LATENT]["inputs"]["width"] = w
graph[NODE_LATENT]["inputs"]["height"] = h
ks = graph[NODE_KSAMPLER]["inputs"]
@@ -333,10 +388,10 @@ def _run_pipeline(
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
@@ -371,23 +426,23 @@ def _move_to_trash(filepath: str):
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):
def _batch_worker(job_id: str, filenames: list, prompts: list[str], poses: list,
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}")
actual_gid = None
try:
person = database.get_person(fname)
# Prefer the source's existing DB group_id; fall back to the caller-supplied
# group_id (which is the gallery gid, potentially stale) or the basename.
if person and person[1]:
actual_gid = person[1]
else:
actual_gid = group_id or naming.get_base_name(fname)
database.upsert_person(fname, group_id=actual_gid)
except Exception as e:
print(f"Error determining/updating group for {fname}: {e}")
actual_gid = group_id or naming.get_base_name(fname)
fpath = os.path.join(output_dir, fname)
if not os.path.exists(fpath):
@@ -395,27 +450,41 @@ def _batch_worker(job_id: str, filenames: list, prompts: list[str], seed: int, m
continue
try:
pil = Image.open(fpath).convert("RGB")
for prompt in prompts:
base_pil = Image.open(fpath).convert("RGB")
for prompt, pose in zip(prompts, poses):
try:
pil = base_pil
# Rotate 180° for poses that work better upside-down
if pose and pose.lower().strip() in ROTATE_180_POSES:
pil = pil.rotate(180)
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
has_bg = True
try:
out_pil = Image.open(io.BytesIO(png))
has_bg = _detect_has_background(out_pil)
except Exception:
pass
try:
embedding = embeddings.generate_embedding(out_path)
database.upsert_person(out_name, filepath=out_path, embedding=embedding, group_id=actual_gid)
database.upsert_person(
out_name, filepath=out_path, embedding=embedding,
group_id=actual_gid, prompt=prompt, pose=pose,
has_background=has_bg,
source_refs=json.dumps([fname]),
)
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}")
@@ -423,7 +492,79 @@ def _batch_worker(job_id: str, filenames: list, prompts: list[str], seed: int, m
except Exception as e:
print(f"Error opening {fname}: {e}")
jobs[job_id]["failed"] += len(prompts)
jobs[job_id]["status"] = "done"
def _multi_ref_worker(job_id: str, filenames: list[str], prompts: list[str], poses: list,
seed: int, max_area: int):
"""Generate one output image per prompt using filenames[0] as primary and the rest as extra refs."""
output_dir = _load_output_dir()
pils = []
for fname in filenames:
fpath = os.path.join(output_dir, fname)
if os.path.exists(fpath):
pils.append((fname, Image.open(fpath).convert("RGB")))
if not pils:
jobs[job_id]["status"] = "done"
return
# Output group: reuse shared group if all sources belong to the same one, else new group
source_groups = set()
for fname, _ in pils:
try:
p = database.get_person(fname)
if p and p[1]:
source_groups.add(p[1])
except Exception:
pass
if len(source_groups) == 1:
output_gid = next(iter(source_groups))
else:
output_gid = f"cg_{uuid.uuid4().hex[:8]}"
primary_fname, primary_pil = pils[0]
extra_pils = [p for _, p in pils[1:]]
for prompt, pose in zip(prompts, poses):
try:
work_pil = primary_pil
if pose and pose.lower().strip() in ROTATE_180_POSES:
work_pil = work_pil.rotate(180)
png = _run_pipeline(work_pil, prompt, seed, max_area, extra_images=extra_pils)
ts = time.strftime("%Y%m%d_%H%M%S")
clean = naming.get_base_name(primary_fname)
out_name = f"{ts}_mr_{clean}"
out_path = os.path.join(output_dir, out_name)
with open(out_path, "wb") as f:
f.write(png)
has_bg = True
try:
out_pil = Image.open(io.BytesIO(png))
has_bg = _detect_has_background(out_pil)
except Exception:
pass
try:
embedding = embeddings.generate_embedding(out_path)
database.upsert_person(out_name, filepath=out_path, embedding=embedding,
group_id=output_gid, prompt=prompt, pose=pose,
has_background=has_bg,
source_refs=json.dumps([f for f, _ in pils]))
except Exception as db_err:
print(f"DB error in multi-ref: {db_err}")
jobs[job_id]["done"] += 1
except Exception as e:
print(f"Error in multi-ref for prompt '{prompt}': {e}")
jobs[job_id]["failed"] += 1
jobs[job_id]["status"] = "done"
@@ -459,24 +600,58 @@ class BatchRequest(BaseModel):
seed: int = -1
max_area: int = 0
group_id: str | None = None
poses: list[str | None] | None = None # pose name per prompt (same index), or None; None entries = no pose
@app.post("/batch")
def start_batch(req: BatchRequest):
prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt
poses = req.poses or [None] * len(prompts)
# Pad poses list to match prompts length
while len(poses) < len(prompts):
poses.append(None)
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),
args=(job_id, req.filenames, prompts, poses, req.seed, req.max_area, req.group_id),
daemon=True,
)
t.start()
return {"job_id": job_id, "total": total_tasks}
class MultiRefRequest(BaseModel):
filenames: list[str] # 23 reference images; first is primary (image1)
prompt: str | list[str]
poses: list[str | None] | None = None
seed: int = -1
max_area: int = 0
@app.post("/multi-ref")
def start_multi_ref(req: MultiRefRequest):
if len(req.filenames) < 2:
raise HTTPException(400, "multi-ref requires at least 2 filenames")
filenames = req.filenames[:3] # cap at 3 (image1/2/3)
prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt
poses = req.poses or [None] * len(prompts)
while len(poses) < len(prompts):
poses.append(None)
job_id = uuid.uuid4().hex[:8]
jobs[job_id] = {"status": "running", "total": len(prompts), "done": 0, "failed": 0}
t = threading.Thread(
target=_multi_ref_worker,
args=(job_id, filenames, prompts, poses, req.seed, req.max_area),
daemon=True,
)
t.start()
return {"job_id": job_id, "total": len(prompts)}
@app.get("/poses")
def get_poses():
return _load_poses()
@@ -497,17 +672,27 @@ def list_images():
# Try to get from DB first
try:
persons = database.list_persons()
# persons is (filename, name, group_id, clip_description)
# persons: (filename, name, group_id, clip_description, prompt, pose, sort_order, group_name, hidden, has_background, source_refs)
db_images = []
for p in persons:
db_images.append({
"filename": p[0],
"name": p[1],
"group_id": p[2],
"clip_description": p[3]
"filename": p[0],
"name": p[1],
"group_id": p[2],
"clip_description":p[3],
"prompt": p[4],
"pose": p[5],
"sort_order": p[6],
"group_name": p[7],
"hidden": bool(p[8]) if p[8] else False,
"has_background": bool(p[9]) if p[9] is not None else True,
"source_refs": p[10],
})
# 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)
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}")
@@ -546,14 +731,7 @@ def tag_image(req: TagRequest):
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))
@@ -564,8 +742,8 @@ 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)
except Exception as e:
raise HTTPException(500, str(e))
@app.post("/names/{filename}")
@@ -575,14 +753,7 @@ def set_name(filename: str, body: dict):
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}
@@ -593,8 +764,8 @@ 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)
except Exception as e:
raise HTTPException(500, str(e))
class MergeRequest(BaseModel):
@@ -610,15 +781,7 @@ def merge_groups(req: MergeRequest):
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}
@@ -633,17 +796,50 @@ def extract_from_group(req: ExtractRequest):
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("/group-names")
def get_group_names():
try:
return database.get_all_group_names()
except Exception as e:
raise HTTPException(500, str(e))
@app.post("/group-names/{group_id}")
def set_group_name(group_id: str, body: dict):
name = body.get("name", "").strip()
try:
database.set_group_name(group_id, name or None)
except Exception as e:
raise HTTPException(500, str(e))
return {"group_id": group_id, "name": name}
@app.get("/groups/{group_id}/order")
def get_group_order(group_id: str):
try:
rows = database.get_group_order(group_id)
return {"group_id": group_id, "filenames": [r[0] for r in rows]}
except Exception as e:
raise HTTPException(500, str(e))
class GroupOrderRequest(BaseModel):
filenames: list[str]
@app.post("/groups/{group_id}/order")
def set_group_order(group_id: str, req: GroupOrderRequest):
try:
database.set_group_order(group_id, req.filenames)
except Exception as e:
raise HTTPException(500, str(e))
return {"group_id": group_id, "filenames": req.filenames}
@app.get("/similar/{filename}")
def get_similar(filename: str, limit: int = 10):
person = database.get_person(filename)
@@ -785,7 +981,7 @@ def upload_image(
# Use default prompt from config
prompt_list = [conf.get("prompt", "high quality, masterpiece")]
group_id = naming.get_base_name(filename)
group_id = f"up_{uuid.uuid4().hex[:8]}" # unique per upload; avoids collisions when pasting generic filenames
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}
@@ -812,6 +1008,31 @@ async def edit(
return Response(content=png, media_type="image/png")
@app.post("/images/{filename}/hidden")
def set_image_hidden(filename: str, body: dict):
hidden = bool(body.get("hidden", False))
try:
database.set_hidden(filename, hidden)
except Exception as e:
raise HTTPException(500, str(e))
return {"filename": filename, "hidden": hidden}
@app.post("/images/{filename}/set-preferred")
def set_image_preferred(filename: str):
"""Make this image sort_order=0 within its group, shifting others to 1,2,..."""
person = database.get_person(filename)
if not person:
raise HTTPException(404, "Image not found")
group_id = person[1]
if not group_id:
raise HTTPException(400, "Image has no group assigned")
rows = database.get_group_order(group_id)
others = [r[0] for r in rows if r[0] != filename]
database.set_group_order(group_id, [filename] + others)
return {"filename": filename, "group_id": group_id}
@app.delete("/images/{filename}")
def delete_image(filename: str):
person = database.get_person(filename)