This commit is contained in:
mike
2026-06-20 02:43:57 +02:00
parent f1f693523b
commit 0176c1b965
16 changed files with 3820 additions and 394 deletions

View File

@@ -62,6 +62,20 @@ 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"}
# WD tagger tags that indicate the subject is wearing clothes
CLOTHING_TAGS = {
"dress", "skirt", "shirt", "top", "pants", "jeans", "jacket", "coat",
"swimsuit", "bikini", "shorts", "uniform", "hoodie", "sweater", "blouse",
"leggings", "stockings", "tights", "lingerie", "miniskirt", "pleated_skirt",
"school_uniform", "maid_dress", "bodysuit", "sailor_uniform", "leotard",
"corset", "kimono", "yukata", "cheongsam", "t-shirt", "tank_top", "crop_top",
"tube_top", "halter_top", "negligee", "nightgown", "pajamas", "trench_coat",
"vest", "bra", "underwear", "panties", "thong", "g-string", "bikini_top",
"bikini_bottom", "one-piece_swimsuit", "sports_bra", "gym_clothes",
}
UNDRESS_PROMPT = "completely nude, bare skin, no clothing, naked body, natural skin texture"
with open(WORKFLOW_PATH, "r", encoding="utf-8") as f:
BASE_WORKFLOW = json.load(f)
@@ -338,6 +352,14 @@ def _detect_has_background(pil: Image.Image) -> bool:
return transparent_px / (pil.width * pil.height) < 0.1
def _detect_has_clothing(tags: list) -> bool | None:
"""Return True if any tag from CLOTHING_TAGS appears above threshold, None if no tags."""
if not tags:
return None
tag_names = {t["tag"] for t in tags}
return bool(tag_names & CLOTHING_TAGS)
def _run_pipeline(
pil: Image.Image,
prompt: str,
@@ -476,10 +498,11 @@ def _batch_worker(job_id: str, filenames: list, prompts: list[str], poses: list,
try:
embedding = embeddings.generate_embedding(out_path)
next_order = database.get_next_sort_order(actual_gid)
database.upsert_person(
out_name, filepath=out_path, embedding=embedding,
group_id=actual_gid, prompt=prompt, pose=pose,
has_background=has_bg,
has_background=has_bg, sort_order=next_order,
source_refs=json.dumps([fname]),
)
except Exception as db_err:
@@ -553,9 +576,10 @@ def _multi_ref_worker(job_id: str, filenames: list[str], prompts: list[str], pos
try:
embedding = embeddings.generate_embedding(out_path)
next_order = database.get_next_sort_order(output_gid)
database.upsert_person(out_name, filepath=out_path, embedding=embedding,
group_id=output_gid, prompt=prompt, pose=pose,
has_background=has_bg,
has_background=has_bg, sort_order=next_order,
source_refs=json.dumps([f for f, _ in pils]))
except Exception as db_err:
print(f"DB error in multi-ref: {db_err}")
@@ -672,9 +696,12 @@ def list_images():
# Try to get from DB first
try:
persons = database.list_persons()
# persons: (filename, name, group_id, clip_description, prompt, pose, sort_order, group_name, hidden, has_background, source_refs)
# persons: (filename, name, group_id, clip_description, prompt, pose, sort_order, group_name, hidden, has_background, source_refs, has_clothing)
db_images = []
for p in persons:
fpath = os.path.join(output_dir, p[0])
if not os.path.exists(fpath):
continue # skip orphan DB records whose file no longer exists
db_images.append({
"filename": p[0],
"name": p[1],
@@ -687,10 +714,10 @@ def list_images():
"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],
"has_clothing": p[11],
})
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,
key=lambda x: os.path.getmtime(os.path.join(output_dir, x["filename"])),
reverse=True,
)
return {"images": db_images}
@@ -723,16 +750,24 @@ def tag_image(req: TagRequest):
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)
has_clothing = _detect_has_clothing(tags)
# Only assign a new name if the image doesn't already have one
existing = database.get_person(req.filename)
auto_name = (existing[0] if existing and existing[0] else None) or 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)
database.upsert_person(
req.filename, filepath=fpath, name=auto_name,
clip_description=clip_desc, tags=tags, embedding=embedding,
group_id=req.group_id, has_clothing=has_clothing,
)
except Exception as db_err:
print(f"Database error during tag: {db_err}")
return {"filename": req.filename, "clip_description": clip_desc, "tags": tags[:30]}
return {"filename": req.filename, "clip_description": clip_desc, "tags": tags[:30], "has_clothing": has_clothing}
except Exception as e:
raise HTTPException(500, str(e))
@@ -865,6 +900,20 @@ def get_similar(filename: str, limit: int = 10):
return {"filename": filename, "similar": similar}
@app.post("/db/cleanup")
def db_cleanup():
"""Delete DB records for files that no longer exist on disk."""
output_dir = _load_output_dir()
persons = database.list_persons()
removed = []
for p in persons:
fpath = os.path.join(output_dir, p[0])
if not os.path.exists(fpath):
database.delete_person(p[0])
removed.append(p[0])
return {"removed": len(removed), "filenames": removed}
@app.get("/health")
def health():
try:
@@ -907,37 +956,43 @@ def _process_upload(file_path: str, filename: str, prompts: list[str], name: str
# 1. CLIP tag the source
tags = _run_tagger(pil.convert("RGB"))
clip_desc = _tags_to_name(tags)
has_clothing = _detect_has_clothing(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
# 3. Register source in DB — sort_order=0 makes it the preferred base image
database.upsert_person(
filename, filepath=file_path, name=auto_name,
clip_description=clip_desc, tags=tags, embedding=embedding,
group_id=group_id, sort_order=0, has_clothing=has_clothing,
)
# 4. Crop if needed
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)
next_order = database.get_next_sort_order(group_id)
database.upsert_person(
out_name, filepath=out_path, name=auto_name,
clip_description=clip_desc, embedding=out_embedding,
group_id=group_id, sort_order=next_order,
)
except Exception as e:
print(f"Error processing prompt '{prompt}' for {filename}: {e}")
@@ -1033,6 +1088,27 @@ def set_image_preferred(filename: str):
return {"filename": filename, "group_id": group_id}
@app.post("/images/{filename}/undress")
def undress_image(filename: str, background_tasks: BackgroundTasks):
"""Queue a generation using the undress prompt on the given image."""
output_dir = _load_output_dir()
fpath = os.path.join(output_dir, filename)
if not os.path.exists(fpath):
raise HTTPException(404, "Image not found")
person = database.get_person(filename)
group_id = person[1] if person and person[1] else naming.get_base_name(filename)
job_id = uuid.uuid4().hex[:8]
jobs[job_id] = {"status": "queued", "done": 0, "failed": 0, "total": 1}
threading.Thread(
target=_batch_worker,
args=(job_id, [filename], [UNDRESS_PROMPT], [None],
random.randint(0, MAX_SEED), MAX_AREA),
kwargs={"group_id": group_id},
daemon=True,
).start()
return {"job_id": job_id, "filename": filename}
@app.delete("/images/{filename}")
def delete_image(filename: str):
person = database.get_person(filename)