aa
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user