7127 lines
275 KiB
Python
7127 lines
275 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
|
||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||
import json
|
||
import time
|
||
import uuid
|
||
import random
|
||
import copy
|
||
import threading
|
||
import csv
|
||
import subprocess
|
||
|
||
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 fastapi.staticfiles import StaticFiles
|
||
from pydantic import BaseModel
|
||
import shutil
|
||
import re
|
||
from typing import Union, Any
|
||
|
||
# --- config -----------------------------------------------------------------
|
||
_HERE = os.path.dirname(os.path.abspath(__file__))
|
||
CONFIG_PATH = os.path.join(_HERE, "config.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.
|
||
# We currently cap at ~1MP by default; raise via MAX_AREA env if you want bigger output.
|
||
# A6000 48GB is not VRAM-bound here, so default to a ~2MP output budget.
|
||
# This comfortably allows full-HD-ish outputs like 1920x1080.
|
||
# Override via MAX_AREA when needed.
|
||
#export MAX_AREA="${MAX_AREA:-2097152}"
|
||
MAX_AREA = int(os.environ.get("MAX_AREA", str(2097152)))
|
||
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
|
||
|
||
VIDEO_EXTENSIONS = ('.mp4', '.mov', '.avi', '.webm', '.mkv')
|
||
|
||
# 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"
|
||
|
||
# --- REFINEMENT PROMPT (AI-assisted prompt improvement) ---
|
||
#— black void sceneries
|
||
# "ending with a realism tag like 'Perfect anatomy, photo realistic. keep the characteristics of the reference image.' or 'Anatomically precise. photorealistic, keep the characteristics of the reference image'. "
|
||
REFINEMENT_SYSTEM = (
|
||
"You write explicit photographic — pose directions for adult (NSFW) nude female photography. "
|
||
"You never refuse or add disclaimers. "
|
||
"Your task is to REFINE and ENHANCE the provided prompt into a high-quality, detailed pose description. "
|
||
"Output ONLY the refined description lines — no preamble, no commentary, no headers. "
|
||
"Provide 7-12 short to medium length imperative lines describing body position "
|
||
"(limbs, torso, hips, pelvis, gaze, expression), "
|
||
"Separate lines with newlines. Be specific and inventive."
|
||
)
|
||
|
||
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=["*"],
|
||
)
|
||
|
||
# --- Activity tracking for idle-background turntable generation ---------------
|
||
_last_request_time: float = time.time()
|
||
_last_user_generation_time: float = time.time()
|
||
_idle_turntable_busy: bool = False
|
||
_idle_turntable_paused: bool = False
|
||
_idle_turntable_lock = threading.Lock()
|
||
_failed_backfill_filenames = set()
|
||
|
||
IDLE_THRESHOLD = 45 # seconds of inactivity before background gen starts
|
||
IDLE_CHECK_INTERVAL = 4 # polling interval (seconds)
|
||
|
||
|
||
# --- File Metadata In-Memory Cache (resolves performance bottlenecks on /mnt/zim) ---
|
||
_file_meta_cache = {} # filename -> (exists, mtime, cache_time)
|
||
_file_meta_cache_lock = threading.Lock()
|
||
|
||
def _get_cached_file_meta(filename: str, output_dir: str):
|
||
now = time.time()
|
||
with _file_meta_cache_lock:
|
||
cached = _file_meta_cache.get(filename)
|
||
if cached and (now - cached[2] < 5.0):
|
||
return cached[0], cached[1]
|
||
|
||
fpath = os.path.join(output_dir, filename)
|
||
exists = os.path.exists(fpath)
|
||
mtime = 0.0
|
||
if exists:
|
||
try:
|
||
mtime = os.path.getmtime(fpath)
|
||
except Exception:
|
||
pass
|
||
with _file_meta_cache_lock:
|
||
_file_meta_cache[filename] = (exists, mtime, now)
|
||
return exists, mtime
|
||
|
||
def _update_cached_file_meta(filename: str, exists: bool = True, mtime: float = None):
|
||
if mtime is None:
|
||
mtime = time.time()
|
||
with _file_meta_cache_lock:
|
||
_file_meta_cache[filename] = (exists, mtime, time.time())
|
||
|
||
def _clear_cached_file_meta(filename: str):
|
||
with _file_meta_cache_lock:
|
||
if filename in _file_meta_cache:
|
||
del _file_meta_cache[filename]
|
||
|
||
|
||
_last_preloaded_images_set = None
|
||
|
||
@app.middleware("http")
|
||
async def _track_activity(request, call_next):
|
||
global _last_request_time
|
||
_last_request_time = time.time()
|
||
return await call_next(request)
|
||
|
||
def _sync_preloaded_images():
|
||
"""Update PRELOADED_IMAGES in car.html (both source and output) to match current DB state."""
|
||
global _last_preloaded_images_set
|
||
try:
|
||
output_dir = _load_output_dir()
|
||
paths = [
|
||
os.path.join(_HERE, "car.html"),
|
||
os.path.join(output_dir, "car.html")
|
||
]
|
||
|
||
persons = database.list_persons(include_archived=False)
|
||
# Only include if file actually exists on disk, using the fast cache
|
||
db_images = []
|
||
for p in persons:
|
||
exists, mtime = _get_cached_file_meta(p[0], output_dir)
|
||
if exists:
|
||
db_images.append(p[0])
|
||
|
||
# Sort by mtime, newest first
|
||
db_images.sort(key=lambda x: _get_cached_file_meta(x, output_dir)[1], reverse=True)
|
||
|
||
# Avoid redundant HTML rewrites if the preloaded set is unchanged
|
||
current_set = set(db_images)
|
||
if _last_preloaded_images_set is not None and _last_preloaded_images_set == current_set:
|
||
return
|
||
_last_preloaded_images_set = current_set
|
||
|
||
images_json = json.dumps(db_images, indent=12).strip()
|
||
# Format for JS insertion
|
||
images_json = images_json.replace('\n', '\n ')
|
||
pattern = r'// --- HYDRATION_START ---.*?// --- HYDRATION_END ---'
|
||
replacement = f'// --- HYDRATION_START ---\n const PRELOADED_IMAGES = {images_json};\n // --- HYDRATION_END ---'
|
||
|
||
for p in paths:
|
||
if not os.path.exists(p): continue
|
||
with open(p, 'r') as f:
|
||
content = f.read()
|
||
new_content = re.sub(pattern, replacement, content, flags=re.DOTALL)
|
||
with open(p, 'w') as f:
|
||
f.write(new_content)
|
||
print(f"[static] Updated {p} with {len(db_images)} preloaded images")
|
||
except Exception as e:
|
||
print(f"[static] Failed to update car.html hydration: {e}")
|
||
|
||
def _sync_frontend():
|
||
for name in ["car.html", "trash.html"]:
|
||
src = os.path.join(_HERE, name)
|
||
if not os.path.exists(src):
|
||
continue
|
||
try:
|
||
dest = os.path.join(_load_output_dir(), name)
|
||
shutil.copy2(src, dest)
|
||
print(f"[{name}] synced → {dest}")
|
||
except Exception as e:
|
||
print(f"[{name}] sync warning: {e}")
|
||
|
||
def _watch_frontend():
|
||
files = ["car.html", "trash.html"]
|
||
last_mtimes = {}
|
||
for name in files:
|
||
src = os.path.join(_HERE, name)
|
||
if os.path.exists(src):
|
||
last_mtimes[name] = os.path.getmtime(src)
|
||
|
||
while True:
|
||
time.sleep(1)
|
||
for name in files:
|
||
src = os.path.join(_HERE, name)
|
||
if not os.path.exists(src): continue
|
||
try:
|
||
mtime = os.path.getmtime(src)
|
||
if mtime != last_mtimes.get(name):
|
||
last_mtimes[name] = mtime
|
||
dest = os.path.join(_load_output_dir(), name)
|
||
shutil.copy2(src, dest)
|
||
print(f"[{name}] change detected → synced to {dest}")
|
||
except Exception:
|
||
pass
|
||
|
||
def _load_wireframe_dir() -> str:
|
||
with open(CONFIG_PATH, "r") as f:
|
||
conf = json.load(f)
|
||
d = conf.get("wireframe_dir", "/mnt/zim/tour-comfy/wireframe")
|
||
return os.path.expanduser(d)
|
||
|
||
|
||
def _load_faceswap_model_path() -> str:
|
||
with open(CONFIG_PATH, "r") as f:
|
||
conf = json.load(f)
|
||
return os.path.expanduser(conf.get("faceswap_model", "~/.insightface/models/inswapper_128.onnx"))
|
||
|
||
|
||
def _run_consistency_check():
|
||
"""Identifies DB records with missing files, files on disk with no DB record, and archived items."""
|
||
try:
|
||
output_dir = _load_output_dir()
|
||
data_dir = os.path.join(output_dir, "_data")
|
||
os.makedirs(data_dir, exist_ok=True)
|
||
|
||
persons = database.list_persons(include_archived=True)
|
||
missing_files = []
|
||
archived_items = []
|
||
missing_group = []
|
||
for p in persons:
|
||
filename = p[0]
|
||
group_id = p[2]
|
||
if not group_id:
|
||
missing_group.append({
|
||
"filename": filename,
|
||
"name": p[1]
|
||
})
|
||
|
||
is_archived = bool(p[14]) if p[14] else False
|
||
if is_archived:
|
||
archived_items.append({
|
||
"filename": filename,
|
||
"group_id": group_id,
|
||
"name": p[1]
|
||
})
|
||
|
||
fpath = os.path.join(output_dir, filename)
|
||
if not os.path.exists(fpath):
|
||
missing_files.append({
|
||
"filename": filename,
|
||
"group_id": group_id,
|
||
"name": p[1]
|
||
})
|
||
|
||
# Untracked files
|
||
db_filenames = set(p[0] for p in persons)
|
||
untracked_files = []
|
||
if os.path.isdir(output_dir):
|
||
for f in os.listdir(output_dir):
|
||
if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp', '.mp4')):
|
||
# Skip internal folders and data files
|
||
if f.startswith('_') or f == 'car.html' or f == 'trash.html':
|
||
continue
|
||
if f not in db_filenames:
|
||
untracked_files.append(f)
|
||
|
||
report = {
|
||
"timestamp": time.time(),
|
||
"missing_files": missing_files,
|
||
"untracked_files": untracked_files,
|
||
"archived_items": archived_items,
|
||
"missing_group": missing_group
|
||
}
|
||
_write_json(os.path.join(data_dir, "inconsistencies.json"), report)
|
||
print(f"[consistency] check complete: {len(missing_files)} missing, {len(untracked_files)} untracked, {len(archived_items)} archived, {len(missing_group)} missing group")
|
||
return report
|
||
except Exception as e:
|
||
print(f"[consistency] check failed: {e}")
|
||
return None
|
||
|
||
|
||
def _consistency_check_daemon():
|
||
"""Runs daily consistency check."""
|
||
# Wait for startup
|
||
time.sleep(30)
|
||
while True:
|
||
_run_consistency_check()
|
||
# Sleep for 24 hours
|
||
time.sleep(86400)
|
||
|
||
|
||
def _idle_turntable_daemon():
|
||
"""
|
||
Background daemon: when the API has been idle > IDLE_THRESHOLD seconds,
|
||
generate the next missing turntable view for the next group. Yields after
|
||
each view so activity check can stop it promptly.
|
||
"""
|
||
global _idle_turntable_busy
|
||
import sys as _sys
|
||
if _HERE not in _sys.path:
|
||
_sys.path.insert(0, _HERE)
|
||
|
||
time.sleep(60) # wait for full startup before touching ComfyUI
|
||
|
||
while True:
|
||
time.sleep(IDLE_CHECK_INTERVAL)
|
||
|
||
if _idle_turntable_paused:
|
||
continue
|
||
if time.time() - _last_user_generation_time < IDLE_THRESHOLD:
|
||
continue
|
||
|
||
try:
|
||
output_dir = _load_output_dir()
|
||
persons = database.list_persons()
|
||
except Exception as e:
|
||
print(f"[turntable-bg] db/config error: {e}")
|
||
continue
|
||
|
||
# Build {group_id: (preferred_filename, best_sort_order)}
|
||
groups: dict = {}
|
||
for row in persons:
|
||
fname, group_id, sort_order = row[0], row[2], row[6]
|
||
if not group_id:
|
||
continue
|
||
if fname.startswith("_turntable/"):
|
||
continue
|
||
if group_id not in groups:
|
||
groups[group_id] = (fname, sort_order)
|
||
else:
|
||
cur_sort = groups[group_id][1]
|
||
if sort_order is not None and (cur_sort is None or sort_order < cur_sort):
|
||
groups[group_id] = (fname, sort_order)
|
||
|
||
import turntable_cache as tc
|
||
|
||
generated_one = False
|
||
for group_id, (preferred_fname, _) in groups.items():
|
||
src_path = os.path.join(output_dir, preferred_fname)
|
||
if not os.path.exists(src_path):
|
||
continue
|
||
|
||
state = tc.load_state(output_dir, group_id)
|
||
|
||
if state is None:
|
||
state = tc.init_state(output_dir, group_id, src_path, preferred_fname)
|
||
elif state.get("completed"):
|
||
continue
|
||
elif state.get("preferred_filename") != preferred_fname:
|
||
# Preferred image changed — restart turntable for this group
|
||
print(f"[turntable-bg] {group_id}: preferred changed, resetting")
|
||
state = tc.init_state(output_dir, group_id, src_path, preferred_fname)
|
||
|
||
deg = tc.next_missing_angle(state)
|
||
if deg is None:
|
||
continue # all angles done but no video yet — build it below
|
||
|
||
# Re-check idle right before the expensive generation
|
||
if time.time() - _last_user_generation_time < IDLE_THRESHOLD or _idle_turntable_paused:
|
||
break
|
||
|
||
print(f"[turntable-bg] {group_id}: rendering {deg:.0f}° "
|
||
f"({len(state['views'])}/{state['n_views']})…")
|
||
|
||
with _idle_turntable_lock:
|
||
_idle_turntable_busy = True
|
||
_write_turntable_static()
|
||
|
||
try:
|
||
from orbit_qwen import yaw_prompt, _autocrop_alpha
|
||
import io as _io
|
||
from PIL import Image as _Image
|
||
|
||
base_pil = _Image.open(src_path).convert("RGB")
|
||
prompt = yaw_prompt(deg)
|
||
png = _run_pipeline(
|
||
base_pil, prompt, state["seed"], MAX_AREA, steps=state["steps"]
|
||
)
|
||
view_pil = _Image.open(_io.BytesIO(png)).convert("RGBA")
|
||
view_pil = _autocrop_alpha(view_pil)
|
||
|
||
views_dir = os.path.join(tc.cache_dir(output_dir, group_id), "views")
|
||
os.makedirs(views_dir, exist_ok=True)
|
||
angle_idx = state["angles"].index(deg)
|
||
vpath = os.path.join(views_dir, f"view_{angle_idx:03d}_{int(deg):03d}deg.png")
|
||
view_pil.save(vpath)
|
||
|
||
tc.mark_view_done(output_dir, group_id, state, deg, vpath)
|
||
n_done = len(state["views"])
|
||
print(f"[turntable-bg] {group_id}: {deg:.0f}° saved "
|
||
f"({n_done}/{state['n_views']})")
|
||
|
||
# Register frame in database for filmstrip visibility + Info tab links
|
||
try:
|
||
vname = os.path.relpath(vpath, output_dir).replace("\\", "/")
|
||
database.upsert_person(
|
||
vname,
|
||
filepath=vpath,
|
||
group_id=group_id,
|
||
prompt=prompt,
|
||
source_refs=json.dumps([preferred_fname]),
|
||
sort_order=200 + angle_idx,
|
||
pose=f"Orbit {int(deg)}°",
|
||
tags=["ORBIT"]
|
||
)
|
||
_update_cached_file_meta(vname, exists=True)
|
||
except Exception as db_err:
|
||
print(f"[turntable-bg] DB registration error: {db_err}")
|
||
|
||
generated_one = True
|
||
_write_turntable_static()
|
||
|
||
if n_done >= state["n_views"]:
|
||
_finalize_turntable(output_dir, group_id, state)
|
||
|
||
except Exception as e:
|
||
import traceback
|
||
print(f"[turntable-bg] {group_id}: error at {deg:.0f}°: {e}")
|
||
print(traceback.format_exc())
|
||
finally:
|
||
with _idle_turntable_lock:
|
||
_idle_turntable_busy = False
|
||
_write_turntable_static()
|
||
|
||
break # one view per cycle; re-check idle on next loop
|
||
|
||
# If no turntable views were generated, check if any legacy image needs metadata backfill
|
||
if not generated_one:
|
||
legacy_candidate = None
|
||
for row in persons:
|
||
fname = row[0]
|
||
content_type = row[12] if len(row) > 12 else None
|
||
people_count = row[19] if len(row) > 19 else None
|
||
if fname.startswith("_turntable/"):
|
||
continue
|
||
if content_type == 'video':
|
||
continue
|
||
if not fname.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
|
||
continue
|
||
fpath = os.path.join(output_dir, fname)
|
||
if not os.path.exists(fpath):
|
||
continue
|
||
if people_count is None and fname not in _failed_backfill_filenames:
|
||
legacy_candidate = fname
|
||
break
|
||
|
||
if legacy_candidate:
|
||
print(f"[metadata-bg] Idle backfill: processing {legacy_candidate}…")
|
||
try:
|
||
res = _process_image_for_metadata(legacy_candidate)
|
||
if res is None:
|
||
_failed_backfill_filenames.add(legacy_candidate)
|
||
except Exception as ex:
|
||
print(f"[metadata-bg] Error processing legacy backfill for {legacy_candidate}: {ex}")
|
||
_failed_backfill_filenames.add(legacy_candidate)
|
||
|
||
|
||
def _finalize_turntable(output_dir: str, group_id: str, state: dict):
|
||
"""Mark state completed (without building the MP4)."""
|
||
import turntable_cache as tc
|
||
try:
|
||
tc.mark_completed(output_dir, group_id, state, "")
|
||
print(f"[turntable-bg] {group_id}: complete (custom frame-loop only)")
|
||
_write_turntable_static()
|
||
except Exception as e:
|
||
import traceback
|
||
print(f"[turntable-bg] {group_id}: finalize error: {e}\n{traceback.format_exc()}")
|
||
|
||
|
||
@app.on_event("startup")
|
||
def on_startup():
|
||
try:
|
||
database.migrate_schema()
|
||
except Exception as e:
|
||
print(f"DB migration warning: {e}")
|
||
_sync_frontend()
|
||
threading.Thread(target=_watch_frontend, daemon=True).start()
|
||
threading.Thread(target=_idle_turntable_daemon, daemon=True).start()
|
||
threading.Thread(target=_consistency_check_daemon, daemon=True).start()
|
||
threading.Thread(target=_privacy_monitor_daemon, daemon=True).start()
|
||
# Mount wireframe static dir for browser video preview
|
||
try:
|
||
wf_dir = _load_wireframe_dir()
|
||
if os.path.isdir(wf_dir):
|
||
app.mount("/wireframe", StaticFiles(directory=wf_dir), name="wireframe")
|
||
print(f"[wireframe] mounted {wf_dir} → /wireframe")
|
||
except Exception as e:
|
||
print(f"[wireframe] mount warning: {e}")
|
||
# Mount output dir so images can be served via HTTP (/output/filename.png)
|
||
try:
|
||
out_dir = _load_output_dir()
|
||
if os.path.isdir(out_dir):
|
||
app.mount("/output", StaticFiles(directory=out_dir), name="output")
|
||
print(f"[output] mounted {out_dir} → /output")
|
||
except Exception as e:
|
||
print(f"[output] mount warning: {e}")
|
||
# Write initial static data files (synchronous — ensures files exist before first request)
|
||
_write_all_static()
|
||
# Trigger pose index backfill
|
||
try:
|
||
if _load_pose_estimator():
|
||
build_pose_index()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
# --- 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. Scale (up or down) to fit area while preserving aspect.
|
||
2. Ensure dimensions are rounded to 16.
|
||
"""
|
||
w, h = pil.width, pil.height
|
||
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."""
|
||
global _last_user_generation_time
|
||
while time.time() < deadline:
|
||
if not _idle_turntable_busy:
|
||
_last_user_generation_time = time.time()
|
||
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", local_files_only=True)
|
||
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()
|
||
with embeddings._gpu_lock:
|
||
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 _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
|
||
|
||
|
||
# --- faceswapper (insightface + inswapper_128) --------------------------------
|
||
# Setup: pip install insightface onnxruntime-gpu opencv-python-headless
|
||
# Download: place inswapper_128.onnx at ~/.insightface/models/inswapper_128.onnx
|
||
# Source: https://huggingface.co/deepinsight/inswapper
|
||
|
||
_faceswapper = None
|
||
_faceswapper_lock = threading.Lock()
|
||
|
||
# Dedicated single-worker pool for face-crop extraction. Running it here
|
||
# instead of via FastAPI BackgroundTasks keeps the heavy insightface inference
|
||
# (and its one-time model load) off the shared request threadpool, so quick
|
||
# endpoints like /order stay responsive right after a "set preferred" click.
|
||
# A single worker also serializes face jobs so a burst can't thrash the GPU.
|
||
from concurrent.futures import ThreadPoolExecutor as _ThreadPoolExecutor
|
||
_face_executor = _ThreadPoolExecutor(max_workers=1, thread_name_prefix="face")
|
||
|
||
_gfpgan = None
|
||
_gfpgan_lock = threading.Lock()
|
||
|
||
|
||
def _load_faceswapper():
|
||
global _faceswapper
|
||
if _faceswapper is not None:
|
||
return _faceswapper
|
||
with _faceswapper_lock:
|
||
if _faceswapper is not None:
|
||
return _faceswapper
|
||
try:
|
||
import insightface
|
||
from insightface.app import FaceAnalysis
|
||
except ImportError:
|
||
raise RuntimeError("insightface not installed. Run: pip install insightface onnxruntime-gpu")
|
||
|
||
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
|
||
app = FaceAnalysis(name='buffalo_l', providers=providers)
|
||
app.prepare(ctx_id=0, det_size=(640, 640))
|
||
|
||
model_path = _load_faceswap_model_path()
|
||
if not os.path.exists(model_path):
|
||
# Try HuggingFace download as fallback
|
||
try:
|
||
import huggingface_hub
|
||
model_path = huggingface_hub.hf_hub_download(
|
||
'deepinsight/inswapper', 'inswapper_128.onnx',
|
||
local_dir=os.path.dirname(model_path),
|
||
local_files_only=True
|
||
)
|
||
print(f"[faceswap] Downloaded inswapper_128.onnx to {model_path}")
|
||
except Exception as de:
|
||
raise RuntimeError(
|
||
f"inswapper_128.onnx not found at {model_path}. "
|
||
f"Download from https://huggingface.co/deepinsight/inswapper and place it there. "
|
||
f"Download error: {de}"
|
||
)
|
||
|
||
swapper = insightface.model_zoo.get_model(model_path, providers=providers)
|
||
_faceswapper = (app, swapper)
|
||
print(f"[faceswap] loaded insightface buffalo_l + inswapper_128 from {model_path}")
|
||
return _faceswapper
|
||
|
||
|
||
def _load_gfpgan():
|
||
"""Lazy-load GFPGAN face restorer. Returns restorer or False if unavailable."""
|
||
global _gfpgan
|
||
if _gfpgan is not None:
|
||
return _gfpgan
|
||
with _gfpgan_lock:
|
||
if _gfpgan is not None:
|
||
return _gfpgan
|
||
try:
|
||
from gfpgan import GFPGANer
|
||
# Main GFPGAN model
|
||
model_path = os.path.expanduser('~/.gfpgan/weights/GFPGANv1.4.pth')
|
||
os.makedirs(os.path.dirname(model_path), exist_ok=True)
|
||
if not os.path.exists(model_path):
|
||
import urllib.request
|
||
print('[gfpgan] Downloading GFPGANv1.4.pth (~333 MB)...')
|
||
tmp = model_path + '.tmp'
|
||
urllib.request.urlretrieve(
|
||
'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth',
|
||
tmp
|
||
)
|
||
os.rename(tmp, model_path)
|
||
# GFPGANer hardcodes facexlib download path to CWD/gfpgan/weights/
|
||
# → change CWD to ~ so models land at ~/gfpgan/weights/ (stable across runs)
|
||
home = os.path.expanduser('~')
|
||
os.makedirs(os.path.join(home, 'gfpgan', 'weights'), exist_ok=True)
|
||
_orig_cwd = os.getcwd()
|
||
os.chdir(home)
|
||
try:
|
||
restorer = GFPGANer(model_path=model_path, upscale=1, arch='clean',
|
||
channel_multiplier=2, bg_upsampler=None)
|
||
finally:
|
||
os.chdir(_orig_cwd)
|
||
_gfpgan = restorer
|
||
print('[gfpgan] GFPGANv1.4 loaded')
|
||
except Exception as e:
|
||
print(f'[gfpgan] not available: {e}')
|
||
_gfpgan = False
|
||
return _gfpgan
|
||
|
||
|
||
def _make_video_poster(video_path: str) -> str | None:
|
||
"""Extract a poster JPG (sibling `<stem>.jpg`) so the gallery can show a
|
||
thumbnail for a video via a plain <img> (file:// can't render <video> as a
|
||
thumb). Returns the poster path on success, else None."""
|
||
import subprocess
|
||
poster_path = os.path.splitext(video_path)[0] + '.jpg'
|
||
try:
|
||
r = subprocess.run([
|
||
'ffmpeg', '-y', '-ss', '1', '-i', video_path,
|
||
'-frames:v', '1', '-q:v', '3', poster_path,
|
||
], capture_output=True, timeout=120)
|
||
if r.returncode == 0 and os.path.exists(poster_path):
|
||
return poster_path
|
||
# -ss 1 can overshoot very short clips; retry from the first frame
|
||
r = subprocess.run([
|
||
'ffmpeg', '-y', '-i', video_path,
|
||
'-frames:v', '1', '-q:v', '3', poster_path,
|
||
], capture_output=True, timeout=120)
|
||
if r.returncode == 0 and os.path.exists(poster_path):
|
||
return poster_path
|
||
except Exception as pe:
|
||
print(f'[poster] failed for {video_path}: {pe}')
|
||
return None
|
||
|
||
|
||
def _faceswap_worker(job_id: str, model_filename: str, video_name: str, enhance: bool = True, preview_scale: float = 1.0):
|
||
"""Frame-by-frame face swap: model face → every face in template video."""
|
||
output_dir = _load_output_dir()
|
||
wireframe_dir = _load_wireframe_dir()
|
||
|
||
try:
|
||
import cv2
|
||
import numpy as np
|
||
app, swapper = _load_faceswapper()
|
||
except Exception as e:
|
||
jobs[job_id]["status"] = "error"
|
||
jobs[job_id]["error"] = str(e)
|
||
return
|
||
|
||
gfpgan_restorer = None
|
||
if enhance:
|
||
try:
|
||
r = _load_gfpgan()
|
||
if r is not False:
|
||
gfpgan_restorer = r
|
||
except Exception:
|
||
pass
|
||
|
||
try:
|
||
# 1. Load source (model) face
|
||
src_path = os.path.join(output_dir, model_filename)
|
||
src_bgr = cv2.imread(src_path)
|
||
if src_bgr is None:
|
||
raise ValueError(f"Cannot read model image: {model_filename}")
|
||
with embeddings._gpu_lock:
|
||
src_faces = app.get(src_bgr)
|
||
if not src_faces:
|
||
raise ValueError(f"No face detected in: {model_filename}")
|
||
# Use the largest face as source
|
||
src_face = max(src_faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
|
||
|
||
# 2. Open template video
|
||
video_path = os.path.join(wireframe_dir, video_name)
|
||
cap = cv2.VideoCapture(video_path)
|
||
if not cap.isOpened():
|
||
raise ValueError(f"Cannot open video: {video_name}")
|
||
|
||
fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
|
||
scale = max(0.1, min(1.0, preview_scale))
|
||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) * scale)
|
||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) * scale)
|
||
# Ensure even dimensions for codec compatibility
|
||
width = width if width % 2 == 0 else width - 1
|
||
height = height if height % 2 == 0 else height - 1
|
||
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
jobs[job_id]["total"] = max(total, 1)
|
||
|
||
# 3. Write frame-swapped temp video
|
||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||
vid_stem = os.path.splitext(video_name)[0]
|
||
dir_part = "" if model_filename.startswith("_turntable/") else os.path.dirname(model_filename)
|
||
basename = os.path.basename(model_filename)
|
||
clean_basename = naming.get_base_name(basename)
|
||
prev_tag = f"_prev{int(scale*100)}" if scale < 1.0 else ""
|
||
tmp_basename = f"{ts}_fs_tmp_{vid_stem}_{clean_basename}{prev_tag}.mp4"
|
||
out_basename = f"{ts}_fs_{vid_stem}_{clean_basename}{prev_tag}.mp4"
|
||
if dir_part:
|
||
tmp_name = f"{dir_part}/{tmp_basename}"
|
||
out_name = f"{dir_part}/{out_basename}"
|
||
else:
|
||
tmp_name = tmp_basename
|
||
out_name = out_basename
|
||
tmp_path = os.path.join(output_dir, tmp_name)
|
||
out_path = os.path.join(output_dir, out_name)
|
||
|
||
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
||
writer = cv2.VideoWriter(tmp_path, fourcc, fps, (width, height))
|
||
|
||
frame_idx = 0
|
||
while True:
|
||
ret, frame = cap.read()
|
||
if not ret:
|
||
break
|
||
with embeddings._gpu_lock:
|
||
tgt_faces = app.get(frame)
|
||
result = frame
|
||
if tgt_faces:
|
||
# Only swap the largest face — avoids false-positive detections
|
||
# (reflections, background faces, face-like textures) causing ghost heads.
|
||
# Ignore faces smaller than 40×40px (1600px²) as likely false positives.
|
||
MIN_FACE_AREA = 1600
|
||
valid = [f for f in tgt_faces
|
||
if (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]) >= MIN_FACE_AREA]
|
||
if valid:
|
||
result = frame.copy()
|
||
best_face = max(valid, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
|
||
try:
|
||
with embeddings._gpu_lock:
|
||
result = swapper.get(result, best_face, src_face, paste_back=True)
|
||
except Exception:
|
||
pass
|
||
if gfpgan_restorer is not None:
|
||
try:
|
||
with embeddings._gpu_lock:
|
||
_, _, result = gfpgan_restorer.enhance(
|
||
result, has_aligned=False, only_center_face=False, paste_back=True
|
||
)
|
||
except Exception:
|
||
pass
|
||
if scale < 1.0:
|
||
result = cv2.resize(result, (width, height), interpolation=cv2.INTER_AREA)
|
||
writer.write(result)
|
||
frame_idx += 1
|
||
jobs[job_id]["done"] = frame_idx
|
||
|
||
cap.release()
|
||
writer.release()
|
||
|
||
# 4. Remux with original audio via ffmpeg
|
||
try:
|
||
import subprocess
|
||
r = subprocess.run([
|
||
'ffmpeg', '-y',
|
||
'-i', tmp_path,
|
||
'-i', video_path,
|
||
'-map', '0:v:0', '-map', '1:a?',
|
||
'-c:v', 'libx264', '-preset', 'fast', '-crf', '18',
|
||
'-c:a', 'aac', '-movflags', '+faststart',
|
||
out_path,
|
||
], capture_output=True, timeout=600)
|
||
if r.returncode == 0:
|
||
os.remove(tmp_path)
|
||
else:
|
||
os.rename(tmp_path, out_path)
|
||
print(f"[faceswap] ffmpeg failed ({r.returncode}), using raw mp4v output")
|
||
except Exception as fe:
|
||
os.rename(tmp_path, out_path)
|
||
print(f"[faceswap] ffmpeg error: {fe}")
|
||
|
||
# 5. Snapshot poster + register output in DB under same group as model
|
||
_make_video_poster(out_path)
|
||
person = database.get_person(model_filename)
|
||
group_id = (person[1] if person and person[1] else naming.get_base_name(os.path.basename(model_filename)))
|
||
database.upsert_person(
|
||
out_name,
|
||
filepath=out_path,
|
||
group_id=group_id,
|
||
content_type='video',
|
||
faceswap_source_video=video_name,
|
||
source_refs=json.dumps([model_filename]),
|
||
sort_order=database.get_next_sort_order(group_id),
|
||
)
|
||
|
||
jobs[job_id]["status"] = "done"
|
||
jobs[job_id]["output"] = out_name
|
||
_invalidate_static()
|
||
|
||
except Exception as e:
|
||
print(f"[faceswap] error: {e}")
|
||
jobs[job_id]["status"] = "error"
|
||
jobs[job_id]["error"] = str(e)
|
||
|
||
|
||
def _faceswap_worker_ff(job_id: str, model_filename: str, video_name: str,
|
||
hair: bool = True, enhance: bool = True, preview_scale: float = 1.0):
|
||
"""High-quality faceswap via FaceFusion CLI (supports hair_swapper + ghost model)."""
|
||
import subprocess as sp
|
||
import sys
|
||
import re as _re
|
||
|
||
output_dir = _load_output_dir()
|
||
wireframe_dir = _load_wireframe_dir()
|
||
|
||
with open(CONFIG_PATH, 'r') as f:
|
||
conf = json.load(f)
|
||
ff_dir = os.path.expanduser(conf.get('facefusion_dir', '~/facefusion'))
|
||
ff_venv = os.path.expanduser(conf.get('facefusion_venv', '~/facefusion-venv'))
|
||
|
||
ff_script = os.path.join(ff_dir, 'facefusion.py')
|
||
ff_py = os.path.join(ff_venv, 'bin', 'python')
|
||
if not os.path.exists(ff_py):
|
||
ff_py = sys.executable
|
||
|
||
if not os.path.exists(ff_script):
|
||
jobs[job_id]['status'] = 'error'
|
||
jobs[job_id]['error'] = (
|
||
f'FaceFusion not found at {ff_dir}. '
|
||
'Run: bash tour-comfy/install_facefusion.sh'
|
||
)
|
||
return
|
||
|
||
src_path = os.path.join(output_dir, model_filename)
|
||
video_path = os.path.join(wireframe_dir, video_name)
|
||
ts = time.strftime('%Y%m%d_%H%M%S')
|
||
vid_stem = os.path.splitext(video_name)[0]
|
||
dir_part = "" if model_filename.startswith("_turntable/") else os.path.dirname(model_filename)
|
||
basename = os.path.basename(model_filename)
|
||
clean_basename = naming.get_base_name(basename)
|
||
scale = max(0.1, min(1.0, preview_scale))
|
||
prev_tag = f'_prev{int(scale*100)}' if scale < 1.0 else ''
|
||
out_basename = f'{ts}_fs_{vid_stem}_{clean_basename}{prev_tag}.mp4'
|
||
if dir_part:
|
||
out_name = f'{dir_part}/{out_basename}'
|
||
else:
|
||
out_name = out_basename
|
||
out_path = os.path.join(output_dir, out_name)
|
||
|
||
processors = ['face_swapper']
|
||
# hair_swapper is not available in this FaceFusion version; use face_enhancer for quality
|
||
if enhance:
|
||
processors.append('face_enhancer')
|
||
|
||
cmd = [
|
||
ff_py, ff_script, 'headless-run',
|
||
'--source-paths', src_path,
|
||
'--target-path', video_path,
|
||
'--output-path', out_path,
|
||
'--processors', *processors,
|
||
'--execution-providers', 'cuda',
|
||
# Limit to 2 threads: each thread owns a cuBLAS handle + workspace; more
|
||
# threads exhausts VRAM when ComfyUI is running concurrently on the same GPU.
|
||
'--execution-thread-count', '2',
|
||
'--face-swapper-model', 'ghost_3_256',
|
||
# The default yolo_face detector at score 0.5 misses the extreme-angle /
|
||
# cropped close-up faces common in these POV template clips, so the swap
|
||
# silently no-ops. scrfd at a lower score + multi-angle detection reliably
|
||
# finds them; 'many' selector swaps every detected face per frame.
|
||
'--face-detector-model', 'scrfd',
|
||
'--face-detector-score', '0.3',
|
||
'--face-detector-angles', '0', '90', '270',
|
||
# 'one' swaps only the single best face per frame; 'many' caused ghost heads
|
||
# by swapping false-positive detections (skin texture, reflections, etc.)
|
||
'--face-selector-mode', 'one',
|
||
]
|
||
if enhance:
|
||
cmd += ['--face-enhancer-model', 'gfpgan_1.4']
|
||
if scale < 1.0:
|
||
# Determine native video width to compute preview width
|
||
try:
|
||
import cv2 as _cv2
|
||
_cap = _cv2.VideoCapture(video_path)
|
||
_native_w = int(_cap.get(_cv2.CAP_PROP_FRAME_WIDTH))
|
||
_cap.release()
|
||
_preview_w = max(2, int(_native_w * scale))
|
||
if _preview_w % 2 != 0:
|
||
_preview_w -= 1
|
||
cmd += ['--output-video-width', str(_preview_w)]
|
||
except Exception:
|
||
pass
|
||
|
||
jobs[job_id]['total'] = 100
|
||
jobs[job_id]['done'] = 0
|
||
|
||
# Ensure CUDA libs are on LD_LIBRARY_PATH for the subprocess (inherited from parent,
|
||
# but also add nvidia package libs as fallback if running outside start_api.sh)
|
||
import site as _site
|
||
_sp_pkgs = next((p for p in _site.getsitepackages() if 'site-packages' in p), '')
|
||
_nv_base = os.path.join(_sp_pkgs, 'nvidia')
|
||
_extra_libs = ':'.join(
|
||
os.path.join(_nv_base, pkg, 'lib')
|
||
for pkg in ('cuda_runtime', 'cublas', 'cudnn', 'curand', 'cufft', 'cusolver', 'cusparse', 'nvjitlink', 'cuda_nvrtc')
|
||
if os.path.isdir(os.path.join(_nv_base, pkg, 'lib'))
|
||
)
|
||
_env = os.environ.copy()
|
||
if _extra_libs:
|
||
_env['LD_LIBRARY_PATH'] = _extra_libs + (':' + _env['LD_LIBRARY_PATH'] if _env.get('LD_LIBRARY_PATH') else '')
|
||
|
||
try:
|
||
output_lines = []
|
||
proc = sp.Popen(
|
||
cmd, cwd=ff_dir, env=_env,
|
||
stdout=sp.PIPE, stderr=sp.PIPE,
|
||
text=True, errors='replace',
|
||
)
|
||
# FaceFusion writes tqdm progress to stderr; stdout carries other output.
|
||
# Parse frame counts from both streams so the UI job counter updates.
|
||
import threading as _thr
|
||
def _parse_progress(line: str):
|
||
m = _re.search(r'(\d+)\s*/\s*(\d+)', line)
|
||
if m:
|
||
done, total = int(m.group(1)), int(m.group(2))
|
||
if total > 0:
|
||
jobs[job_id]['done'] = done
|
||
jobs[job_id]['total'] = total
|
||
def _drain_stderr():
|
||
for ln in proc.stderr:
|
||
output_lines.append(ln.rstrip())
|
||
print(f'[facefusion] {ln.rstrip()}')
|
||
_parse_progress(ln)
|
||
_thr.Thread(target=_drain_stderr, daemon=True).start()
|
||
for line in proc.stdout:
|
||
print(f'[facefusion] {line.rstrip()}')
|
||
_parse_progress(line)
|
||
proc.wait()
|
||
|
||
if proc.returncode != 0:
|
||
tail = '\n'.join(output_lines[-10:])
|
||
raise RuntimeError(f'FaceFusion exited with code {proc.returncode}: {tail}')
|
||
if not os.path.exists(out_path):
|
||
raise RuntimeError('FaceFusion produced no output file')
|
||
|
||
_make_video_poster(out_path)
|
||
person = database.get_person(model_filename)
|
||
group_id = (person[1] if person and person[1] else naming.get_base_name(os.path.basename(model_filename)))
|
||
database.upsert_person(
|
||
out_name, filepath=out_path, group_id=group_id,
|
||
content_type='video', faceswap_source_video=video_name,
|
||
source_refs=json.dumps([model_filename]),
|
||
sort_order=database.get_next_sort_order(group_id),
|
||
)
|
||
_invalidate_static()
|
||
jobs[job_id]['status'] = 'done'
|
||
jobs[job_id]['output'] = out_name
|
||
|
||
except Exception as e:
|
||
print(f'[faceswap-ff] error: {e}')
|
||
jobs[job_id]['status'] = 'error'
|
||
jobs[job_id]['error'] = str(e)
|
||
|
||
|
||
# --- 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_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] = {"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] = {"text": " ".join(current_desc).strip(), "beta": current_beta}
|
||
|
||
return poses
|
||
|
||
|
||
def _save_poses(poses: dict) -> None:
|
||
"""Rewrite poses.md from a {name: {text, beta}} dict, round-tripping _load_poses' format.
|
||
|
||
Each pose is written as a ``# Name`` (or ``# Name (beta)``) header followed by its body.
|
||
Body sentences separated by '. ' are written on their own lines to match the existing
|
||
hand-authored style.
|
||
"""
|
||
poses_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "poses.md")
|
||
blocks = []
|
||
for name, entry in poses.items():
|
||
name = str(name).strip()
|
||
if not name:
|
||
continue
|
||
if isinstance(entry, dict):
|
||
text = str(entry.get("text", "")).strip()
|
||
beta = bool(entry.get("beta"))
|
||
else:
|
||
text = str(entry).strip()
|
||
beta = False
|
||
header = f"# {name}{' (beta)' if beta else ''}"
|
||
# Split into readable lines on sentence boundaries without losing the period.
|
||
body_lines = [s.strip() for s in re.split(r'(?<=\.)\s+', text) if s.strip()]
|
||
blocks.append(header + "\n" + "\n".join(body_lines))
|
||
with open(poses_path, "w", encoding="utf-8") as f:
|
||
f.write("\n\n".join(blocks) + ("\n" if blocks else ""))
|
||
|
||
|
||
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 _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,
|
||
seed: int = -1,
|
||
max_area: int = 0,
|
||
steps: int = 4,
|
||
cfg: float = 1.0,
|
||
sampler_name: str = "euler_ancestral",
|
||
scheduler: str = "beta",
|
||
extra_images: list = None, # additional PIL images wired to image2, image3
|
||
) -> bytes:
|
||
global _last_user_generation_time
|
||
if not _idle_turntable_busy:
|
||
_last_user_generation_time = time.time()
|
||
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
|
||
|
||
# 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]
|
||
|
||
# ── background-removal routing ────────────────────────────────────────────
|
||
# Two configurable strategies (config.json key "bg_removal"):
|
||
#
|
||
# "rembg" (default) — strip transparent keyword → Qwen renders a natural
|
||
# scene → rembg (U2Net) separates person from any complex background.
|
||
#
|
||
# "sam2" — replace transparent keyword with "black background" → Qwen
|
||
# renders a solid black BG → SAM2 bbox segmentation on a black image
|
||
# works perfectly because the contrast is maximal.
|
||
#
|
||
# Either way, explicit "black background" in the prompt always routes to
|
||
# SAM2 (the user already set up the ideal SAM2 input).
|
||
# ─────────────────────────────────────────────────────────────────────────
|
||
_TRANSPARENT_KWS = ["transparent background", "no background",
|
||
"remove background", "alpha channel"]
|
||
_BLACK_BG_KWS = ["black background"]
|
||
|
||
with open(CONFIG_PATH) as _cf:
|
||
_bg_conf = json.load(_cf)
|
||
bg_method = _bg_conf.get("bg_removal", "rembg") # "rembg" | "sam2"
|
||
|
||
is_transparent = any(kw in prompt.lower() for kw in _TRANSPARENT_KWS)
|
||
is_black_bg = any(kw in prompt.lower() for kw in _BLACK_BG_KWS)
|
||
post_process = None # "rembg" | "sam2"
|
||
|
||
if is_transparent:
|
||
if bg_method == "sam2":
|
||
# Swap "transparent background" → "black background" so Qwen renders
|
||
# a pure-black BG that SAM2 can segment with maximal contrast.
|
||
cleaned = prompt
|
||
for kw in _TRANSPARENT_KWS:
|
||
cleaned = re.sub(re.escape(kw), "black background", cleaned, flags=re.IGNORECASE)
|
||
# Collapse duplicates if multiple keywords matched
|
||
cleaned = re.sub(r"(?i)(black background[\s,]*){2,}", "black background, ", cleaned)
|
||
cleaned = re.sub(r",\s*,", ",", cleaned).strip(", ")
|
||
graph[NODE_POSITIVE]["inputs"]["prompt"] = cleaned
|
||
graph[NODE_NEGATIVE]["inputs"]["prompt"] = (
|
||
"real background, outdoor scene, indoor scene, gradient, "
|
||
"colored background, watermark, deformed anatomy"
|
||
)
|
||
post_process = "sam2"
|
||
else:
|
||
# Strip the keyword so Qwen renders a natural scene; rembg handles
|
||
# any background complexity reliably.
|
||
cleaned = prompt
|
||
for kw in _TRANSPARENT_KWS:
|
||
cleaned = re.sub(re.escape(kw), "", cleaned, flags=re.IGNORECASE)
|
||
cleaned = re.sub(r",\s*,", ",", cleaned)
|
||
cleaned = re.sub(r",\s*$", "", cleaned.strip()).strip(", ")
|
||
graph[NODE_POSITIVE]["inputs"]["prompt"] = cleaned
|
||
graph[NODE_NEGATIVE]["inputs"]["prompt"] = "deformed anatomy, watermark, logo"
|
||
post_process = "rembg"
|
||
|
||
elif is_black_bg:
|
||
# Prompt already specifies a black background — ideal SAM2 input.
|
||
# Route to SAM2 regardless of the configured bg_removal method.
|
||
post_process = "sam2"
|
||
|
||
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)
|
||
if not _idle_turntable_busy:
|
||
_last_user_generation_time = time.time()
|
||
png_bytes = _comfy_fetch_image(outputs)
|
||
|
||
if post_process == "sam2":
|
||
# Input has a black background (Qwen was told "black background").
|
||
# Use threshold-derived bbox so SAM2 gets a person-specific hint
|
||
# rather than the full frame — full-frame bbox inverts the mask on
|
||
# black-bg images because the large dark region scores higher.
|
||
png_bytes = _apply_transparency_black_bg(png_bytes)
|
||
elif post_process == "rembg":
|
||
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 = os.path.expanduser(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}")
|
||
|
||
|
||
# --- static data files -------------------------------------------------------
|
||
|
||
_static_write_lock = threading.Lock()
|
||
_invalidate_timer: "threading.Timer | None" = None
|
||
_invalidate_timer_lock = threading.Lock()
|
||
_privacy_locked = False
|
||
_privacy_lock = threading.Lock()
|
||
|
||
|
||
def _privacy_monitor_daemon():
|
||
"""Monitors GNOME/Freedesktop screen lock via D-Bus gdbus monitor."""
|
||
global _privacy_locked
|
||
print("[privacy] Starting monitor daemon...")
|
||
|
||
# We try to monitor both GNOME and Freedesktop ScreenSaver
|
||
destinations = [
|
||
("org.gnome.ScreenSaver", "/org/gnome/ScreenSaver"),
|
||
("org.freedesktop.ScreenSaver", "/org/freedesktop/ScreenSaver")
|
||
]
|
||
|
||
def monitor(dest, obj_path):
|
||
global _privacy_locked
|
||
cmd = ["gdbus", "monitor", "--session", "--dest", dest, "--object-path", obj_path]
|
||
try:
|
||
proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
|
||
for line in proc.stdout:
|
||
# GNOME: ActiveChanged (true,) or (false,)
|
||
# Freedesktop: Locked or ActiveChanged
|
||
if "ActiveChanged" in line or "Locked" in line:
|
||
is_locked = "(true,)" in line or "true" in line.lower()
|
||
with _privacy_lock:
|
||
if _privacy_locked != is_locked:
|
||
_privacy_locked = is_locked
|
||
print(f"[privacy] System lock state change detected via {dest}: {is_locked}")
|
||
_write_all_static()
|
||
except Exception as e:
|
||
print(f"[privacy] Monitor error for {dest}: {e}")
|
||
|
||
for dest, path in destinations:
|
||
threading.Thread(target=monitor, args=(dest, path), daemon=True).start()
|
||
|
||
|
||
def _write_json(path: str, data) -> None:
|
||
"""Atomic JSON write: write to .tmp then os.replace (no partial reads)."""
|
||
tmp = path + ".tmp"
|
||
with open(tmp, "w") as f:
|
||
json.dump(data, f)
|
||
os.replace(tmp, path)
|
||
|
||
|
||
def _write_all_static() -> None:
|
||
"""Regenerate <output_dir>/_data/{images,names,groups,group-names,videos,poses}.json.
|
||
|
||
Called on startup and after every mutation so the browser can\
|
||
fetch pre-computed JSON instead of hitting the DB on each page load.
|
||
Uses a lock so concurrent invalidations don't race each other.
|
||
"""
|
||
with _static_write_lock:
|
||
try:
|
||
output_dir = _load_output_dir()
|
||
data_dir = os.path.join(output_dir, "_data")
|
||
os.makedirs(data_dir, exist_ok=True)
|
||
|
||
# Single DB query reused for images / names / groups
|
||
persons = database.list_persons(include_archived=True)
|
||
|
||
# images.json — all images (frontend filters by .archived field)
|
||
db_images = []
|
||
archived_count = 0
|
||
for p in persons:
|
||
exists, mtime = _get_cached_file_meta(p[0], output_dir)
|
||
if not exists:
|
||
continue
|
||
|
||
is_archived = bool(p[14]) if p[14] else False
|
||
if is_archived:
|
||
archived_count += 1
|
||
|
||
tags_val = p[16]
|
||
tags_list = []
|
||
if tags_val:
|
||
if isinstance(tags_val, str):
|
||
try:
|
||
tags_list = json.loads(tags_val)
|
||
except Exception:
|
||
tags_list = []
|
||
elif isinstance(tags_val, list):
|
||
tags_list = tags_val
|
||
|
||
obj_val = p[22]
|
||
obj_list = []
|
||
if obj_val:
|
||
if isinstance(obj_val, str):
|
||
try:
|
||
obj_list = json.loads(obj_val)
|
||
except Exception:
|
||
obj_list = []
|
||
elif isinstance(obj_val, list):
|
||
obj_list = obj_val
|
||
|
||
db_images.append({
|
||
"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],
|
||
"has_clothing": p[11],
|
||
"content_type": p[12] or "image",
|
||
"faceswap_source_video": p[13],
|
||
"archived": is_archived,
|
||
"is_source": bool(p[15]) if p[15] else False,
|
||
"tags": tags_list,
|
||
"pose_description": p[17],
|
||
"pose_skeleton": p[18],
|
||
"people_count": p[19],
|
||
"anatomical_completeness": p[20],
|
||
"facial_direction": p[21],
|
||
"objects": obj_list,
|
||
"description": p[23],
|
||
"filepath": p[24],
|
||
})
|
||
print(f"[static] write_all: {len(db_images)} total images, {archived_count} archived")
|
||
try:
|
||
db_images.sort(
|
||
key=lambda x: _get_cached_file_meta(x["filename"], output_dir)[1],
|
||
reverse=True,
|
||
)
|
||
except Exception:
|
||
pass
|
||
_write_json(os.path.join(data_dir, "images.json"), {"images": db_images})
|
||
|
||
# names.json — {filename: name}
|
||
_write_json(os.path.join(data_dir, "names.json"),
|
||
{p[0]: p[1] for p in persons if p[1]})
|
||
|
||
# groups.json — {filename: group_id}
|
||
_write_json(os.path.join(data_dir, "groups.json"),
|
||
{p[0]: p[2] for p in persons if p[2]})
|
||
|
||
# group-names.json — {group_id: display_name}
|
||
_write_json(os.path.join(data_dir, "group-names.json"),
|
||
database.get_all_group_names())
|
||
|
||
# videos.json — wireframe template videos and grouped assets
|
||
wireframe_dir = _load_wireframe_dir()
|
||
_write_json(os.path.join(data_dir, "videos.json"),
|
||
_get_grouped_wireframes(wireframe_dir))
|
||
|
||
# poses.json — pose library from poses.md
|
||
_write_json(os.path.join(data_dir, "poses.json"), _load_poses())
|
||
|
||
# system_status.json — for privacy lock etc
|
||
_write_json(os.path.join(data_dir, "system_status.json"), {"locked": _privacy_locked})
|
||
|
||
# config.json — current config snapshot
|
||
try:
|
||
with open(CONFIG_PATH, "r") as _f:
|
||
_write_json(os.path.join(data_dir, "config.json"), json.load(_f))
|
||
except Exception:
|
||
pass
|
||
|
||
# Update car.html PRELOADED_IMAGES
|
||
_sync_preloaded_images()
|
||
|
||
# Group-specific JSON and HTML generation
|
||
car_html_src = os.path.join(_HERE, "car.html")
|
||
template_content = ""
|
||
if os.path.exists(car_html_src):
|
||
try:
|
||
with open(car_html_src, "r", encoding="utf-8") as f:
|
||
template_content = f.read()
|
||
except Exception as template_err:
|
||
print(f"[static] Error reading car.html template: {template_err}")
|
||
|
||
# Extract unique group IDs from db_images
|
||
gids = set(p["group_id"] for p in db_images if p.get("group_id"))
|
||
for gid in gids:
|
||
if not gid:
|
||
continue
|
||
# Sanitize the group_id to prevent directory traversal / invalid characters in filenames
|
||
sanitized_gid = re.sub(r'[^a-zA-Z0-9_\-]', '_', gid)
|
||
|
||
# Filter images belonging to this group
|
||
group_images = [p for p in db_images if p.get("group_id") == gid]
|
||
|
||
# Write group json data: _data/group_{sanitized_gid}.json
|
||
group_json_path = os.path.join(data_dir, f"group_{sanitized_gid}.json")
|
||
_write_json(group_json_path, {"images": group_images})
|
||
|
||
# Write shoot specific html: shoot_{sanitized_gid}.html next to car.html
|
||
if template_content:
|
||
# Replace `const LOAD_ONLY_GROUP_ID = null;` with `const LOAD_ONLY_GROUP_ID = '{gid}';`
|
||
escaped_gid = gid.replace("'", "\\'").replace('"', '\\"')
|
||
replaced_content = template_content.replace(
|
||
"const LOAD_ONLY_GROUP_ID = null;",
|
||
f"const LOAD_ONLY_GROUP_ID = '{escaped_gid}';"
|
||
)
|
||
|
||
group_html_path = os.path.join(output_dir, f"shoot_{sanitized_gid}.html")
|
||
try:
|
||
with open(group_html_path, "w", encoding="utf-8") as f:
|
||
f.write(replaced_content)
|
||
except Exception as html_err:
|
||
print(f"[static] Error writing group html for {gid}: {html_err}")
|
||
|
||
except Exception as e:
|
||
print(f"[static] write_all error: {e}")
|
||
|
||
# Turntable static is cheap and independent; write outside the main lock
|
||
_write_turntable_static()
|
||
|
||
|
||
def _write_turntable_static() -> None:
|
||
"""Write _data/turntables.json with per-group turntable state + frame URLs."""
|
||
try:
|
||
import turntable_cache as tc
|
||
output_dir = _load_output_dir()
|
||
data_dir = os.path.join(output_dir, "_data")
|
||
os.makedirs(data_dir, exist_ok=True)
|
||
|
||
# Load group names from DB for display
|
||
try:
|
||
group_names = database.get_all_group_names()
|
||
except Exception:
|
||
group_names = {}
|
||
|
||
# Fetch active and existing group IDs and filenames from the database (excluding turntable assets)
|
||
all_db_gids = set()
|
||
active_db_gids = set()
|
||
all_db_filenames = set()
|
||
active_db_filenames = set()
|
||
try:
|
||
conn = database.get_db_connection()
|
||
cur = conn.cursor()
|
||
cur.execute("""
|
||
SELECT filename, group_id, archived, hidden
|
||
FROM person
|
||
WHERE NOT (filename LIKE '_turntable/%')
|
||
""")
|
||
for fname, gid, archived, hidden in cur.fetchall():
|
||
is_archived = bool(archived) if archived is not None else False
|
||
is_hidden = bool(hidden) if hidden is not None else False
|
||
|
||
all_db_filenames.add(fname)
|
||
if not is_archived and not is_hidden:
|
||
active_db_filenames.add(fname)
|
||
|
||
if gid:
|
||
all_db_gids.add(gid)
|
||
if not is_archived and not is_hidden:
|
||
active_db_gids.add(gid)
|
||
cur.close()
|
||
database._put_db_connection(conn)
|
||
except Exception as db_err:
|
||
print(f"[static] DB check error: {db_err}")
|
||
|
||
turntables = []
|
||
for gid in tc.list_cached_group_ids(output_dir):
|
||
is_orphaned = False
|
||
is_active = False
|
||
|
||
if gid.startswith("file_"):
|
||
fname = gid[5:]
|
||
if all_db_filenames and fname not in all_db_filenames:
|
||
is_orphaned = True
|
||
elif active_db_filenames and fname in active_db_filenames:
|
||
is_active = True
|
||
else:
|
||
if all_db_gids and gid not in all_db_gids:
|
||
is_orphaned = True
|
||
elif active_db_gids and gid in active_db_gids:
|
||
is_active = True
|
||
|
||
# 1. Truly orphaned turntable caches (no associated base images or files in DB)
|
||
if is_orphaned:
|
||
print(f"[static] Deleting truly orphaned turntable cache and DB records for group/file {gid}")
|
||
tc.delete_state(output_dir, gid)
|
||
try:
|
||
conn = database.get_db_connection()
|
||
cur = conn.cursor()
|
||
cur.execute("DELETE FROM person WHERE group_id = %s", (gid,))
|
||
conn.commit()
|
||
cur.close()
|
||
database._put_db_connection(conn)
|
||
except Exception as del_err:
|
||
print(f"[static] Failed to clean DB records for orphaned group/file {gid}: {del_err}")
|
||
continue
|
||
|
||
# 2. Inactive turntable caches (associated group/file is fully archived or hidden)
|
||
if not is_active:
|
||
# Keep cache on disk but do not write to turntables.json (hide from UI)
|
||
continue
|
||
|
||
state = tc.load_state(output_dir, gid)
|
||
if not state:
|
||
continue
|
||
# Build ordered frame URL list (relative to output_dir, served via /output/)
|
||
angles = state.get("angles", [])
|
||
views = state.get("views", {})
|
||
frames = []
|
||
for deg in angles:
|
||
dk = tc.deg_key(deg)
|
||
if dk in views:
|
||
abs_path = views[dk]
|
||
if os.path.exists(abs_path):
|
||
rel = os.path.relpath(abs_path, output_dir)
|
||
frames.append(rel.replace("\\", "/"))
|
||
|
||
video_rel = None
|
||
vp = state.get("video_path")
|
||
if vp and os.path.exists(vp):
|
||
video_rel = os.path.relpath(vp, output_dir).replace("\\", "/")
|
||
|
||
turntables.append({
|
||
"group_id": gid,
|
||
"group_name": group_names.get(gid, ""),
|
||
"preferred_filename": state.get("preferred_filename", ""),
|
||
"completed": state.get("completed", False),
|
||
"n_done": len(frames),
|
||
"n_total": state.get("n_views", 24),
|
||
"frames": frames,
|
||
"video_path": video_rel,
|
||
"started_at": state.get("started_at"),
|
||
"completed_at": state.get("completed_at"),
|
||
})
|
||
|
||
# Sort: complete first, then by most views done
|
||
turntables.sort(key=lambda t: (-int(t["completed"]), -t["n_done"]))
|
||
_write_json(os.path.join(data_dir, "turntables.json"), {"turntables": turntables})
|
||
|
||
summary = tc.get_status_summary(output_dir)
|
||
n_complete = sum(1 for v in summary.values() if v["completed"])
|
||
status_payload = {
|
||
"groups": summary,
|
||
"n_complete": n_complete,
|
||
"n_total": len(summary),
|
||
"is_generating": _idle_turntable_busy,
|
||
"is_paused": _idle_turntable_paused,
|
||
"idle_seconds": round(time.time() - _last_user_generation_time, 1),
|
||
}
|
||
_write_json(os.path.join(data_dir, "turntable_status.json"), status_payload)
|
||
except Exception as e:
|
||
print(f"[static] write_turntable error: {e}")
|
||
|
||
|
||
def _invalidate_static() -> None:
|
||
"""Coalesce rapid invalidation calls — restarts a 0.3 s debounce timer each time.
|
||
At most one _write_all_static() runs per quiet window, preventing thread floods
|
||
during batch jobs that call this after every single image."""
|
||
global _invalidate_timer
|
||
with _invalidate_timer_lock:
|
||
if _invalidate_timer is not None:
|
||
_invalidate_timer.cancel()
|
||
t = threading.Timer(0.3, _write_all_static)
|
||
t.daemon = True
|
||
t.start()
|
||
_invalidate_timer = t
|
||
|
||
|
||
# -----------------------------------------------------------------------------
|
||
|
||
def _batch_worker(job_id: str, filenames: list, prompts: list[str], poses: list,
|
||
seed: int, max_area: int, group_id: str | None = None,
|
||
wireframe_ref: str | None = None, wireframe_time: float = 0.5,
|
||
pad_top: Any = 0, pad_right: Any = 0,
|
||
pad_bottom: Any = 0, pad_left: Any = 0, pad_fill: str = "transparent",
|
||
pad_outpaint: bool = False):
|
||
output_dir = _load_output_dir()
|
||
for fname in filenames:
|
||
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(os.path.basename(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(os.path.basename(fname))
|
||
|
||
fpath = os.path.join(output_dir, fname)
|
||
if not os.path.exists(fpath):
|
||
jobs[job_id]["failed"] += len(prompts)
|
||
continue
|
||
|
||
try:
|
||
base_pil = Image.open(fpath).convert("RGB")
|
||
if any([pad_top, pad_right, pad_bottom, pad_left]):
|
||
base_pil = _apply_manual_pad(base_pil, pad_top, pad_right, pad_bottom, pad_left, pad_fill)
|
||
|
||
# Extract wireframe pose reference frame once per filename
|
||
pose_guide_pil = None
|
||
if wireframe_ref:
|
||
try:
|
||
wf_path = os.path.join(_load_wireframe_dir(), wireframe_ref)
|
||
cap = cv2.VideoCapture(wf_path)
|
||
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
target_frame = max(0, min(total_frames - 1, int(total_frames * wireframe_time)))
|
||
cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame)
|
||
ret, frame = cap.read()
|
||
cap.release()
|
||
if ret:
|
||
pose_guide_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||
if any([pad_top, pad_right, pad_bottom, pad_left]):
|
||
pose_guide_pil = _apply_manual_pad(pose_guide_pil, pad_top, pad_right, pad_bottom, pad_left, pad_fill)
|
||
print(f"[batch] using wireframe {wireframe_ref} frame {target_frame}/{total_frames}")
|
||
except Exception as wf_err:
|
||
print(f"[batch] wireframe extract error: {wf_err}")
|
||
|
||
for prompt, pose in zip(prompts, poses):
|
||
if jobs[job_id].get("cancelled"):
|
||
return
|
||
try:
|
||
try:
|
||
database.save_db_prompt("pose-prompt", prompt, {
|
||
"pose": pose,
|
||
"seed": seed,
|
||
"max_area": max_area,
|
||
"wireframe_ref": wireframe_ref,
|
||
"wireframe_time": wireframe_time,
|
||
"pad_top": pad_top,
|
||
"pad_right": pad_right,
|
||
"pad_bottom": pad_bottom,
|
||
"pad_left": pad_left,
|
||
"pad_fill": pad_fill,
|
||
"pad_outpaint": pad_outpaint
|
||
})
|
||
except Exception as db_err:
|
||
print(f"[batch] failed to save to prompt table: {db_err}")
|
||
|
||
pil = base_pil
|
||
actual_prompt = prompt
|
||
if pad_outpaint:
|
||
out_instr = "Naturally outpaint and extend the borders of the image to complete the scene."
|
||
if not actual_prompt.strip():
|
||
actual_prompt = out_instr
|
||
elif out_instr.lower() not in actual_prompt.lower():
|
||
actual_prompt = f"{actual_prompt}. {out_instr}"
|
||
|
||
# Rotate 180° for poses that work better upside-down
|
||
if pose and pose.lower().strip() in ROTATE_180_POSES:
|
||
pil = pil.rotate(180)
|
||
|
||
extra_imgs = [pose_guide_pil] if pose_guide_pil else None
|
||
png = _run_pipeline(pil, actual_prompt, seed, max_area, extra_images=extra_imgs)
|
||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||
dir_part = "" if fname.startswith("_turntable/") else os.path.dirname(fname)
|
||
basename = os.path.basename(fname)
|
||
clean_basename = naming.get_base_name(basename)
|
||
new_basename = f"{ts}_{clean_basename}"
|
||
if dir_part:
|
||
out_name = f"{dir_part}/{new_basename}"
|
||
else:
|
||
out_name = new_basename
|
||
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)
|
||
next_order = database.get_next_sort_order(actual_gid)
|
||
database.upsert_person(
|
||
out_name, filepath=out_path, embedding=embedding,
|
||
group_id=actual_gid, prompt=actual_prompt, pose=pose,
|
||
has_background=has_bg, sort_order=next_order,
|
||
source_refs=json.dumps([fname]),
|
||
)
|
||
_update_cached_file_meta(out_name, exists=True)
|
||
except Exception as db_err:
|
||
print(f"Database error in batch worker: {db_err}")
|
||
|
||
jobs[job_id]["latest_output"] = out_name
|
||
jobs[job_id]["done"] += 1
|
||
# Regenerate static JSON so the frontend's polling picks up the
|
||
# new image immediately (progressive refresh, not just at the end).
|
||
_invalidate_static()
|
||
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"
|
||
_invalidate_static()
|
||
|
||
|
||
def _multi_ref_worker(job_id: str, filenames: list[str], prompts: list[str], poses: list,
|
||
seed: int, max_area: int,
|
||
pad_top: Any = 0, pad_right: Any = 0,
|
||
pad_bottom: Any = 0, pad_left: Any = 0,
|
||
pad_fill: str = "black", pad_outpaint: bool = False):
|
||
"""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
|
||
|
||
# Apply padding to images if requested
|
||
if any([pad_top, pad_right, pad_bottom, pad_left]):
|
||
pils = [(name, _apply_manual_pad(pil, pad_top, pad_right, pad_bottom, pad_left, pad_fill)) for name, pil in pils]
|
||
|
||
# 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:
|
||
try:
|
||
database.save_db_prompt("pose-prompt", prompt, {
|
||
"pose": pose,
|
||
"seed": seed,
|
||
"max_area": max_area,
|
||
"filenames": filenames,
|
||
"pad_top": pad_top,
|
||
"pad_right": pad_right,
|
||
"pad_bottom": pad_bottom,
|
||
"pad_left": pad_left,
|
||
"pad_fill": pad_fill,
|
||
"pad_outpaint": pad_outpaint
|
||
})
|
||
except Exception as db_err:
|
||
print(f"[multi-ref] failed to save to prompt table: {db_err}")
|
||
work_pil = primary_pil
|
||
actual_prompt = prompt
|
||
if pad_outpaint:
|
||
out_instr = "Naturally outpaint and extend the borders of the image to complete the scene."
|
||
if not actual_prompt.strip():
|
||
actual_prompt = out_instr
|
||
elif out_instr.lower() not in actual_prompt.lower():
|
||
actual_prompt = f"{actual_prompt}. {out_instr}"
|
||
|
||
if pose and pose.lower().strip() in ROTATE_180_POSES:
|
||
work_pil = work_pil.rotate(180)
|
||
|
||
png = _run_pipeline(work_pil, actual_prompt, seed, max_area, extra_images=extra_pils)
|
||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||
dir_part = "" if primary_fname.startswith("_turntable/") else os.path.dirname(primary_fname)
|
||
basename = os.path.basename(primary_fname)
|
||
clean_basename = naming.get_base_name(basename)
|
||
new_basename = f"{ts}_mr_{clean_basename}"
|
||
if dir_part:
|
||
out_name = f"{dir_part}/{new_basename}"
|
||
else:
|
||
out_name = new_basename
|
||
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)
|
||
next_order = database.get_next_sort_order(output_gid)
|
||
database.upsert_person(out_name, filepath=out_path, embedding=embedding,
|
||
group_id=output_gid, prompt=actual_prompt, pose=pose,
|
||
has_background=has_bg, sort_order=next_order,
|
||
source_refs=json.dumps([f for f, _ in pils]))
|
||
_update_cached_file_meta(out_name, exists=True)
|
||
except Exception as db_err:
|
||
print(f"DB error in multi-ref: {db_err}")
|
||
|
||
jobs[job_id]["latest_output"] = out_name
|
||
jobs[job_id]["done"] += 1
|
||
# Regenerate static JSON so the frontend's polling picks up the new
|
||
# image immediately (progressive refresh, matching _batch_worker).
|
||
_invalidate_static()
|
||
except Exception as e:
|
||
print(f"Error in multi-ref for prompt '{prompt}': {e}")
|
||
jobs[job_id]["failed"] += 1
|
||
|
||
jobs[job_id]["status"] = "done"
|
||
_invalidate_static()
|
||
|
||
|
||
# --- routes -----------------------------------------------------------------
|
||
|
||
class ConfigUpdate(BaseModel):
|
||
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.seed is not None:
|
||
conf["seed"] = update.seed
|
||
with open(CONFIG_PATH, "w") as f:
|
||
json.dump(conf, f, indent=2)
|
||
_invalidate_static()
|
||
return {"seed": conf["seed"]}
|
||
|
||
|
||
class SavePromptRequest(BaseModel):
|
||
type: str
|
||
prompt_text: str
|
||
metadata: dict | None = None
|
||
|
||
|
||
@app.post("/prompts")
|
||
def api_save_prompt(req: SavePromptRequest):
|
||
database.save_db_prompt(req.type, req.prompt_text, req.metadata)
|
||
return {"status": "success"}
|
||
|
||
|
||
@app.get("/prompts")
|
||
def api_list_prompts(type: str | None = None, limit: int = 100):
|
||
return database.list_db_prompts(type, limit)
|
||
|
||
|
||
class GroupArchiveRequest(BaseModel):
|
||
filenames: list[str]
|
||
|
||
class RepairRequest(BaseModel):
|
||
action: str # "delete_record", "import_file", "restore", "delete_permanently", "assign_group"
|
||
filename: str
|
||
group_id: str | None = None
|
||
|
||
class BatchRequest(BaseModel):
|
||
filenames: list[str]
|
||
prompt: str | list[str]
|
||
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
|
||
wireframe_ref: str | None = None # wireframe video name to use as pose guide (image2 slot)
|
||
wireframe_time: float = 0.5 # normalized time (0–1) to extract the pose frame from
|
||
pad_top: int | float | str = 0
|
||
pad_right: int | float | str = 0
|
||
pad_bottom: int | float | str = 0
|
||
pad_left: int | float | str = 0
|
||
pad_fill: str = "black"
|
||
pad_outpaint: bool = False
|
||
|
||
|
||
@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, "cancelled": False}
|
||
t = threading.Thread(
|
||
target=_batch_worker,
|
||
args=(job_id, req.filenames, prompts, poses, req.seed, req.max_area, req.group_id),
|
||
kwargs={
|
||
"wireframe_ref": req.wireframe_ref, "wireframe_time": req.wireframe_time,
|
||
"pad_top": req.pad_top, "pad_right": req.pad_right,
|
||
"pad_bottom": req.pad_bottom, "pad_left": req.pad_left,
|
||
"pad_fill": req.pad_fill,
|
||
"pad_outpaint": req.pad_outpaint,
|
||
},
|
||
daemon=True,
|
||
)
|
||
t.start()
|
||
return {"job_id": job_id, "total": total_tasks}
|
||
|
||
|
||
@app.delete("/batch/{job_id}")
|
||
def cancel_batch(job_id: str):
|
||
if job_id not in jobs:
|
||
raise HTTPException(404, "Job not found")
|
||
jobs[job_id]["cancelled"] = True
|
||
jobs[job_id]["status"] = "cancelled"
|
||
return {"status": "cancelled", "job_id": job_id}
|
||
|
||
|
||
class MultiRefRequest(BaseModel):
|
||
filenames: list[str] # 2–3 reference images; first is primary (image1)
|
||
prompt: str | list[str]
|
||
poses: list[str | None] | None = None
|
||
seed: int = -1
|
||
max_area: int = 0
|
||
pad_top: int | float | str = 0
|
||
pad_right: int | float | str = 0
|
||
pad_bottom: int | float | str = 0
|
||
pad_left: int | float | str = 0
|
||
pad_fill: str = "black"
|
||
pad_outpaint: bool = False
|
||
|
||
|
||
@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),
|
||
kwargs={
|
||
"pad_top": req.pad_top, "pad_right": req.pad_right,
|
||
"pad_bottom": req.pad_bottom, "pad_left": req.pad_left,
|
||
"pad_fill": req.pad_fill, "pad_outpaint": req.pad_outpaint
|
||
},
|
||
daemon=True,
|
||
)
|
||
t.start()
|
||
return {"job_id": job_id, "total": len(prompts)}
|
||
|
||
|
||
class RefineRequest(BaseModel):
|
||
prompt: str = ""
|
||
filename: str | None = None
|
||
user_instruction: str | None = None
|
||
field: str | None = None
|
||
|
||
|
||
@app.post("/refine-prompt")
|
||
def refine_prompt(req: RefineRequest):
|
||
"""Refine or create a prompt/description using the external uncensored chat LLM."""
|
||
context_str = ""
|
||
if req.filename:
|
||
try:
|
||
person = database.get_person(req.filename)
|
||
if person:
|
||
# person columns: SELECT name, group_id, tags, embedding, clip_description, filepath,
|
||
# prompt, pose, sort_order, group_name, hidden, has_background, source_refs,
|
||
# has_clothing, is_source, pose_description, pose_skeleton,
|
||
# people_count, anatomical_completeness, facial_direction, objects
|
||
pose_desc = person[15]
|
||
people_count = person[17]
|
||
anatomical_completeness = person[18]
|
||
facial_direction = person[19]
|
||
objects_val = person[20]
|
||
|
||
context_parts = []
|
||
if pose_desc:
|
||
context_parts.append(f"Pose details: {pose_desc}")
|
||
if people_count is not None:
|
||
context_parts.append(f"Subject count: {people_count} person(s)")
|
||
if anatomical_completeness is not None:
|
||
context_parts.append(f"Anatomical completeness: {'complete/full body' if anatomical_completeness else 'partial/closeup'}")
|
||
if facial_direction:
|
||
context_parts.append(f"Gaze and facial direction: {facial_direction}")
|
||
|
||
if objects_val:
|
||
try:
|
||
if isinstance(objects_val, str):
|
||
objs = json.loads(objects_val)
|
||
else:
|
||
objs = objects_val
|
||
if objs:
|
||
tag_names = [o["tag"] for o in objs if isinstance(o, dict) and "tag" in o]
|
||
if tag_names:
|
||
context_parts.append(f"Detected elements/objects in scene: {', '.join(tag_names)}")
|
||
except Exception as parse_err:
|
||
print(f"[refine-prompt] failed to parse objects: {parse_err}")
|
||
|
||
if context_parts:
|
||
context_str = "\n".join(context_parts)
|
||
except Exception as db_err:
|
||
print(f"[refine-prompt] database error for {req.filename}: {db_err}")
|
||
|
||
field_label = req.field if req.field else "text"
|
||
|
||
if req.prompt:
|
||
if req.user_instruction:
|
||
user_content = f"Refine this {field_label} according to the following instructions:\n{req.user_instruction}\n\nOriginal {field_label}:\n{req.prompt}"
|
||
else:
|
||
user_content = f"Refine this {field_label} to make it more descriptive, detailed, and high-quality.\n\nOriginal {field_label}:\n{req.prompt}"
|
||
else:
|
||
# Prompt/text is empty, so we are creating one from scratch!
|
||
if req.user_instruction:
|
||
user_content = f"Create a new {field_label} from scratch according to these instructions:\n{req.user_instruction}"
|
||
else:
|
||
user_content = f"Create a new detailed and high-quality {field_label} from scratch."
|
||
|
||
if context_str:
|
||
user_content += f"\n\nUse the following image context details to ensure the output matches the reference characteristics closely:\n{context_str}"
|
||
|
||
# Use the same API as gen_poses.py
|
||
llm_api = "http://192.168.1.160:8001/v1/chat/completions"
|
||
payload = {
|
||
"model": "dphn/Dolphin3.0-Mistral-24B",
|
||
"messages": [
|
||
{"role": "system", "content": REFINEMENT_SYSTEM},
|
||
{"role": "user", "content": user_content}
|
||
],
|
||
"temperature": 0.8,
|
||
"max_tokens": 1024
|
||
}
|
||
|
||
try:
|
||
r = requests.post(llm_api, json=payload, timeout=90)
|
||
r.raise_for_status()
|
||
data = r.json()
|
||
refined = data["choices"][0]["message"]["content"].strip()
|
||
try:
|
||
database.save_db_prompt("refine", refined, {
|
||
"original": req.prompt,
|
||
"filename": req.filename,
|
||
"field": req.field
|
||
})
|
||
except Exception as db_err:
|
||
print(f"[refine-prompt] failed to save to prompt table: {db_err}")
|
||
return {"refined": refined}
|
||
except Exception as e:
|
||
print(f"Refinement error: {e}")
|
||
raise HTTPException(500, f"LLM refinement failed: {str(e)}")
|
||
|
||
|
||
class UpdatePromptRequest(BaseModel):
|
||
prompt: str
|
||
|
||
|
||
class UpdateDescriptionRequest(BaseModel):
|
||
description: str
|
||
|
||
|
||
@app.post("/images/{filename:path}/update-description")
|
||
def update_description(filename: str, req: UpdateDescriptionRequest):
|
||
try:
|
||
person = database.get_person(filename)
|
||
if not person:
|
||
raise HTTPException(404, "Image not found in database")
|
||
database.upsert_person(filename, description=req.description)
|
||
_invalidate_static()
|
||
return {"status": "success", "filename": filename, "description": req.description}
|
||
except Exception as e:
|
||
raise HTTPException(500, str(e))
|
||
|
||
|
||
@app.post("/images/{filename:path}/reverse-engineer")
|
||
def reverse_engineer(filename: str):
|
||
person = database.get_person(filename)
|
||
if not person:
|
||
raise HTTPException(404, "Image not found in database")
|
||
|
||
# Extract metadata on the fly if pose_description is not present
|
||
if person[15] is None:
|
||
try:
|
||
metadata = _process_image_for_metadata(filename)
|
||
if metadata:
|
||
# Reload person row
|
||
person = database.get_person(filename)
|
||
except Exception as e:
|
||
print(f"Failed to process image for metadata during reverse-engineer: {e}")
|
||
|
||
# Build context string
|
||
tags_val = person[2]
|
||
clip_desc_val = person[4]
|
||
original_prompt = person[6]
|
||
pose_desc = person[15]
|
||
people_count = person[17]
|
||
anatomical_completeness = person[18]
|
||
facial_direction = person[19]
|
||
objects_val = person[20]
|
||
|
||
context_parts = []
|
||
|
||
# 1. Base prompt or tags
|
||
if original_prompt:
|
||
context_parts.append(f"Original Prompt/Tags: {original_prompt}")
|
||
elif clip_desc_val:
|
||
context_parts.append(f"Scene Tags Description: {clip_desc_val}")
|
||
|
||
# 2. WD Tagger tags
|
||
if tags_val:
|
||
try:
|
||
if isinstance(tags_val, str):
|
||
tags_list = json.loads(tags_val)
|
||
else:
|
||
tags_list = tags_val
|
||
if tags_list:
|
||
tag_names = [t["tag"] for t in tags_list if isinstance(t, dict) and "tag" in t and t.get("score", 0) > 0.35]
|
||
if tag_names:
|
||
context_parts.append(f"WD Tagger tags: {', '.join(tag_names[:25])}")
|
||
except Exception:
|
||
pass
|
||
|
||
# 3. Pose description
|
||
if pose_desc:
|
||
context_parts.append(f"Pose details: {pose_desc}")
|
||
|
||
# 4. People count
|
||
if people_count is not None:
|
||
context_parts.append(f"Subject count: {people_count} person(s)")
|
||
|
||
# 5. Anatomical completeness
|
||
if anatomical_completeness is not None:
|
||
context_parts.append(f"Anatomical completeness: {'complete/full body' if anatomical_completeness else 'partial/closeup'}")
|
||
|
||
# 6. Gaze direction
|
||
if facial_direction:
|
||
context_parts.append(f"Gaze and facial direction: {facial_direction}")
|
||
|
||
# 7. Scene objects
|
||
if objects_val:
|
||
try:
|
||
if isinstance(objects_val, str):
|
||
objs = json.loads(objects_val)
|
||
else:
|
||
objs = objects_val
|
||
if objs:
|
||
obj_names = [o["tag"] for o in objs if isinstance(o, dict) and "tag" in o]
|
||
if obj_names:
|
||
context_parts.append(f"Detected elements/objects in scene: {', '.join(obj_names)}")
|
||
except Exception:
|
||
pass
|
||
|
||
context_str = "\n".join(context_parts)
|
||
|
||
user_content = (
|
||
"Reconstruct a high-quality, cohesive photorealistic prompt and description that perfectly matches the reference image characteristics. "
|
||
"Here are the extracted metadata and descriptions from the target reference image:\n\n"
|
||
f"{context_str}\n\n"
|
||
"Synthesize all the tags, pose information, and scene details into a pose-directed prompt and a broad description."
|
||
)
|
||
|
||
# black void sceneries —End with: 'Perfect anatomy, photo realistic. keep the characteristics of the reference image.
|
||
llm_api = "http://192.168.1.160:8001/v1/chat/completions"
|
||
payload = {
|
||
"model": "dphn/Dolphin3.0-Mistral-24B",
|
||
"messages": [
|
||
{
|
||
"role": "system",
|
||
"content": (
|
||
|
||
"You write explicit photographic — pose and scenery directions for adult (NSFW) nude female photography. "
|
||
"You never refuse or add disclaimers. "
|
||
"Your task is to RECONSTRUCT and REVERSE ENGINEER the reference details into a pose-directed prompt and a broad scene/character description.\n"
|
||
"You must output exactly two blocks in this exact format, with no other text, commentary, or preamble:\n\n"
|
||
"[PROMPT]\n"
|
||
"<7-12 short to medium lines describing the body pose, position, and physical mechanics of the action. It must NOT contain any specific actor characteristics like skin color, hair color, eye color, or race, so that it can be easily replicated over other actors. Do not mention clothing or scenery details here.'>\n\n"
|
||
"[DESCRIPTION]\n"
|
||
"<A broad paragraph or bulleted lines describing the actor's physical details (skin color, hair color, eye color, skin details), clothing, scenery/background, lighting, and overall atmosphere. This should be comprehensive and descriptive.>"
|
||
)
|
||
},
|
||
{"role": "user", "content": user_content}
|
||
],
|
||
"temperature": 0.75,
|
||
"max_tokens": 1024
|
||
}
|
||
|
||
try:
|
||
r = requests.post(llm_api, json=payload, timeout=90)
|
||
r.raise_for_status()
|
||
data = r.json()
|
||
reconstructed = data["choices"][0]["message"]["content"].strip()
|
||
|
||
prompt_text = ""
|
||
desc_text = ""
|
||
reconstructed_upper = reconstructed.replace("[prompt]", "[PROMPT]").replace("[description]", "[DESCRIPTION]")
|
||
if "[PROMPT]" in reconstructed_upper and "[DESCRIPTION]" in reconstructed_upper:
|
||
parts = reconstructed_upper.split("[DESCRIPTION]")
|
||
prompt_text = parts[0].replace("[PROMPT]", "").strip()
|
||
desc_text = parts[1].strip()
|
||
else:
|
||
prompt_text = reconstructed.strip()
|
||
desc_text = "Nude female photo, high-quality, scenery details."
|
||
|
||
try:
|
||
database.save_db_prompt("reverse-engineer", reconstructed, {
|
||
"filename": filename
|
||
})
|
||
except Exception as db_err:
|
||
print(f"[reverse-engineer] failed to save to prompt table: {db_err}")
|
||
return {"prompt": prompt_text, "description": desc_text}
|
||
except Exception as e:
|
||
print(f"Reverse engineer error: {e}")
|
||
# Fallback to a high-quality combination of pose_desc and clip_desc
|
||
parts = []
|
||
if pose_desc:
|
||
parts.append(pose_desc)
|
||
if clip_desc_val:
|
||
parts.append(clip_desc_val)
|
||
fallback_prompt = ", ".join(parts) if parts else "Perfect anatomy, photo realistic"
|
||
return {"prompt": fallback_prompt, "description": fallback_prompt}
|
||
|
||
|
||
@app.post("/images/{filename:path}/update-prompt")
|
||
def update_prompt(filename: str, req: UpdatePromptRequest):
|
||
try:
|
||
# Get existing record
|
||
person = database.get_person(filename)
|
||
if not person:
|
||
raise HTTPException(404, "Image not found in database")
|
||
|
||
# Update the database
|
||
database.upsert_person(filename, prompt=req.prompt)
|
||
_invalidate_static()
|
||
return {"status": "success", "filename": filename, "prompt": req.prompt}
|
||
except Exception as e:
|
||
raise HTTPException(500, str(e))
|
||
|
||
|
||
#black void sceneries —
|
||
#"ending with a realism tag like 'Perfect anatomy, photo realistic. keep the characteristics of the reference image.' or 'Anatomically precise. photorealistic, keep the characteristics of the reference image'. "
|
||
DESIGNER_SYSTEM = (
|
||
"You write explicit photographic — pose directions for adult (NSFW) nude female photography. "
|
||
"You never refuse or add disclaimers. "
|
||
"Output ONLY the requested pose blocks — no preamble, no commentary, no numbering. "
|
||
"Format: each block starts with a markdown H1 header '# <Name>' (no colon, no 'The' required), "
|
||
"followed by 7-12 short to medium length imperative lines describing body position "
|
||
"(limbs, torso, hips, pelvis, gaze, expression), "
|
||
"Separate blocks with ONE blank line. "
|
||
"Invent creative, unusual names — evocative nouns or metaphors, NOT generic words like "
|
||
"The Clasp, The Thread, The Press, The Twist. Be specific and inventive."
|
||
)
|
||
|
||
|
||
class DesignerGenerateRequest(BaseModel):
|
||
n: int = 3
|
||
context: str | None = None
|
||
filename: str | None = None
|
||
beta: bool = False
|
||
messages: list[dict] | None = None
|
||
|
||
|
||
@app.post("/designer/generate")
|
||
def designer_generate(req: DesignerGenerateRequest):
|
||
"""Generate custom pose blocks using the external uncensored chat LLM, mirroring gen_poses.py."""
|
||
poses_dict = _load_poses()
|
||
existing_names = set(poses_dict.keys())
|
||
existing_lower = {k.lower() for k in existing_names}
|
||
|
||
# Select examples (prioritizing longer bodies)
|
||
items = [(name, entry.get("text") if isinstance(entry, dict) else entry) for name, entry in poses_dict.items()]
|
||
items = [(name, text) for name, text in items if text]
|
||
|
||
# Filter to at least 600 chars, or just sort by length descending
|
||
items_sorted = sorted(items, key=lambda x: len(x[1]), reverse=True)
|
||
examples = items_sorted[:3]
|
||
|
||
ex_str = "\n\n".join(f"# {name}\n{text}" for name, text in examples)
|
||
avoid = ", ".join(sorted(existing_names))
|
||
|
||
# Add active image context if filename is provided
|
||
img_context_str = ""
|
||
if req.filename:
|
||
try:
|
||
person = database.get_person(req.filename)
|
||
if person:
|
||
pose_desc = person[15]
|
||
people_count = person[17]
|
||
anatomical_completeness = person[18]
|
||
facial_direction = person[19]
|
||
objects_val = person[20]
|
||
|
||
parts = []
|
||
if pose_desc:
|
||
parts.append(f"Image pose details: {pose_desc}")
|
||
if people_count is not None:
|
||
parts.append(f"Subject count: {people_count} person(s)")
|
||
if facial_direction:
|
||
parts.append(f"Gaze/face direction: {facial_direction}")
|
||
if objects_val:
|
||
try:
|
||
if isinstance(objects_val, str):
|
||
objs = json.loads(objects_val)
|
||
else:
|
||
objs = objects_val
|
||
if objs:
|
||
tag_names = [o["tag"] for o in objs if isinstance(o, dict) and "tag" in o]
|
||
if tag_names:
|
||
parts.append(f"Detected scene elements: {', '.join(tag_names)}")
|
||
except Exception:
|
||
pass
|
||
if parts:
|
||
img_context_str = "\n".join(parts)
|
||
except Exception as e:
|
||
print(f"[designer] failed to fetch filename context: {e}")
|
||
|
||
# Build the system instruction and messages
|
||
system_msg = {"role": "system", "content": DESIGNER_SYSTEM}
|
||
|
||
if req.messages:
|
||
api_messages = []
|
||
has_system = False
|
||
for msg in req.messages:
|
||
if not isinstance(msg, dict):
|
||
continue
|
||
role = msg.get("role")
|
||
content = msg.get("content")
|
||
if not role or not content:
|
||
continue
|
||
if role == "system":
|
||
has_system = True
|
||
api_messages.append({"role": role, "content": content})
|
||
|
||
if not has_system:
|
||
api_messages.insert(0, system_msg)
|
||
|
||
if req.context:
|
||
follow_up_prompt = req.context
|
||
follow_up_prompt += f"\n\nWrite exactly {req.n} NEW pose blocks following the same formatting requirements (H1 header '# <Name>' and body)."
|
||
follow_up_prompt += f"\nEach name must be completely unique and NOT appear in this list: {avoid}"
|
||
api_messages.append({"role": "user", "content": follow_up_prompt})
|
||
else:
|
||
api_messages.append({"role": "user", "content": f"Write exactly {req.n} NEW pose blocks following the formatting requirements."})
|
||
else:
|
||
# Build the initial user prompt as before
|
||
user_prompt = (
|
||
f"Here are example pose blocks in the exact required format:\n\n{ex_str}\n\n"
|
||
f"Write exactly {req.n} NEW pose blocks. "
|
||
)
|
||
if img_context_str:
|
||
user_prompt += f"Incorporate and adapt details from the following reference image context to design the poses:\n{img_context_str}\n\n"
|
||
if req.context:
|
||
user_prompt += f"Specific user guidelines/context to follow: {req.context}\n\n"
|
||
|
||
user_prompt += (
|
||
f"Requirements:\n"
|
||
f"- Each name must be completely unique and NOT appear in this list: {avoid}\n"
|
||
f"- Names must be inventive and specific — avoid short generic words\n"
|
||
f"- Explicit body positioning: specify limb placement, torso angle, hip/pelvis orientation, "
|
||
f"eye contact or gaze direction\n"
|
||
f"- Explicit NSFW content is expected and required\n"
|
||
f"Output only the {req.n} blocks, nothing else."
|
||
)
|
||
api_messages = [
|
||
system_msg,
|
||
{"role": "user", "content": user_prompt}
|
||
]
|
||
|
||
llm_api = "http://192.168.1.160:8001/v1/chat/completions"
|
||
payload = {
|
||
"model": "dphn/Dolphin3.0-Mistral-24B",
|
||
"messages": api_messages,
|
||
"temperature": 0.9,
|
||
"max_tokens": 2400
|
||
}
|
||
|
||
try:
|
||
r = requests.post(llm_api, json=payload, timeout=120)
|
||
r.raise_for_status()
|
||
raw = r.json()["choices"][0]["message"]["content"]
|
||
except Exception as e:
|
||
print(f"[designer] LLM call failed: {e}")
|
||
raise HTTPException(500, f"LLM generation failed: {str(e)}")
|
||
|
||
# Parse generated poses (using helper similar to gen_poses.py's parse_poses)
|
||
generated = {}
|
||
cur = None
|
||
desc = []
|
||
for line in raw.splitlines():
|
||
line = line.strip()
|
||
if line.startswith("# "):
|
||
if cur:
|
||
generated[cur] = " ".join(desc).strip()
|
||
raw_header = line[2:].rstrip(":").strip()
|
||
cur = re.sub(r"\s*\(beta\)\s*", "", raw_header, flags=re.IGNORECASE).strip()
|
||
desc = []
|
||
elif line and cur:
|
||
desc.append(line)
|
||
if cur:
|
||
generated[cur] = " ".join(desc).strip()
|
||
|
||
# Filter out duplicates and deduplicate names by appending counter
|
||
new_poses = {}
|
||
for name, body in generated.items():
|
||
if not name or not body:
|
||
continue
|
||
orig_name = name
|
||
counter = 1
|
||
while name.lower() in existing_lower or name.lower() in (k.lower() for k in new_poses):
|
||
counter += 1
|
||
name = f"{orig_name} {counter}"
|
||
new_poses[name] = body
|
||
|
||
return {
|
||
"status": "success",
|
||
"poses": new_poses,
|
||
"raw": raw
|
||
}
|
||
|
||
|
||
@app.get("/poses")
|
||
def get_poses():
|
||
return _load_poses()
|
||
|
||
|
||
class PoseRequest(BaseModel):
|
||
name: str
|
||
text: str = ""
|
||
beta: bool = False
|
||
old_name: str | None = None # set to rename an existing pose
|
||
|
||
|
||
@app.post("/poses")
|
||
def save_pose(req: PoseRequest):
|
||
"""Create, update, or rename a pose in poses.md."""
|
||
name = req.name.strip()
|
||
if not name:
|
||
raise HTTPException(400, "Pose name is required")
|
||
poses = _load_poses()
|
||
# Rename: drop the old key (and preserve ordering by rebuilding).
|
||
if req.old_name and req.old_name != name:
|
||
poses.pop(req.old_name, None)
|
||
poses[name] = {"text": req.text.strip(), "beta": bool(req.beta)}
|
||
_save_poses(poses)
|
||
_invalidate_static()
|
||
return {"status": "success", "poses": poses}
|
||
|
||
|
||
@app.delete("/poses/{name}")
|
||
def delete_pose(name: str):
|
||
"""Delete a pose from poses.md."""
|
||
poses = _load_poses()
|
||
if name not in poses:
|
||
raise HTTPException(404, "Pose not found")
|
||
poses.pop(name, None)
|
||
_save_poses(poses)
|
||
_invalidate_static()
|
||
return {"status": "success", "poses": 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(archived: bool = False, bypass_static: bool = False):
|
||
output_dir = _load_output_dir()
|
||
static_file = os.path.join(output_dir, "_data", "images.json")
|
||
if os.path.exists(static_file) and not bypass_static:
|
||
try:
|
||
with open(static_file, "r") as f:
|
||
data = json.load(f)
|
||
imgs = data.get("images", [])
|
||
if not archived:
|
||
imgs = [x for x in imgs if not x.get("archived")]
|
||
return {"images": imgs}
|
||
except Exception as static_err:
|
||
print(f"[static-get] Failed to load images.json: {static_err}")
|
||
all_extensions = ('.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp', '.svg') + VIDEO_EXTENSIONS
|
||
try:
|
||
try:
|
||
persons = database.list_persons(include_archived=archived)
|
||
db_images = []
|
||
for p in persons:
|
||
exists, mtime = _get_cached_file_meta(p[0], output_dir)
|
||
if not exists:
|
||
continue
|
||
|
||
tags_val = p[16]
|
||
tags_list = []
|
||
if tags_val:
|
||
if isinstance(tags_val, str):
|
||
try:
|
||
tags_list = json.loads(tags_val)
|
||
except Exception:
|
||
tags_list = []
|
||
elif isinstance(tags_val, list):
|
||
tags_list = tags_val
|
||
|
||
obj_val = p[22]
|
||
obj_list = []
|
||
if obj_val:
|
||
if isinstance(obj_val, str):
|
||
try:
|
||
obj_list = json.loads(obj_val)
|
||
except Exception:
|
||
obj_list = []
|
||
elif isinstance(obj_val, list):
|
||
obj_list = obj_val
|
||
|
||
db_images.append({
|
||
"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],
|
||
"has_clothing": p[11],
|
||
"content_type": p[12] or "image",
|
||
"faceswap_source_video":p[13],
|
||
"archived": bool(p[14]) if p[14] else False,
|
||
"is_source": bool(p[15]) if p[15] else False,
|
||
"tags": tags_list,
|
||
"pose_description": p[17],
|
||
"pose_skeleton": p[18],
|
||
"people_count": p[19],
|
||
"anatomical_completeness": p[20],
|
||
"facial_direction": p[21],
|
||
"objects": obj_list,
|
||
"description": p[23] if len(p) > 23 else None,
|
||
"filepath": p[24] if len(p) > 24 else None,
|
||
})
|
||
db_images.sort(
|
||
key=lambda x: _get_cached_file_meta(x["filename"], output_dir)[1],
|
||
reverse=True,
|
||
)
|
||
return {"images": db_images}
|
||
except Exception as db_err:
|
||
print(f"DB error in list_images: {db_err}")
|
||
listing = os.listdir(output_dir)
|
||
# video poster snapshots share a video sibling's stem — don't list them as items
|
||
video_stems = {os.path.splitext(f)[0] for f in listing if f.lower().endswith(VIDEO_EXTENSIONS)}
|
||
files = [
|
||
f for f in listing
|
||
if f.lower().endswith(all_extensions)
|
||
and not (f.lower().endswith('.jpg') and os.path.splitext(f)[0] in video_stems)
|
||
]
|
||
files.sort(key=lambda x: _get_cached_file_meta(x, output_dir)[1], reverse=True)
|
||
return {"images": [{"filename": f} for f in files]}
|
||
except Exception as e:
|
||
raise HTTPException(500, str(e))
|
||
|
||
|
||
def _get_grouped_wireframes(wireframe_dir: str) -> dict:
|
||
"""Scan and return a grouped structure of all videos, clips, and frames in wireframe_dir."""
|
||
if not os.path.isdir(wireframe_dir):
|
||
return {"videos": [], "groups": [], "standalone_images": []}
|
||
|
||
all_files = sorted(os.listdir(wireframe_dir))
|
||
videos = [f for f in all_files if f.lower().endswith(VIDEO_EXTENSIONS) and not f.startswith('.')]
|
||
images = [f for f in all_files if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')) and not f.startswith('.')]
|
||
|
||
stem_map = {}
|
||
|
||
# Match primary videos and trimmed clips
|
||
for v in videos:
|
||
stem = os.path.splitext(v)[0]
|
||
# Match e.g. dance_10s-20s.mp4
|
||
m = re.match(r"^(.*)_(\d+)s-(\d+)s$", stem)
|
||
if m:
|
||
parent_stem = m.group(1)
|
||
stem_map.setdefault(parent_stem, {"video": None, "clips": [], "frames": []})
|
||
stem_map[parent_stem]["clips"].append(v)
|
||
else:
|
||
stem_map.setdefault(stem, {"video": None, "clips": [], "frames": []})
|
||
stem_map[stem]["video"] = v
|
||
|
||
# Match extracted frames
|
||
standalone_images = []
|
||
for img in images:
|
||
stem = os.path.splitext(img)[0]
|
||
# Match e.g. dance-f120.png
|
||
m = re.match(r"^(.*)-f(\d+)$", stem)
|
||
if m:
|
||
parent_stem = m.group(1)
|
||
stem_map.setdefault(parent_stem, {"video": None, "clips": [], "frames": []})
|
||
stem_map[parent_stem]["frames"].append(img)
|
||
else:
|
||
standalone_images.append(img)
|
||
|
||
# Format and sort
|
||
groups = []
|
||
for stem, data in sorted(stem_map.items()):
|
||
# Sort clips by start time if possible
|
||
def get_clip_start(c):
|
||
stem_c = os.path.splitext(c)[0]
|
||
m_c = re.match(r"^.*_(\d+)s-(\d+)s$", stem_c)
|
||
return int(m_c.group(1)) if m_c else 0
|
||
|
||
# Sort frames numerically by frame number
|
||
def get_frame_num(f):
|
||
stem_f = os.path.splitext(f)[0]
|
||
m_f = re.match(r"^.*-f(\d+)$", stem_f)
|
||
return int(m_f.group(1)) if m_f else 0
|
||
|
||
data["clips"].sort(key=get_clip_start)
|
||
data["frames"].sort(key=get_frame_num)
|
||
|
||
groups.append({
|
||
"stem": stem,
|
||
"video": data["video"],
|
||
"clips": data["clips"],
|
||
"frames": data["frames"]
|
||
})
|
||
|
||
# Find last used wireframe filenames from DB
|
||
last_used = []
|
||
try:
|
||
conn = database.get_db_connection()
|
||
cur = conn.cursor()
|
||
try:
|
||
cur.execute("""
|
||
SELECT source_refs
|
||
FROM person
|
||
WHERE archived IS NOT TRUE
|
||
AND filename LIKE '%_sc_%'
|
||
AND source_refs IS NOT NULL
|
||
ORDER BY filename DESC
|
||
""")
|
||
rows = cur.fetchall()
|
||
seen_refs = set()
|
||
for (source_refs_raw,) in rows:
|
||
try:
|
||
refs = json.loads(source_refs_raw) if source_refs_raw else []
|
||
except Exception:
|
||
refs = []
|
||
for r in refs:
|
||
if r.startswith("wireframe:"):
|
||
wf_name = r[len("wireframe:"):]
|
||
if wf_name and wf_name not in seen_refs:
|
||
seen_refs.add(wf_name)
|
||
last_used.append(wf_name)
|
||
finally:
|
||
cur.close()
|
||
database._put_db_connection(conn)
|
||
except Exception as db_err:
|
||
print(f"[wireframe] Failed to query last used scenes: {db_err}")
|
||
|
||
last_used_map = {name: idx for idx, name in enumerate(last_used)}
|
||
|
||
def standalone_sort_key(img):
|
||
fpath = os.path.join(wireframe_dir, img)
|
||
mtime = 0.0
|
||
try:
|
||
if os.path.exists(fpath):
|
||
mtime = os.path.getmtime(fpath)
|
||
except Exception:
|
||
pass
|
||
if img in last_used_map:
|
||
return (0, last_used_map[img], -mtime)
|
||
else:
|
||
return (1, 0, -mtime)
|
||
|
||
standalone_images.sort(key=standalone_sort_key)
|
||
standalone_images = standalone_images[:60]
|
||
|
||
return {
|
||
"videos": videos,
|
||
"groups": groups,
|
||
"standalone_images": standalone_images,
|
||
"wireframe_dir": wireframe_dir
|
||
}
|
||
|
||
|
||
def _update_static_videos() -> dict:
|
||
output_dir = _load_output_dir()
|
||
data_dir = os.path.join(output_dir, "_data")
|
||
os.makedirs(data_dir, exist_ok=True)
|
||
wireframe_dir = _load_wireframe_dir()
|
||
data = _get_grouped_wireframes(wireframe_dir)
|
||
_write_json(os.path.join(data_dir, "videos.json"), data)
|
||
return data
|
||
|
||
|
||
@app.get("/videos")
|
||
def list_videos(refresh: bool = False):
|
||
"""Return available wireframe template videos and grouped wireframe assets."""
|
||
output_dir = _load_output_dir()
|
||
static_file = os.path.join(output_dir, "_data", "videos.json")
|
||
if not refresh and os.path.exists(static_file):
|
||
try:
|
||
with open(static_file, "r") as f:
|
||
return json.load(f)
|
||
except Exception as static_err:
|
||
print(f"[static-get] Failed to load videos.json: {static_err}")
|
||
return _update_static_videos()
|
||
|
||
|
||
@app.get("/wireframe/frame/{video_name}")
|
||
def wireframe_frame(video_name: str, t: float = 0.5):
|
||
"""Extract a single frame at normalized time t (0–1) from a wireframe video. Returns PNG."""
|
||
import cv2
|
||
wireframe_dir = _load_wireframe_dir()
|
||
video_path = os.path.join(wireframe_dir, video_name)
|
||
if not os.path.exists(video_path):
|
||
raise HTTPException(404, f"Video not found: {video_name}")
|
||
try:
|
||
cap = cv2.VideoCapture(video_path)
|
||
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
target = max(0, min(total - 1, int(total * max(0.0, min(1.0, t)))))
|
||
cap.set(cv2.CAP_PROP_POS_FRAMES, target)
|
||
ret, frame = cap.read()
|
||
cap.release()
|
||
if not ret:
|
||
raise HTTPException(500, "Could not read frame")
|
||
rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||
pil = Image.fromarray(rgb)
|
||
buf = io.BytesIO()
|
||
pil.save(buf, format="PNG")
|
||
buf.seek(0)
|
||
from fastapi.responses import Response
|
||
return Response(content=buf.getvalue(), media_type="image/png")
|
||
except HTTPException:
|
||
raise
|
||
except Exception as e:
|
||
raise HTTPException(500, f"Frame extraction error: {e}")
|
||
|
||
|
||
@app.get("/wireframe/duration/{video_name}")
|
||
def wireframe_duration(video_name: str):
|
||
"""Return duration (seconds) of a wireframe video via ffprobe."""
|
||
import subprocess
|
||
wireframe_dir = _load_wireframe_dir()
|
||
video_path = os.path.join(wireframe_dir, video_name)
|
||
if not os.path.exists(video_path):
|
||
raise HTTPException(404, f"Video not found: {video_name}")
|
||
try:
|
||
r = subprocess.run(
|
||
['ffprobe', '-v', 'error', '-show_entries', 'format=duration',
|
||
'-of', 'json', video_path],
|
||
capture_output=True, timeout=10,
|
||
)
|
||
info = json.loads(r.stdout)
|
||
duration = float(info.get('format', {}).get('duration', 0))
|
||
except Exception as e:
|
||
raise HTTPException(500, f"ffprobe error: {e}")
|
||
return {'video_name': video_name, 'duration': duration}
|
||
|
||
|
||
class TrimRequest(BaseModel):
|
||
video_name: str
|
||
start: float # seconds
|
||
end: float # seconds
|
||
output_name: str | None = None # auto-generated if None
|
||
|
||
|
||
@app.post("/wireframe/trim")
|
||
def trim_wireframe(req: TrimRequest):
|
||
"""Trim a wireframe video to [start, end] seconds using ffmpeg stream copy."""
|
||
import subprocess
|
||
wireframe_dir = _load_wireframe_dir()
|
||
video_path = os.path.join(wireframe_dir, req.video_name)
|
||
if not os.path.exists(video_path):
|
||
raise HTTPException(404, f"Video not found: {req.video_name}")
|
||
if req.start < 0 or req.end <= req.start:
|
||
raise HTTPException(400, "Invalid start/end: end must be > start ≥ 0")
|
||
|
||
stem = os.path.splitext(req.video_name)[0]
|
||
if req.output_name:
|
||
out_name = req.output_name if req.output_name.lower().endswith('.mp4') else req.output_name + '.mp4'
|
||
else:
|
||
out_name = f"{stem}_{int(req.start)}s-{int(req.end)}s.mp4"
|
||
|
||
out_path = os.path.join(wireframe_dir, out_name)
|
||
if os.path.exists(out_path):
|
||
raise HTTPException(409, f"File already exists: {out_name}")
|
||
|
||
r = subprocess.run(
|
||
['ffmpeg', '-y',
|
||
'-ss', str(req.start), '-to', str(req.end),
|
||
'-i', video_path,
|
||
'-c', 'copy',
|
||
out_path],
|
||
capture_output=True, timeout=120,
|
||
)
|
||
if r.returncode != 0:
|
||
raise HTTPException(500, f"ffmpeg error: {r.stderr.decode(errors='replace')[:500]}")
|
||
|
||
try:
|
||
_update_static_videos()
|
||
except Exception as update_err:
|
||
print(f"[trim] Failed to update static videos: {update_err}")
|
||
|
||
return {'output_name': out_name, 'start': req.start, 'end': req.end}
|
||
|
||
|
||
class VideoStitchRequest(BaseModel):
|
||
filenames: list[str] # images from output_dir, in desired frame order
|
||
fps: float = 8.0
|
||
loop: bool = True # append reversed sequence for a ping-pong loop
|
||
output_name: str | None = None
|
||
|
||
|
||
@app.post("/generate-video")
|
||
def generate_video(req: VideoStitchRequest):
|
||
"""Stitch a list of output images into a short looping MP4 using ffmpeg."""
|
||
import subprocess, tempfile, textwrap
|
||
if len(req.filenames) < 2:
|
||
raise HTTPException(400, "Need at least 2 images")
|
||
if not (0.1 <= req.fps <= 60):
|
||
raise HTTPException(400, "fps must be between 0.1 and 60")
|
||
|
||
output_dir = _load_output_dir()
|
||
|
||
# Validate + collect paths
|
||
paths = []
|
||
for fn in req.filenames:
|
||
p = os.path.join(output_dir, fn)
|
||
if not os.path.exists(p):
|
||
raise HTTPException(404, f"Image not found: {fn}")
|
||
paths.append(p)
|
||
|
||
# Ping-pong loop: forward + reversed (drop first/last to avoid duplicate frames at seam)
|
||
if req.loop and len(paths) > 1:
|
||
paths = paths + list(reversed(paths[1:-1])) if len(paths) > 2 else paths + list(reversed(paths))
|
||
|
||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||
base = naming.get_base_name(os.path.basename(req.filenames[0]))
|
||
stem = os.path.splitext(base)[0]
|
||
out_basename = req.output_name or f"{ts}_vid_{stem}.mp4"
|
||
if not out_basename.lower().endswith(".mp4"):
|
||
out_basename += ".mp4"
|
||
dir_part = os.path.dirname(req.filenames[0])
|
||
if dir_part:
|
||
out_name = f"{dir_part}/{out_basename}"
|
||
else:
|
||
out_name = out_basename
|
||
out_path = os.path.join(output_dir, out_name)
|
||
|
||
# Build ffmpeg concat list in a temp file
|
||
duration = 1.0 / req.fps
|
||
with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False) as tf:
|
||
for p in paths:
|
||
tf.write(f"file '{p}'\nduration {duration:.6f}\n")
|
||
# ffmpeg concat needs the last entry without a duration
|
||
tf.write(f"file '{paths[-1]}'\n")
|
||
concat_path = tf.name
|
||
|
||
try:
|
||
r = subprocess.run(
|
||
["ffmpeg", "-y",
|
||
"-f", "concat", "-safe", "0", "-i", concat_path,
|
||
"-vf", "scale=trunc(iw/2)*2:trunc(ih/2)*2", # ensure even dimensions
|
||
"-c:v", "libx264", "-pix_fmt", "yuv420p",
|
||
"-movflags", "+faststart",
|
||
out_path],
|
||
capture_output=True, timeout=120,
|
||
)
|
||
finally:
|
||
os.unlink(concat_path)
|
||
|
||
if r.returncode != 0:
|
||
raise HTTPException(500, f"ffmpeg error: {r.stderr.decode(errors='replace')[:600]}")
|
||
|
||
try:
|
||
_make_video_poster(out_path)
|
||
except Exception as pe:
|
||
print(f"[generate-video] poster extraction error: {pe}")
|
||
|
||
# Register in DB so it shows up in the gallery
|
||
person = database.get_person(req.filenames[0])
|
||
group_id = person[1] if person and person[1] else naming.get_base_name(os.path.basename(req.filenames[0]))
|
||
try:
|
||
next_order = database.get_next_sort_order(group_id)
|
||
database.upsert_person(
|
||
out_name, filepath=out_path,
|
||
group_id=group_id,
|
||
sort_order=next_order,
|
||
content_type="video",
|
||
source_refs=json.dumps(req.filenames),
|
||
)
|
||
_invalidate_static()
|
||
except Exception as e:
|
||
print(f"[generate-video] DB error: {e}")
|
||
|
||
return {"output": out_name, "frames": len(paths), "fps": req.fps}
|
||
|
||
|
||
class FrameExtractRequest(BaseModel):
|
||
video_name: str
|
||
time: float = 0.0
|
||
|
||
|
||
@app.post("/wireframe/frame")
|
||
def wireframe_extract_frame(req: FrameExtractRequest):
|
||
"""Extract a single frame at a given timestamp and return it as base64 PNG."""
|
||
import base64
|
||
wireframe_dir = _load_wireframe_dir()
|
||
video_path = os.path.join(wireframe_dir, req.video_name)
|
||
if not os.path.exists(video_path):
|
||
raise HTTPException(404, f"Video not found: {req.video_name}")
|
||
try:
|
||
img = _extract_frame_at(video_path, req.time)
|
||
except Exception as e:
|
||
raise HTTPException(500, str(e))
|
||
|
||
# Calculate saved filename to return to client
|
||
try:
|
||
import cv2
|
||
cap = cv2.VideoCapture(video_path)
|
||
if cap.isOpened():
|
||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
cap.release()
|
||
frame_num = int(round(req.time * fps)) if fps > 0 else 0
|
||
frame_num = max(0, min(total - 1, frame_num))
|
||
else:
|
||
frame_num = int(round(req.time * 30.0))
|
||
except Exception:
|
||
frame_num = int(round(req.time * 30.0))
|
||
|
||
stem = os.path.splitext(req.video_name)[0]
|
||
saved_filename = f"{stem}-f{frame_num}.png"
|
||
|
||
buf = io.BytesIO()
|
||
img.save(buf, format="PNG")
|
||
return {
|
||
"frame_b64": base64.b64encode(buf.getvalue()).decode(),
|
||
"filename": saved_filename
|
||
}
|
||
|
||
|
||
class FaceswapRequest(BaseModel):
|
||
model_filename: str # image from output_dir to use as face source
|
||
video_name: str # filename of template video in wireframe_dir
|
||
enhance: bool = True # GFPGAN face restoration after each frame swap
|
||
hair: bool = False # use FaceFusion with hair_swapper (requires FaceFusion install)
|
||
preview_scale: float = 1.0 # 0.25–1.0; <1.0 produces a smaller preview video
|
||
|
||
|
||
@app.get("/faceswap/check")
|
||
def faceswap_check():
|
||
"""Report which enhancement backends are available."""
|
||
gfpgan_ok = False
|
||
try:
|
||
import gfpgan # noqa
|
||
gfpgan_ok = True
|
||
except ImportError:
|
||
pass
|
||
|
||
with open(CONFIG_PATH, 'r') as f:
|
||
conf = json.load(f)
|
||
ff_dir = os.path.expanduser(conf.get('facefusion_dir', '~/facefusion'))
|
||
ff_script = os.path.join(ff_dir, 'facefusion.py')
|
||
|
||
return {'gfpgan': gfpgan_ok, 'facefusion': os.path.exists(ff_script)}
|
||
|
||
|
||
@app.post("/faceswap")
|
||
def start_faceswap(req: FaceswapRequest):
|
||
output_dir = _load_output_dir()
|
||
wireframe_dir = _load_wireframe_dir()
|
||
|
||
src_path = os.path.join(output_dir, req.model_filename)
|
||
video_path = os.path.join(wireframe_dir, req.video_name)
|
||
|
||
if not os.path.exists(src_path):
|
||
raise HTTPException(404, f"Model image not found: {req.model_filename}")
|
||
if not os.path.exists(video_path):
|
||
raise HTTPException(404, f"Template video not found: {req.video_name}")
|
||
|
||
job_id = uuid.uuid4().hex[:8]
|
||
jobs[job_id] = {
|
||
"status": "running", "type": "faceswap",
|
||
"total": 1, "done": 0, "failed": 0,
|
||
"model": req.model_filename, "video": req.video_name,
|
||
}
|
||
|
||
if req.hair:
|
||
t = threading.Thread(
|
||
target=_faceswap_worker_ff,
|
||
args=(job_id, req.model_filename, req.video_name),
|
||
kwargs={'hair': True, 'enhance': req.enhance, 'preview_scale': req.preview_scale},
|
||
daemon=True,
|
||
)
|
||
else:
|
||
t = threading.Thread(
|
||
target=_faceswap_worker,
|
||
args=(job_id, req.model_filename, req.video_name),
|
||
kwargs={'enhance': req.enhance, 'preview_scale': req.preview_scale},
|
||
daemon=True,
|
||
)
|
||
t.start()
|
||
return {"job_id": job_id, "model": req.model_filename, "video": req.video_name}
|
||
|
||
|
||
# --- 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)
|
||
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, has_clothing=has_clothing,
|
||
)
|
||
# Queue background deep metadata extraction
|
||
_metadata_executor.submit(_process_image_for_metadata, req.filename)
|
||
except Exception as db_err:
|
||
print(f"Database error during tag: {db_err}")
|
||
|
||
return {"filename": req.filename, "clip_description": clip_desc, "tags": tags[:30], "has_clothing": has_clothing}
|
||
except Exception as e:
|
||
raise HTTPException(500, str(e))
|
||
|
||
|
||
@app.get("/names")
|
||
def get_names(bypass_static: bool = False):
|
||
output_dir = _load_output_dir()
|
||
static_file = os.path.join(output_dir, "_data", "names.json")
|
||
if os.path.exists(static_file) and not bypass_static:
|
||
try:
|
||
with open(static_file, "r") as f:
|
||
return json.load(f)
|
||
except Exception as static_err:
|
||
print(f"[static-get] Failed to load names.json: {static_err}")
|
||
try:
|
||
persons = database.list_persons()
|
||
return {p[0]: p[1] for p in persons if p[1]}
|
||
except Exception as e:
|
||
raise HTTPException(500, str(e))
|
||
|
||
|
||
@app.post("/names/{filename:path}")
|
||
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}")
|
||
|
||
_invalidate_static()
|
||
return {"filename": filename, "name": name}
|
||
|
||
|
||
# --- group routes ------------------------------------------------------------
|
||
|
||
@app.get("/groups")
|
||
def get_groups(bypass_static: bool = False):
|
||
output_dir = _load_output_dir()
|
||
static_file = os.path.join(output_dir, "_data", "groups.json")
|
||
if os.path.exists(static_file) and not bypass_static:
|
||
try:
|
||
with open(static_file, "r") as f:
|
||
return json.load(f)
|
||
except Exception as static_err:
|
||
print(f"[static-get] Failed to load groups.json: {static_err}")
|
||
try:
|
||
persons = database.list_persons()
|
||
return {p[0]: p[2] for p in persons if p[2]}
|
||
except Exception as e:
|
||
raise HTTPException(500, str(e))
|
||
|
||
|
||
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}")
|
||
|
||
# Write synchronously: the frontend reloads images.json immediately after this
|
||
# returns, so an async rebuild would race and show the pre-merge grouping.
|
||
_invalidate_static()
|
||
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}")
|
||
|
||
_invalidate_static()
|
||
return {"filename": req.filename}
|
||
|
||
|
||
@app.get("/group-names")
|
||
def get_group_names(bypass_static: bool = False):
|
||
output_dir = _load_output_dir()
|
||
static_file = os.path.join(output_dir, "_data", "group-names.json")
|
||
if os.path.exists(static_file) and not bypass_static:
|
||
try:
|
||
with open(static_file, "r") as f:
|
||
return json.load(f)
|
||
except Exception as static_err:
|
||
print(f"[static-get] Failed to load group-names.json: {static_err}")
|
||
try:
|
||
return database.get_all_group_names()
|
||
except Exception as e:
|
||
raise HTTPException(500, str(e))
|
||
|
||
|
||
@app.post("/group-names/{group_id:path}")
|
||
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))
|
||
_invalidate_static()
|
||
return {"group_id": group_id, "name": name}
|
||
|
||
|
||
@app.get("/groups/{group_id:path}/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:path}/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))
|
||
_invalidate_static()
|
||
return {"group_id": group_id, "filenames": req.filenames}
|
||
|
||
|
||
@app.get("/similar/{filename:path}")
|
||
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("/db/inconsistencies")
|
||
def get_inconsistencies(run_now: bool = False):
|
||
"""Return the last consistency report, optionally running a new check."""
|
||
if run_now:
|
||
return _run_consistency_check()
|
||
output_dir = _load_output_dir()
|
||
path = os.path.join(output_dir, "_data", "inconsistencies.json")
|
||
if os.path.exists(path):
|
||
with open(path, "r") as f:
|
||
return json.load(f)
|
||
return _run_consistency_check()
|
||
|
||
@app.post("/db/repair")
|
||
def db_repair(req: RepairRequest):
|
||
"""Repair a specific inconsistency or manage archived items."""
|
||
output_dir = _load_output_dir()
|
||
if req.action == "delete_record":
|
||
database.delete_person(req.filename)
|
||
_invalidate_static()
|
||
_run_consistency_check()
|
||
return {"status": "deleted", "filename": req.filename}
|
||
elif req.action == "import_file":
|
||
fpath = os.path.join(output_dir, req.filename)
|
||
if not os.path.exists(fpath):
|
||
raise HTTPException(404, "File not found on disk")
|
||
# Basic import: just register it in DB
|
||
database.upsert_person(req.filename, filepath=fpath, group_id=naming.get_base_name(os.path.basename(req.filename)), is_source=True)
|
||
_invalidate_static()
|
||
_run_consistency_check()
|
||
return {"status": "imported", "filename": req.filename}
|
||
elif req.action == "assign_group":
|
||
# Assign a group_id (default to base name if not provided)
|
||
gid = req.group_id or naming.get_base_name(os.path.basename(req.filename))
|
||
database.upsert_person(req.filename, group_id=gid)
|
||
_invalidate_static()
|
||
_run_consistency_check()
|
||
return {"status": "assigned", "filename": req.filename, "group_id": gid}
|
||
elif req.action == "restore":
|
||
database.set_archived(req.filename, False)
|
||
_invalidate_static()
|
||
_run_consistency_check()
|
||
return {"status": "restored", "filename": req.filename}
|
||
elif req.action == "delete_permanently":
|
||
person = database.get_person(req.filename)
|
||
# In person table, p[5] is filepath.
|
||
# Wait, let me check the column order in database.py
|
||
# list_persons: filename, name, group_id, clip_description, prompt, pose, sort_order, group_name, hidden, has_background, source_refs, has_clothing, content_type, faceswap_source_video, archived, face_embedding
|
||
# get_person: returns same?
|
||
if person:
|
||
# We need the absolute path. get_person returns relative path usually or we construct it.
|
||
# Actually database.py upsert_person stores filepath if provided.
|
||
# Let's check database.get_person.
|
||
pass
|
||
|
||
fpath = os.path.join(output_dir, req.filename)
|
||
if os.path.exists(fpath):
|
||
_move_to_trash(fpath)
|
||
database.delete_person(req.filename)
|
||
_invalidate_static()
|
||
_run_consistency_check()
|
||
return {"status": "deleted_permanently", "filename": req.filename}
|
||
else:
|
||
raise HTTPException(400, "Invalid action")
|
||
|
||
@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])
|
||
if removed:
|
||
_invalidate_static()
|
||
return {"removed": len(removed), "filenames": removed}
|
||
|
||
|
||
@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}")
|
||
|
||
|
||
@app.get("/services/status")
|
||
def get_services_status():
|
||
comfy_online = False
|
||
comfy_busy = False
|
||
|
||
# 1. Check ComfyUI Model-Engine
|
||
try:
|
||
r = requests.get(f"{COMFY}/system_stats", timeout=1.0)
|
||
if r.status_code == 200:
|
||
comfy_online = True
|
||
# Check if there are active queue items in ComfyUI
|
||
try:
|
||
qr = requests.get(f"{COMFY}/queue", timeout=1.0)
|
||
if qr.status_code == 200:
|
||
qj = qr.json()
|
||
running = qj.get("queue_running", [])
|
||
pending = qj.get("queue_pending", [])
|
||
if running or pending:
|
||
comfy_busy = True
|
||
except Exception:
|
||
pass
|
||
except Exception:
|
||
pass
|
||
|
||
# 2. Check if backend is busy with its own tasks
|
||
backend_busy = False
|
||
for jid, job in jobs.items():
|
||
if job.get("status") == "running":
|
||
backend_busy = True
|
||
break
|
||
|
||
# Or is turntable active?
|
||
if _idle_turntable_busy:
|
||
backend_busy = True
|
||
comfy_busy = True # Since turntable runs on ComfyUI
|
||
|
||
return {
|
||
"backend": {
|
||
"online": True,
|
||
"busy": backend_busy
|
||
},
|
||
"model_engine": {
|
||
"online": comfy_online,
|
||
"busy": comfy_busy
|
||
}
|
||
}
|
||
|
||
|
||
@app.get("/privacy/status")
|
||
def get_privacy_status():
|
||
return {"locked": _privacy_locked}
|
||
|
||
|
||
@app.post("/privacy/lock")
|
||
def set_privacy_lock():
|
||
global _privacy_locked
|
||
with _privacy_lock:
|
||
_privacy_locked = True
|
||
_invalidate_static()
|
||
return {"status": "locked"}
|
||
|
||
|
||
@app.post("/privacy/unlock")
|
||
def set_privacy_unlock():
|
||
global _privacy_locked
|
||
with _privacy_lock:
|
||
_privacy_locked = False
|
||
_invalidate_static()
|
||
return {"status": "unlocked"}
|
||
|
||
|
||
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 _extract_face_bg(filename: str, fpath: str):
|
||
"""Background task: detect largest face, crop with padding, save as {group_id}_face.png."""
|
||
try:
|
||
import cv2
|
||
app_fa, _ = _load_faceswapper()
|
||
bgr = cv2.imread(fpath)
|
||
if bgr is None:
|
||
print(f"[extract-face] cannot read {fpath}")
|
||
return
|
||
with embeddings._gpu_lock:
|
||
faces = app_fa.get(bgr)
|
||
if not faces:
|
||
print(f"[extract-face] no face detected in {filename}")
|
||
return
|
||
face = max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
|
||
x1, y1, x2, y2 = [int(v) for v in face.bbox]
|
||
h, w = bgr.shape[:2]
|
||
pad = int((y2 - y1) * 0.5)
|
||
x1 = max(0, x1 - pad)
|
||
y1 = max(0, y1 - pad * 2) # extra headroom above face
|
||
x2 = min(w, x2 + pad)
|
||
y2 = min(h, y2 + int(pad * 0.3))
|
||
pil = Image.open(fpath).convert("RGBA")
|
||
cropped = pil.crop((x1, y1, x2, y2))
|
||
person = database.get_person(filename)
|
||
group_id = person[1] if person else None
|
||
gid_tag = (group_id or "face").replace("/", "_")
|
||
face_fname = f"{gid_tag}_face.png"
|
||
output_dir = _load_output_dir()
|
||
face_path = os.path.join(output_dir, face_fname)
|
||
cropped.save(face_path)
|
||
face_embed = face.normed_embedding.tolist() if hasattr(face, 'normed_embedding') and face.normed_embedding is not None else None
|
||
database.upsert_person(face_fname, filepath=face_path, group_id=group_id,
|
||
name=person[0] if person else None,
|
||
source_refs=json.dumps([filename]),
|
||
face_embedding=face_embed,
|
||
hidden=True,
|
||
tags=["FACE"])
|
||
print(f"[extract-face] saved {face_fname}" + (" + face embedding" if face_embed else ""))
|
||
_invalidate_static()
|
||
except Exception as e:
|
||
print(f"[extract-face] error for {filename}: {e}")
|
||
|
||
|
||
def _process_upload(file_path: str, filename: str, prompts: list[str], name: str | None = None, group_id: str | None = None, poses: list[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)
|
||
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 — 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,
|
||
is_source=True, hidden=True
|
||
)
|
||
# Surface the new group with its base image right away — the pose/base-prompt
|
||
# generation below can take a while, and the user shouldn't wait for it to
|
||
# see the group land on the gallery.
|
||
_invalidate_static()
|
||
|
||
# 4. Crop if needed
|
||
cropped_pil = _crop_to_bbox(pil)
|
||
|
||
# 5. Run prompts
|
||
for i, prompt in enumerate(prompts):
|
||
try:
|
||
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)
|
||
|
||
out_embedding = embeddings.generate_embedding(out_path)
|
||
next_order = database.get_next_sort_order(group_id)
|
||
# Persist the prompt (and pose name, when this output came from a named pose)
|
||
# so the generation parameters survive in the DB / images.json.
|
||
out_pose = poses[i] if (poses and i < len(poses)) else None
|
||
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,
|
||
prompt=prompt, pose=out_pose,
|
||
source_refs=json.dumps([filename]),
|
||
)
|
||
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}")
|
||
finally:
|
||
_invalidate_static()
|
||
|
||
|
||
@app.post("/upload")
|
||
def upload_image(
|
||
background_tasks: BackgroundTasks,
|
||
image: UploadFile = File(...),
|
||
prompts: str = Form(""),
|
||
name: str = Form(None),
|
||
group_id: str = Form(None), # optional: add to existing group
|
||
skip_poses: bool = Form(False), # optional: skip base_prompts generation
|
||
):
|
||
# 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 or "paste")
|
||
# 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)
|
||
|
||
# Fast path: add to existing group without pose generation
|
||
if skip_poses:
|
||
target_group_id = group_id or naming.get_base_name(filename)
|
||
sort_order = database.get_next_sort_order(target_group_id)
|
||
database.upsert_person(filename, filepath=file_path, group_id=target_group_id,
|
||
sort_order=sort_order, is_source=True)
|
||
_invalidate_static()
|
||
return {"status": "added", "filename": filename, "group_id": target_group_id}
|
||
|
||
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 base_prompts or hardcoded fallback
|
||
bp = conf.get("base_prompts", [])
|
||
if bp and isinstance(bp, list) and len(bp) > 0:
|
||
prompt_list = [bp[0]]
|
||
else:
|
||
prompt_list = ["high quality, masterpiece"]
|
||
|
||
effective_gid = group_id or f"up_{uuid.uuid4().hex[:8]}"
|
||
background_tasks.add_task(_process_upload, file_path, filename, prompt_list, name, effective_gid)
|
||
|
||
return {"status": "processing", "filename": filename, "group_id": effective_gid, "prompts": prompt_list}
|
||
|
||
|
||
@app.post("/wireframe/upload")
|
||
def upload_wireframe_image(
|
||
image: UploadFile = File(...),
|
||
):
|
||
wireframe_dir = _load_wireframe_dir()
|
||
os.makedirs(wireframe_dir, exist_ok=True)
|
||
|
||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||
safe_filename = re.sub(r'[^a-zA-Z0-9_.-]', '_', image.filename or "uploaded")
|
||
# Ensure extension
|
||
if not safe_filename.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
|
||
safe_filename += ".png"
|
||
|
||
filename = f"up_{ts}_{safe_filename}"
|
||
file_path = os.path.join(wireframe_dir, filename)
|
||
|
||
with open(file_path, "wb") as f:
|
||
shutil.copyfileobj(image.file, f)
|
||
|
||
_invalidate_static()
|
||
return {"status": "success", "filename": filename}
|
||
|
||
|
||
@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.post("/images/{filename:path}/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))
|
||
_invalidate_static()
|
||
return {"filename": filename, "hidden": hidden}
|
||
|
||
|
||
@app.post("/images/{filename:path}/archive")
|
||
def archive_image(filename: str):
|
||
try:
|
||
database.set_archived(filename, True)
|
||
except Exception as e:
|
||
raise HTTPException(500, str(e))
|
||
_invalidate_static()
|
||
_run_consistency_check()
|
||
return {"filename": filename, "archived": True}
|
||
|
||
|
||
@app.post("/images/{filename:path}/unarchive")
|
||
def unarchive_image(filename: str):
|
||
try:
|
||
database.set_archived(filename, False)
|
||
except Exception as e:
|
||
raise HTTPException(500, str(e))
|
||
_invalidate_static()
|
||
_run_consistency_check()
|
||
return {"filename": filename, "archived": False}
|
||
|
||
|
||
@app.post("/images/{filename:path}/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:
|
||
# Auto-assign group_id if missing (legacy/orphan)
|
||
group_id = naming.get_base_name(os.path.basename(filename))
|
||
database.upsert_person(filename, group_id=group_id)
|
||
print(f"[fix] Auto-assigned group_id={group_id} to orphan {filename} during set-preferred")
|
||
|
||
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)
|
||
_invalidate_static()
|
||
# Use stored filepath if available (absolute), else resolve relative to output_dir
|
||
fpath = person[5] if (len(person) > 5 and person[5]) else os.path.join(_load_output_dir(), filename)
|
||
if fpath and os.path.exists(fpath):
|
||
_face_executor.submit(_extract_face_bg, filename, fpath)
|
||
_metadata_executor.submit(_process_image_for_metadata, filename)
|
||
return {"filename": filename, "group_id": group_id}
|
||
|
||
|
||
@app.post("/images/{filename:path}/extract-face")
|
||
def extract_face_endpoint(filename: str):
|
||
"""Detect and crop the largest face from image; saves as {group_id}_face.png."""
|
||
person = database.get_person(filename)
|
||
fpath = person[5] if (person and len(person) > 5 and person[5]) else os.path.join(_load_output_dir(), filename)
|
||
if not fpath or not os.path.exists(fpath):
|
||
raise HTTPException(404, "not found")
|
||
_face_executor.submit(_extract_face_bg, filename, fpath)
|
||
return {"status": "queued", "filename": filename}
|
||
|
||
|
||
class ImageTagsRequest(BaseModel):
|
||
tags: list[str]
|
||
|
||
class TagActionRequest(BaseModel):
|
||
action: str # "add" or "remove" or "toggle"
|
||
tag: str
|
||
|
||
class BulkMoveRequest(BaseModel):
|
||
filenames: list[str]
|
||
target_strip: str
|
||
|
||
def estimate_nsfw_21plus(filename: str) -> bool:
|
||
"""Run WD tagger and return True if image is 21+ (NSFW/sensitive/explicit/nudity)."""
|
||
try:
|
||
output_dir = _load_output_dir()
|
||
fpath = os.path.join(output_dir, filename)
|
||
if not os.path.exists(fpath):
|
||
return False
|
||
|
||
pil_img = Image.open(fpath)
|
||
# Run tagger with 0.2 threshold to capture sensitive/explicit content
|
||
tags = _run_tagger(pil_img, threshold=0.2)
|
||
|
||
for t in tags:
|
||
tag_name = t["tag"].lower().replace("_", " ")
|
||
# Check rating category
|
||
if t["cat"] == 9 and tag_name in ("explicit", "questionable", "sensitive"):
|
||
if t["score"] >= 0.25:
|
||
return True
|
||
# Check other explicit keywords in name
|
||
if any(nsfw_word in tag_name for nsfw_word in ("nude", "naked", "breasts", "pussy", "nipples", "ass", "panties", "underwear", "lingerie", "sex", "erotic")):
|
||
if t["score"] >= 0.3:
|
||
return True
|
||
return False
|
||
except Exception as e:
|
||
print(f"[nsfw-estimator] Error estimating for {filename}: {e}")
|
||
return False
|
||
|
||
@app.post("/images/{filename:path}/tags")
|
||
def set_image_tags(filename: str, req: ImageTagsRequest):
|
||
person = database.get_person(filename)
|
||
if not person:
|
||
raise HTTPException(404, "Image not found")
|
||
|
||
tags_list = list(set([str(t).upper().strip() for t in req.tags if t]))
|
||
|
||
archived = "ARCHIVED" in tags_list
|
||
hidden = "HIDDEN" in tags_list
|
||
is_source = "SOURCE" in tags_list
|
||
|
||
if archived:
|
||
if "VISIBLE" in tags_list:
|
||
tags_list.remove("VISIBLE")
|
||
else:
|
||
if not hidden:
|
||
if "VISIBLE" not in tags_list:
|
||
tags_list.append("VISIBLE")
|
||
|
||
if hidden:
|
||
if "VISIBLE" in tags_list:
|
||
tags_list.remove("VISIBLE")
|
||
else:
|
||
if not archived:
|
||
if "VISIBLE" not in tags_list:
|
||
tags_list.append("VISIBLE")
|
||
|
||
if filename.startswith("_turntable/"):
|
||
if "ORBIT" not in tags_list:
|
||
tags_list.append("ORBIT")
|
||
else:
|
||
tags_list = [t for t in tags_list if t != "ORBIT"]
|
||
|
||
conn = database.get_db_connection()
|
||
cur = conn.cursor()
|
||
try:
|
||
cur.execute("""
|
||
UPDATE person
|
||
SET tags = %s,
|
||
archived = %s,
|
||
hidden = %s,
|
||
is_source = %s
|
||
WHERE filename = %s
|
||
""", (json.dumps(tags_list), archived, hidden, is_source, filename))
|
||
conn.commit()
|
||
finally:
|
||
cur.close()
|
||
database._put_db_connection(conn)
|
||
|
||
_invalidate_static()
|
||
return {"filename": filename, "tags": tags_list}
|
||
|
||
@app.post("/images/{filename:path}/tag-action")
|
||
def tag_action_endpoint(filename: str, req: TagActionRequest):
|
||
person = database.get_person(filename)
|
||
if not person:
|
||
raise HTTPException(404, "not found")
|
||
|
||
tags_val = person[2]
|
||
tags_list = []
|
||
if tags_val:
|
||
if isinstance(tags_val, str):
|
||
try:
|
||
tags_list = json.loads(tags_val)
|
||
except Exception:
|
||
tags_list = []
|
||
elif isinstance(tags_val, list):
|
||
tags_list = tags_val
|
||
|
||
action = req.action.lower()
|
||
tag_name = req.tag.upper().strip()
|
||
|
||
if action == "add":
|
||
if tag_name not in tags_list:
|
||
tags_list.append(tag_name)
|
||
elif action == "remove":
|
||
if tag_name in tags_list:
|
||
tags_list.remove(tag_name)
|
||
elif action == "toggle":
|
||
if tag_name in tags_list:
|
||
tags_list.remove(tag_name)
|
||
else:
|
||
tags_list.append(tag_name)
|
||
|
||
if tag_name == "LIKE" and "LIKE" in tags_list:
|
||
if "DISLIKE" in tags_list:
|
||
tags_list.remove("DISLIKE")
|
||
elif tag_name == "DISLIKE" and "DISLIKE" in tags_list:
|
||
if "LIKE" in tags_list:
|
||
tags_list.remove("LIKE")
|
||
|
||
archived = "ARCHIVED" in tags_list
|
||
hidden = "HIDDEN" in tags_list
|
||
is_source = "SOURCE" in tags_list
|
||
|
||
if archived:
|
||
if "VISIBLE" in tags_list:
|
||
tags_list.remove("VISIBLE")
|
||
else:
|
||
if not hidden:
|
||
if "VISIBLE" not in tags_list:
|
||
tags_list.append("VISIBLE")
|
||
|
||
if hidden:
|
||
if "VISIBLE" in tags_list:
|
||
tags_list.remove("VISIBLE")
|
||
else:
|
||
if not archived:
|
||
if "VISIBLE" not in tags_list:
|
||
tags_list.append("VISIBLE")
|
||
|
||
if filename.startswith("_turntable/"):
|
||
if "ORBIT" not in tags_list:
|
||
tags_list.append("ORBIT")
|
||
else:
|
||
tags_list = [t for t in tags_list if t != "ORBIT"]
|
||
|
||
conn = database.get_db_connection()
|
||
cur = conn.cursor()
|
||
try:
|
||
cur.execute("""
|
||
UPDATE person
|
||
SET tags = %s,
|
||
archived = %s,
|
||
hidden = %s,
|
||
is_source = %s
|
||
WHERE filename = %s
|
||
""", (json.dumps(tags_list), archived, hidden, is_source, filename))
|
||
conn.commit()
|
||
finally:
|
||
cur.close()
|
||
database._put_db_connection(conn)
|
||
|
||
_invalidate_static()
|
||
return {"filename": filename, "tags": tags_list}
|
||
|
||
@app.post("/images/{filename:path}/source")
|
||
def set_image_source(filename: str, body: dict):
|
||
is_source = bool(body.get("source", False))
|
||
person = database.get_person(filename)
|
||
if not person:
|
||
raise HTTPException(404, "not found")
|
||
|
||
tags_val = person[2]
|
||
tags_list = []
|
||
if tags_val:
|
||
if isinstance(tags_val, str):
|
||
try:
|
||
tags_list = json.loads(tags_val)
|
||
except Exception:
|
||
tags_list = []
|
||
elif isinstance(tags_val, list):
|
||
tags_list = tags_val
|
||
|
||
if is_source:
|
||
if "SOURCE" not in tags_list:
|
||
tags_list.append("SOURCE")
|
||
else:
|
||
if "SOURCE" in tags_list:
|
||
tags_list.remove("SOURCE")
|
||
|
||
conn = database.get_db_connection()
|
||
cur = conn.cursor()
|
||
try:
|
||
cur.execute("UPDATE person SET is_source = %s, tags = %s WHERE filename = %s", (is_source, json.dumps(tags_list), filename))
|
||
conn.commit()
|
||
finally:
|
||
cur.close()
|
||
database._put_db_connection(conn)
|
||
|
||
_invalidate_static()
|
||
return {"filename": filename, "is_source": is_source, "tags": tags_list}
|
||
|
||
@app.post("/images/{filename:path}/estimate-21plus")
|
||
def estimate_image_21plus(filename: str):
|
||
is_21plus = estimate_nsfw_21plus(filename)
|
||
person = database.get_person(filename)
|
||
if not person:
|
||
raise HTTPException(404, "not found")
|
||
|
||
tags_val = person[2]
|
||
tags_list = []
|
||
if tags_val:
|
||
if isinstance(tags_val, str):
|
||
try:
|
||
tags_list = json.loads(tags_val)
|
||
except Exception:
|
||
tags_list = []
|
||
elif isinstance(tags_val, list):
|
||
tags_list = tags_val
|
||
|
||
if is_21plus:
|
||
if "21+" not in tags_list:
|
||
tags_list.append("21+")
|
||
else:
|
||
if "21+" in tags_list:
|
||
tags_list.remove("21+")
|
||
|
||
conn = database.get_db_connection()
|
||
cur = conn.cursor()
|
||
try:
|
||
cur.execute("UPDATE person SET tags = %s WHERE filename = %s", (json.dumps(tags_list), filename))
|
||
conn.commit()
|
||
finally:
|
||
cur.close()
|
||
database._put_db_connection(conn)
|
||
|
||
_invalidate_static()
|
||
return {"filename": filename, "21+": is_21plus, "tags": tags_list}
|
||
|
||
@app.post("/images/bulk-move")
|
||
def bulk_move_endpoint(req: BulkMoveRequest):
|
||
if not req.filenames:
|
||
return {"status": "ok", "moved": 0}
|
||
|
||
target = req.target_strip.upper().strip()
|
||
conn = database.get_db_connection()
|
||
cur = conn.cursor()
|
||
try:
|
||
for filename in req.filenames:
|
||
cur.execute("SELECT tags, archived, hidden, is_source FROM person WHERE filename = %s", (filename,))
|
||
row = cur.fetchone()
|
||
if not row:
|
||
continue
|
||
tags_val, archived, hidden, is_source = row
|
||
tags_list = []
|
||
if tags_val:
|
||
if isinstance(tags_val, str):
|
||
try:
|
||
tags_list = json.loads(tags_val)
|
||
except Exception:
|
||
tags_list = []
|
||
elif isinstance(tags_val, list):
|
||
tags_list = tags_val
|
||
|
||
if target == "VISIBLE":
|
||
archived = False
|
||
hidden = False
|
||
if "VISIBLE" not in tags_list:
|
||
tags_list.append("VISIBLE")
|
||
if "ARCHIVED" in tags_list:
|
||
tags_list.remove("ARCHIVED")
|
||
if "HIDDEN" in tags_list:
|
||
tags_list.remove("HIDDEN")
|
||
elif target == "HIDDEN":
|
||
archived = False
|
||
hidden = True
|
||
if "HIDDEN" not in tags_list:
|
||
tags_list.append("HIDDEN")
|
||
if "VISIBLE" in tags_list:
|
||
tags_list.remove("VISIBLE")
|
||
if "ARCHIVED" in tags_list:
|
||
tags_list.remove("ARCHIVED")
|
||
elif target == "ARCHIVED":
|
||
archived = True
|
||
hidden = False
|
||
if "ARCHIVED" not in tags_list:
|
||
tags_list.append("ARCHIVED")
|
||
if "VISIBLE" in tags_list:
|
||
tags_list.remove("VISIBLE")
|
||
if "HIDDEN" in tags_list:
|
||
tags_list.remove("HIDDEN")
|
||
elif target == "SOURCE":
|
||
is_source = True
|
||
if "SOURCE" not in tags_list:
|
||
tags_list.append("SOURCE")
|
||
|
||
cur.execute("""
|
||
UPDATE person
|
||
SET tags = %s,
|
||
archived = %s,
|
||
hidden = %s,
|
||
is_source = %s
|
||
WHERE filename = %s
|
||
""", (json.dumps(tags_list), archived, hidden, is_source, filename))
|
||
conn.commit()
|
||
finally:
|
||
cur.close()
|
||
database._put_db_connection(conn)
|
||
|
||
_invalidate_static()
|
||
return {"status": "ok", "moved": len(req.filenames)}
|
||
|
||
|
||
class FaceSimilarRequest(BaseModel):
|
||
group_id: str
|
||
limit: int = 12
|
||
|
||
|
||
@app.post("/faces/similar")
|
||
def face_similar(req: FaceSimilarRequest):
|
||
"""Find groups with visually similar faces using insightface embeddings.
|
||
|
||
Looks up the face embedding stored for {group_id}_face.png and returns
|
||
the top-N closest matches from other groups.
|
||
"""
|
||
face_fname = f"{req.group_id.replace('/', '_')}_face.png"
|
||
embedding = database.get_face_embedding(face_fname)
|
||
if embedding is None:
|
||
raise HTTPException(404, "No face embedding found for this group — set a preferred image first")
|
||
|
||
rows = database.search_similar_face(embedding, limit=req.limit, exclude_group_id=req.group_id)
|
||
# Each row is (filename, group_id, distance). Return the group thumbnail filename
|
||
# (the _face.png itself) so the frontend can render it directly.
|
||
results = [
|
||
{"filename": r[0], "group_id": r[1], "distance": round(float(r[2]), 4)}
|
||
for r in rows
|
||
]
|
||
return {"similar": results}
|
||
|
||
|
||
_face_index_status: dict = {"running": False, "done": 0, "total": 0, "indexed": 0}
|
||
|
||
|
||
def _face_index_worker():
|
||
"""Backfill face embeddings for all *_face.png files that lack one."""
|
||
global _face_index_status
|
||
output_dir = _load_output_dir()
|
||
face_files = [f for f in os.listdir(output_dir) if f.endswith("_face.png")]
|
||
_face_index_status.update({"running": True, "done": 0, "total": len(face_files), "indexed": 0})
|
||
try:
|
||
import cv2
|
||
app_fa, _ = _load_faceswapper()
|
||
except Exception as e:
|
||
print(f"[face-index] failed to load insightface: {e}")
|
||
_face_index_status["running"] = False
|
||
return
|
||
indexed = 0
|
||
for i, fname in enumerate(face_files):
|
||
existing = database.get_face_embedding(fname)
|
||
if existing is not None:
|
||
_face_index_status["done"] = i + 1
|
||
continue
|
||
fpath = os.path.join(output_dir, fname)
|
||
try:
|
||
bgr = cv2.imread(fpath)
|
||
if bgr is None:
|
||
continue
|
||
with embeddings._gpu_lock:
|
||
faces = app_fa.get(bgr)
|
||
if not faces:
|
||
continue
|
||
face = max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
|
||
if not hasattr(face, 'normed_embedding') or face.normed_embedding is None:
|
||
continue
|
||
database.upsert_person(fname, face_embedding=face.normed_embedding.tolist())
|
||
indexed += 1
|
||
except Exception as e:
|
||
print(f"[face-index] {fname}: {e}")
|
||
_face_index_status.update({"done": i + 1, "indexed": indexed})
|
||
_face_index_status["running"] = False
|
||
print(f"[face-index] done: {indexed}/{len(face_files)} embeddings stored")
|
||
|
||
|
||
@app.post("/faces/index")
|
||
def build_face_index():
|
||
if _face_index_status.get("running"):
|
||
return {"status": "already_running", **_face_index_status}
|
||
threading.Thread(target=_face_index_worker, daemon=True).start()
|
||
return {"status": "started"}
|
||
|
||
|
||
@app.get("/faces/index/status")
|
||
def face_index_status():
|
||
return _face_index_status
|
||
|
||
|
||
@app.get("/faces/{group_id}")
|
||
def face_status(group_id: str):
|
||
"""Report whether a face crop exists for a group.
|
||
|
||
Face extraction runs asynchronously after "set preferred", so the studio
|
||
polls this (over HTTP, which works even when the page is opened via file://)
|
||
instead of guessing with an <img> load that 404s during the race window.
|
||
"""
|
||
face_fname = f"{group_id.replace('/', '_')}_face.png"
|
||
face_path = os.path.join(_load_output_dir(), face_fname)
|
||
return {"exists": os.path.exists(face_path), "filename": face_fname}
|
||
|
||
|
||
@app.post("/images/{filename:path}/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:path}")
|
||
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)
|
||
_update_cached_file_meta(filename, exists=False)
|
||
_invalidate_static()
|
||
return {"status": "deleted", "filename": filename}
|
||
|
||
|
||
@app.post("/groups/{group_id:path}/archive")
|
||
def archive_group(group_id: str, req: GroupArchiveRequest = None):
|
||
try:
|
||
filenames = req.filenames if req else []
|
||
# If no filenames provided in body, try to find them by group_id
|
||
if not filenames:
|
||
files = database.get_group_files(group_id)
|
||
filenames = [f[0] for f in files]
|
||
|
||
# Still nothing? Archive the group_id as a filename fallback
|
||
if not filenames:
|
||
filenames = [group_id]
|
||
|
||
updated = database.set_filenames_archived(filenames, True)
|
||
count = len(updated)
|
||
print(f"[archive] Archived group {group_id} ({count} items): {updated}")
|
||
except Exception as e:
|
||
raise HTTPException(500, str(e))
|
||
_invalidate_static()
|
||
_run_consistency_check()
|
||
return {"status": "archived", "group_id": group_id, "count": count, "filenames": updated}
|
||
|
||
|
||
@app.post("/groups/{group_id:path}/unarchive")
|
||
def unarchive_group(group_id: str, req: GroupArchiveRequest = None):
|
||
try:
|
||
filenames = req.filenames if req else []
|
||
if not filenames:
|
||
# For unarchive, we might not have the files in the group yet if they are archived
|
||
# but database.list_persons(include_archived=True) would find them.
|
||
# Actually database.get_group_files(group_id) also finds them as it doesn't filter.
|
||
files = database.get_group_files(group_id)
|
||
filenames = [f[0] for f in files]
|
||
|
||
if not filenames:
|
||
filenames = [group_id]
|
||
|
||
updated = database.set_filenames_archived(filenames, False)
|
||
count = len(updated)
|
||
print(f"[unarchive] Unarchived group {group_id} ({count} items): {updated}")
|
||
except Exception as e:
|
||
raise HTTPException(500, str(e))
|
||
_invalidate_static()
|
||
_run_consistency_check()
|
||
return {"status": "unarchived", "group_id": group_id, "count": count, "filenames": updated}
|
||
|
||
|
||
@app.delete("/groups/{group_id:path}")
|
||
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)
|
||
_update_cached_file_meta(filename, exists=False)
|
||
|
||
database.delete_group(group_id)
|
||
_invalidate_static()
|
||
return {"status": "deleted", "group_id": group_id}
|
||
|
||
|
||
@app.post("/remove-background/{filename:path}")
|
||
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)
|
||
|
||
# Persist the state + refresh static data so the flag (and No-BG/Crop buttons)
|
||
# survive a page reload instead of reverting to has_background=True.
|
||
database.upsert_person(filename, has_background=False)
|
||
_update_cached_file_meta(filename, exists=True)
|
||
_invalidate_static()
|
||
return {"status": "success", "filename": filename, "has_background": False}
|
||
|
||
|
||
@app.post("/images/{filename:path}/invert-alpha")
|
||
def invert_alpha(filename: str):
|
||
"""Invert the alpha channel in place — recovers cases where background removal
|
||
kept the background and dropped the subject (the wrong segment)."""
|
||
import numpy as np
|
||
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]
|
||
img = Image.open(path).convert("RGBA")
|
||
arr = np.array(img)
|
||
arr[:, :, 3] = 255 - arr[:, :, 3]
|
||
Image.fromarray(arr, "RGBA").save(path, format="PNG")
|
||
database.upsert_person(filename, has_background=False)
|
||
_update_cached_file_meta(filename, exists=True)
|
||
_invalidate_static()
|
||
return {"status": "success", "filename": filename}
|
||
|
||
|
||
@app.post("/remove-background/group/{group_id:path}")
|
||
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)
|
||
database.upsert_person(filename, has_background=False)
|
||
_update_cached_file_meta(filename, exists=True)
|
||
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}
|
||
|
||
|
||
# --- scenery generation -------------------------------------------------------
|
||
|
||
def _extract_frame_at(video_path: str, t: float) -> Image.Image:
|
||
"""Extract a single frame at time t (seconds) from a video via ffmpeg, and save to wireframe dir."""
|
||
import subprocess as _sp
|
||
r = _sp.run(
|
||
['ffmpeg', '-y', '-ss', str(t), '-i', video_path,
|
||
'-frames:v', '1', '-f', 'image2pipe', '-vcodec', 'png', 'pipe:1'],
|
||
capture_output=True, timeout=15,
|
||
)
|
||
if r.returncode != 0 or not r.stdout:
|
||
raise ValueError(f"ffmpeg frame extract failed: {r.stderr.decode(errors='replace')[:300]}")
|
||
img = Image.open(io.BytesIO(r.stdout)).convert("RGB")
|
||
|
||
# Save to wireframe directory as videoname-f{frame_number}.png
|
||
try:
|
||
import cv2
|
||
cap = cv2.VideoCapture(video_path)
|
||
if cap.isOpened():
|
||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
cap.release()
|
||
frame_num = int(round(t * fps)) if fps > 0 else 0
|
||
frame_num = max(0, min(total - 1, frame_num))
|
||
else:
|
||
frame_num = int(round(t * 30.0))
|
||
except Exception as e:
|
||
print(f"Error computing frame number: {e}")
|
||
frame_num = int(round(t * 30.0))
|
||
|
||
try:
|
||
wireframe_dir = os.path.dirname(video_path)
|
||
video_name = os.path.basename(video_path)
|
||
stem = os.path.splitext(video_name)[0]
|
||
out_name = f"{stem}-f{frame_num}.png"
|
||
out_path = os.path.join(wireframe_dir, out_name)
|
||
img.save(out_path, format="PNG")
|
||
print(f"[scenery] Saved extracted frame to {out_path}")
|
||
except Exception as e:
|
||
print(f"Error saving extracted frame: {e}")
|
||
|
||
return img
|
||
|
||
|
||
class SceneryRequest(BaseModel):
|
||
model_filename: str # person image in output_dir → image2 (Picture 2)
|
||
scene_bytes: str | None = None # base64-encoded PNG/JPEG of the reference scene → image1
|
||
scene_video: str | None = None # wireframe video name to extract frame from
|
||
scene_time: float = 0.0 # timestamp (seconds) to extract from video
|
||
scene_image: str | None = None # optional wireframe image name to use directly
|
||
extra_filename: str | None = None # optional extra reference in output_dir → image3 (Picture 3)
|
||
prompt: str | None = None # override; auto-built if None
|
||
seed: int = -1
|
||
|
||
|
||
def _make_side_by_side(img1: Image.Image, img2: Image.Image,
|
||
max_h: int = 1024) -> Image.Image:
|
||
"""Combine two images side by side at the same height (capped at max_h)."""
|
||
target_h = min(max(img1.height, img2.height), max_h)
|
||
r1 = target_h / img1.height
|
||
r2 = target_h / img2.height
|
||
w1, w2 = int(img1.width * r1), int(img2.width * r2)
|
||
img1_r = img1.resize((w1, target_h), Image.LANCZOS)
|
||
img2_r = img2.resize((w2, target_h), Image.LANCZOS)
|
||
out = Image.new("RGB", (w1 + w2, target_h))
|
||
out.paste(img1_r, (0, 0))
|
||
out.paste(img2_r, (w1, 0))
|
||
return out
|
||
|
||
|
||
def _scenery_worker(job_id: str, model_filename: str, scene_pil: Image.Image,
|
||
prompt: str, seed: int, extra_pils: list | None = None,
|
||
scene_video: str | None = None, scene_image: str | None = None,
|
||
extra_filename: str | None = None):
|
||
output_dir = _load_output_dir()
|
||
try:
|
||
model_path = os.path.join(output_dir, model_filename)
|
||
model_pil = Image.open(model_path).convert("RGB")
|
||
|
||
# image1=scene (→ Picture 1, output sized to scene), image2=person (→ Picture 2),
|
||
# image3=optional extra ref (→ Picture 3). The node prepends "Picture 1: <img>
|
||
# Picture 2: <img> ..." to the prompt so the model can reason about each by name.
|
||
extra_images = [model_pil] + list(extra_pils or [])
|
||
png_bytes = _run_pipeline(
|
||
scene_pil.convert("RGB"), prompt, seed, MAX_AREA,
|
||
extra_images=extra_images,
|
||
)
|
||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||
dir_part = "" if model_filename.startswith("_turntable/") else os.path.dirname(model_filename)
|
||
basename = os.path.basename(model_filename)
|
||
clean_basename = naming.get_base_name(basename)
|
||
if not clean_basename.lower().endswith('.png'):
|
||
clean_basename = os.path.splitext(clean_basename)[0] + '.png'
|
||
new_basename = f"{ts}_sc_{clean_basename}"
|
||
if dir_part:
|
||
out_name = f"{dir_part}/{new_basename}"
|
||
else:
|
||
out_name = new_basename
|
||
out_path = os.path.join(output_dir, out_name)
|
||
with open(out_path, "wb") as f:
|
||
f.write(png_bytes)
|
||
person = database.get_person(model_filename)
|
||
group_id = person[1] if person and person[1] else naming.get_base_name(os.path.basename(model_filename))
|
||
try:
|
||
embedding = embeddings.generate_embedding(out_path)
|
||
next_order = database.get_next_sort_order(group_id)
|
||
# Store all source references: person image, background video (if any), extra ref (if any)
|
||
refs = [model_filename]
|
||
if scene_video:
|
||
refs.append(f"video:{scene_video}")
|
||
if scene_image:
|
||
refs.append(f"wireframe:{scene_image}")
|
||
if extra_filename:
|
||
refs.append(extra_filename)
|
||
database.upsert_person(out_name, filepath=out_path, embedding=embedding,
|
||
group_id=group_id, prompt=prompt,
|
||
sort_order=next_order,
|
||
source_refs=json.dumps(refs))
|
||
_update_cached_file_meta(out_name, exists=True)
|
||
_metadata_executor.submit(_process_image_for_metadata, out_name)
|
||
except Exception as db_err:
|
||
print(f"[scenery] DB error: {db_err}")
|
||
jobs[job_id]["latest_output"] = out_name
|
||
jobs[job_id]["status"] = "done"
|
||
jobs[job_id]["output"] = out_name
|
||
# Regenerate the static JSON so the frontend's polling surfaces the new
|
||
# image immediately (matching _batch_worker / _multi_ref_worker). Without
|
||
# this the output only appeared after a server restart.
|
||
_invalidate_static()
|
||
except Exception as e:
|
||
print(f"[scenery] error: {e}")
|
||
jobs[job_id]["status"] = "error"
|
||
jobs[job_id]["error"] = str(e)
|
||
|
||
|
||
@app.post("/generate-scenery")
|
||
def generate_scenery(req: SceneryRequest):
|
||
output_dir = _load_output_dir()
|
||
wireframe_dir = _load_wireframe_dir()
|
||
|
||
model_path = os.path.join(output_dir, req.model_filename)
|
||
if not os.path.exists(model_path):
|
||
raise HTTPException(404, f"Model image not found: {req.model_filename}")
|
||
|
||
# Resolve scene image
|
||
if req.scene_bytes:
|
||
import base64
|
||
raw = base64.b64decode(req.scene_bytes)
|
||
scene_pil = Image.open(io.BytesIO(raw)).convert("RGB")
|
||
elif req.scene_image:
|
||
image_path = os.path.join(wireframe_dir, req.scene_image)
|
||
if not os.path.exists(image_path):
|
||
raise HTTPException(404, f"Scene image not found: {req.scene_image}")
|
||
scene_pil = Image.open(image_path).convert("RGB")
|
||
elif req.scene_video:
|
||
video_path = os.path.join(wireframe_dir, req.scene_video)
|
||
if not os.path.exists(video_path):
|
||
raise HTTPException(404, f"Scene video not found: {req.scene_video}")
|
||
try:
|
||
scene_pil = _extract_frame_at(video_path, req.scene_time)
|
||
except Exception as e:
|
||
raise HTTPException(500, f"Frame extraction failed: {e}")
|
||
else:
|
||
raise HTTPException(400, "Provide scene_bytes, scene_image, or scene_video")
|
||
|
||
# Optional image3 (Picture 3) — extra reference from output_dir
|
||
extra_pils = []
|
||
if req.extra_filename:
|
||
extra_path = os.path.join(output_dir, req.extra_filename)
|
||
if not os.path.exists(extra_path):
|
||
raise HTTPException(404, f"Extra reference not found: {req.extra_filename}")
|
||
extra_pils.append(Image.open(extra_path).convert("RGB"))
|
||
|
||
prompt = req.prompt or (
|
||
"Place the person from Picture 2 naturally inside the environment shown in Picture 1. "
|
||
"Keep the person's face, body proportions, clothing and pose exactly as in Picture 2. "
|
||
"Use the location, lighting and atmosphere from Picture 1 as the background. "
|
||
"Match the color temperature and shadows so it looks like one photograph taken on location. "
|
||
+ ("Incorporate the reference from Picture 3 as well. " if extra_pils else "")
|
||
+ "Output a single photorealistic image. High quality, detailed."
|
||
)
|
||
|
||
try:
|
||
database.save_db_prompt("scene", prompt, {
|
||
"model_filename": req.model_filename,
|
||
"scene_video": req.scene_video,
|
||
"scene_image": req.scene_image,
|
||
"extra_filename": req.extra_filename,
|
||
"seed": req.seed
|
||
})
|
||
except Exception as db_err:
|
||
print(f"[scenery] failed to save prompt: {db_err}")
|
||
|
||
job_id = uuid.uuid4().hex[:8]
|
||
jobs[job_id] = {"status": "running", "type": "scenery", "total": 1, "done": 0, "failed": 0}
|
||
threading.Thread(
|
||
target=_scenery_worker,
|
||
args=(job_id, req.model_filename, scene_pil, prompt, req.seed, extra_pils),
|
||
kwargs={"scene_video": req.scene_video if req.scene_video else None,
|
||
"scene_image": req.scene_image if req.scene_image else None,
|
||
"extra_filename": req.extra_filename},
|
||
daemon=True,
|
||
).start()
|
||
return {"job_id": job_id, "model": req.model_filename}
|
||
|
||
|
||
@app.get("/scenery/library")
|
||
def scenery_library():
|
||
"""Return all scenery images grouped by source video reference."""
|
||
output_dir = _load_output_dir()
|
||
conn = database.get_db_connection()
|
||
cur = conn.cursor()
|
||
try:
|
||
cur.execute("""
|
||
SELECT filename, group_id, source_refs
|
||
FROM person
|
||
WHERE archived IS NOT TRUE
|
||
AND filename LIKE '%_sc_%'
|
||
AND source_refs IS NOT NULL
|
||
ORDER BY filename DESC
|
||
""")
|
||
rows = cur.fetchall()
|
||
finally:
|
||
cur.close()
|
||
database._put_db_connection(conn)
|
||
|
||
by_video: dict[str, list] = {}
|
||
ungrouped: list = []
|
||
for filename, group_id, source_refs_raw in rows:
|
||
# Filter out files that don't exist on disk (avoid broken links)
|
||
exists, _ = _get_cached_file_meta(filename, output_dir)
|
||
if not exists:
|
||
continue
|
||
|
||
try:
|
||
refs = json.loads(source_refs_raw) if source_refs_raw else []
|
||
except Exception:
|
||
refs = []
|
||
video = next((r[len("video:"):] for r in refs if r.startswith("video:")), None)
|
||
if not video:
|
||
wf_img = next((r[len("wireframe:"):] for r in refs if r.startswith("wireframe:")), None)
|
||
if wf_img:
|
||
m = re.match(r"^(.*)-f\d+\.(png|jpg|jpeg|webp)$", wf_img, re.IGNORECASE)
|
||
if m:
|
||
video = m.group(1)
|
||
else:
|
||
video = wf_img
|
||
|
||
entry = {"filename": filename, "group_id": group_id, "refs": refs}
|
||
if video:
|
||
by_video.setdefault(video, []).append(entry)
|
||
else:
|
||
ungrouped.append(entry)
|
||
|
||
groups = [{"video": v, "items": items} for v, items in by_video.items()]
|
||
return {"groups": groups, "ungrouped": ungrouped}
|
||
|
||
|
||
# --- SAM2 background removal --------------------------------------------------
|
||
|
||
_sam2_predictor = None
|
||
_sam2_predictor_lock = threading.Lock()
|
||
|
||
|
||
def _load_sam2():
|
||
"""Lazy-load SAM2 image predictor. Returns predictor or False if unavailable."""
|
||
global _sam2_predictor
|
||
if _sam2_predictor is not None:
|
||
return _sam2_predictor
|
||
with _sam2_predictor_lock:
|
||
if _sam2_predictor is not None:
|
||
return _sam2_predictor
|
||
try:
|
||
from sam2.build_sam import build_sam2
|
||
from sam2.sam2_image_predictor import SAM2ImagePredictor
|
||
with open(CONFIG_PATH) as f:
|
||
conf = json.load(f)
|
||
ckpt = os.path.expanduser(conf.get("sam2_checkpoint", "~/.sam/sam2.1_hiera_base_plus.pt"))
|
||
cfg = conf.get("sam2_config", "configs/sam2.1/sam2.1_hiera_t.yaml")
|
||
if not os.path.exists(ckpt):
|
||
raise FileNotFoundError(f"SAM2 checkpoint not found: {ckpt}")
|
||
model = build_sam2(cfg, ckpt, device="cuda")
|
||
_sam2_predictor = SAM2ImagePredictor(model)
|
||
print(f"[sam2] loaded from {ckpt}")
|
||
except Exception as e:
|
||
print(f"[sam2] not available: {e}")
|
||
_sam2_predictor = False
|
||
return _sam2_predictor
|
||
|
||
|
||
def _clean_mask(mask):
|
||
"""
|
||
Remove small isolated components and thin border artifacts from a boolean mask.
|
||
`mask` should be a 2D numpy boolean array (H, W).
|
||
"""
|
||
try:
|
||
import cv2
|
||
import numpy as np
|
||
except ImportError:
|
||
return mask
|
||
|
||
h, w = mask.shape
|
||
mask_uint8 = mask.astype(np.uint8) * 255
|
||
|
||
# Connectivity 8 to keep diagonal hair strands connected
|
||
num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(mask_uint8, connectivity=8)
|
||
|
||
if num_labels <= 1:
|
||
return mask
|
||
|
||
# label 0 is background
|
||
areas = stats[1:, cv2.CC_STAT_AREA]
|
||
if len(areas) == 0:
|
||
return mask
|
||
|
||
largest_label = np.argmax(areas) + 1
|
||
largest_area = areas[largest_label - 1]
|
||
|
||
new_mask_uint8 = np.zeros_like(mask_uint8)
|
||
|
||
# Always keep the largest component
|
||
new_mask_uint8[labels == largest_label] = 255
|
||
|
||
for i in range(1, num_labels):
|
||
if i == largest_label:
|
||
continue
|
||
|
||
x, y, width, height, area = stats[i]
|
||
|
||
# Artifact detection heuristics:
|
||
# 1. Thin border lines (common artifacts at extreme edges)
|
||
is_border_artifact = False
|
||
if (x == 0 or x + width == w) and width <= 2:
|
||
is_border_artifact = True
|
||
if (y == 0 or y + height == h) and height <= 2:
|
||
is_border_artifact = True
|
||
|
||
# 2. Very small isolated noise
|
||
is_small_noise = area < (largest_area * 0.001) or area < 15
|
||
|
||
if is_border_artifact or is_small_noise:
|
||
continue
|
||
|
||
# Keep significant secondary components
|
||
new_mask_uint8[labels == i] = 255
|
||
|
||
return new_mask_uint8 > 0
|
||
|
||
|
||
def _person_mask_score(mask, h: int, w: int):
|
||
"""Rate how much `mask` looks like a centered subject vs. the background.
|
||
|
||
The subject (person) sits in the middle of the frame and rarely fills the
|
||
corners; the background is the opposite — it hugs the corners and is sparse
|
||
in the center. So `center_cov - corner_cov` is strongly positive for a
|
||
correct person mask and negative when SAM2 has selected the background
|
||
instead (the inverted-mask failure mode).
|
||
|
||
Returns (score, center_cov, corner_cov), all floats in [-1, 1] / [0, 1].
|
||
"""
|
||
import numpy as np
|
||
cy0, cy1 = int(h * 0.30), int(h * 0.70)
|
||
cx0, cx1 = int(w * 0.30), int(w * 0.70)
|
||
center_cov = float(mask[cy0:cy1, cx0:cx1].mean())
|
||
cs = max(1, int(min(h, w) * 0.10))
|
||
corner_cov = float(np.mean([
|
||
mask[:cs, :cs].mean(), mask[:cs, -cs:].mean(),
|
||
mask[-cs:, :cs].mean(), mask[-cs:, -cs:].mean(),
|
||
]))
|
||
return center_cov - corner_cov, center_cov, corner_cov
|
||
|
||
|
||
def _apply_transparency_sam2(png_bytes: bytes) -> bytes:
|
||
"""Remove background with SAM2 bbox segmentation; fallback to rembg.
|
||
|
||
Uses a near-full-frame bbox so SAM2 finds the largest foreground object
|
||
(the person) regardless of rotation or pose. This works well because
|
||
"transparent background" is stripped from the Qwen prompt upstream, so the
|
||
model renders a solid real background — giving SAM2 clear contrast to work
|
||
with. Point prompts were tried but produced holes in ¾-rotated poses
|
||
because the spine-column seeds land on background when the body is offset.
|
||
"""
|
||
predictor = _load_sam2()
|
||
if predictor is False:
|
||
return _apply_transparency(png_bytes)
|
||
try:
|
||
import numpy as np
|
||
import torch
|
||
img = Image.open(io.BytesIO(png_bytes)).convert("RGB")
|
||
arr = np.array(img)
|
||
h, w = arr.shape[:2]
|
||
|
||
# Near-full-frame bbox — 1 % margin so hair / shoes are inside the hint.
|
||
# SAM2 treats this as "find the prominent object within this region".
|
||
box = np.array([[int(w * 0.01), int(h * 0.01),
|
||
int(w * 0.99), int(h * 0.99)]], dtype=np.float32)
|
||
|
||
with embeddings._gpu_lock:
|
||
with torch.inference_mode():
|
||
predictor.set_image(arr)
|
||
masks, scores, _ = predictor.predict(
|
||
box=box,
|
||
multimask_output=True,
|
||
)
|
||
|
||
if masks is None or len(masks) == 0:
|
||
print("[sam2] no masks returned, falling back to rembg")
|
||
return _apply_transparency(png_bytes)
|
||
|
||
# Pick the candidate that best matches a centered subject rather than
|
||
# blindly trusting argmax(scores): on busy or low-contrast backgrounds
|
||
# the top-confidence SAM2 mask is sometimes the background itself.
|
||
# Combine the centered-subject prior with SAM2's own confidence.
|
||
best = None
|
||
best_rank = -1e9
|
||
for i in range(len(masks)):
|
||
m = masks[i].astype(bool)
|
||
psc, _, _ = _person_mask_score(m, h, w)
|
||
rank = psc + 0.10 * float(scores[i])
|
||
if rank > best_rank:
|
||
best_rank, best = rank, m
|
||
|
||
# Inversion guard — the user's hint: the model is in the center. If the
|
||
# chosen mask still covers the corners more than the center, SAM2 picked
|
||
# the background; flip the alpha so the person stays opaque.
|
||
psc, ccov, kcov = _person_mask_score(best, h, w)
|
||
if psc < 0:
|
||
print(f"[sam2] mask inverted (center {ccov:.0%} < corners {kcov:.0%}) — flipping alpha")
|
||
best = ~best
|
||
|
||
# Remove small isolated artifacts and border lines
|
||
best = _clean_mask(best)
|
||
|
||
# Sanity check: a person should cover 5 %–92 % of the frame
|
||
coverage = float(best.sum()) / (h * w)
|
||
if coverage < 0.05 or coverage > 0.92:
|
||
print(f"[sam2] mask coverage {coverage:.1%} out of range, falling back to rembg")
|
||
return _apply_transparency(png_bytes)
|
||
|
||
mask_np = best.astype(np.uint8) * 255
|
||
|
||
# Soft anti-aliased edge (radius 1 keeps accessory detail)
|
||
try:
|
||
from PIL import ImageFilter
|
||
alpha_img = Image.fromarray(mask_np, mode="L")
|
||
alpha_img = alpha_img.filter(ImageFilter.GaussianBlur(radius=1))
|
||
except Exception:
|
||
alpha_img = Image.fromarray(mask_np, mode="L")
|
||
|
||
rgba = img.convert("RGBA")
|
||
r, g, b, _ = rgba.split()
|
||
out = Image.merge("RGBA", (r, g, b, alpha_img))
|
||
buf = io.BytesIO()
|
||
out.save(buf, format="PNG")
|
||
print(f"[sam2] mask OK ({coverage:.1%} coverage)")
|
||
return buf.getvalue()
|
||
except Exception as e:
|
||
print(f"[sam2] inference error, falling back to rembg: {e}")
|
||
return _apply_transparency(png_bytes)
|
||
|
||
|
||
def _apply_transparency_black_bg(png_bytes: bytes) -> bytes:
|
||
"""Background removal for black-background Qwen output (bg_removal=sam2 mode).
|
||
|
||
Strategy:
|
||
1. Threshold: any pixel with max-channel > 25 is person (non-black).
|
||
This correctly identifies the subject regardless of pose or rotation.
|
||
2. Derive a tight person bounding-box from the threshold mask.
|
||
3. Run SAM2 with that box for sub-pixel edge refinement.
|
||
Accept SAM2 result only when its coverage is close (±30 pp) to the
|
||
threshold estimate — this rejects the inverted-mask failure mode where
|
||
SAM2 picks the large dark region as the "object".
|
||
4. Fall back to the threshold mask (Gaussian-blurred edges) if SAM2
|
||
is unavailable, errors, or diverges.
|
||
|
||
Do NOT use the full-frame bbox here: on black-background images the large
|
||
dark region scores higher than the person, causing SAM2 to invert the mask.
|
||
"""
|
||
import numpy as np
|
||
import torch
|
||
from PIL import ImageFilter
|
||
|
||
img = Image.open(io.BytesIO(png_bytes)).convert("RGB")
|
||
arr = np.array(img)
|
||
h, w = arr.shape[:2]
|
||
|
||
# Step 1 — threshold: non-black pixels are the person
|
||
is_person = np.max(arr, axis=2) > 25
|
||
thresh_cov = float(is_person.sum()) / (h * w)
|
||
print(f"[bg-black] threshold coverage: {thresh_cov:.1%}")
|
||
|
||
if not is_person.any():
|
||
print("[bg-black] all-black image — falling back to rembg")
|
||
return _apply_transparency(png_bytes)
|
||
|
||
# Step 2 — tight bounding box from threshold
|
||
rows = np.any(is_person, axis=1)
|
||
cols = np.any(is_person, axis=0)
|
||
rmin = int(np.where(rows)[0][0]); rmax = int(np.where(rows)[0][-1])
|
||
cmin = int(np.where(cols)[0][0]); cmax = int(np.where(cols)[0][-1])
|
||
margin = int(min(h, w) * 0.02)
|
||
x1 = max(0, cmin - margin); y1 = max(0, rmin - margin)
|
||
x2 = min(w, cmax + margin); y2 = min(h, rmax + margin)
|
||
|
||
# Step 3 — SAM2 with the person-specific bbox
|
||
predictor = _load_sam2()
|
||
if predictor is not False:
|
||
box = np.array([[x1, y1, x2, y2]], dtype=np.float32)
|
||
try:
|
||
with embeddings._gpu_lock:
|
||
with torch.inference_mode():
|
||
predictor.set_image(arr)
|
||
masks, scores, _ = predictor.predict(box=box, multimask_output=True)
|
||
|
||
if masks is not None and len(masks) > 0:
|
||
best = masks[int(np.argmax(scores))]
|
||
sam_cov = float(best.sum()) / (h * w)
|
||
print(f"[bg-black] SAM2 coverage: {sam_cov:.1%}")
|
||
|
||
if 0.03 < sam_cov < 0.95 and abs(sam_cov - thresh_cov) < 0.30:
|
||
best = _clean_mask(best)
|
||
mask_np = best.astype(np.uint8) * 255
|
||
alpha_img = Image.fromarray(mask_np, "L").filter(ImageFilter.GaussianBlur(radius=1))
|
||
rgba = img.convert("RGBA"); r, g, b, _ = rgba.split()
|
||
out = Image.merge("RGBA", (r, g, b, alpha_img))
|
||
buf = io.BytesIO(); out.save(buf, "PNG")
|
||
print(f"[bg-black] SAM2 accepted ✓")
|
||
return buf.getvalue()
|
||
else:
|
||
print(f"[bg-black] SAM2 diverged ({sam_cov:.1%} vs {thresh_cov:.1%}) — threshold fallback")
|
||
except Exception as e:
|
||
print(f"[bg-black] SAM2 error: {e} — threshold fallback")
|
||
|
||
# Step 4 — fallback: threshold mask with soft edge blur
|
||
print("[bg-black] using threshold mask")
|
||
is_person = _clean_mask(is_person)
|
||
mask_np = is_person.astype(np.uint8) * 255
|
||
alpha_img = Image.fromarray(mask_np, "L").filter(ImageFilter.GaussianBlur(radius=2))
|
||
rgba = img.convert("RGBA"); r, g, b, _ = rgba.split()
|
||
out = Image.merge("RGBA", (r, g, b, alpha_img))
|
||
buf = io.BytesIO(); out.save(buf, "PNG")
|
||
return buf.getvalue()
|
||
|
||
|
||
@app.post("/remove-background-sam/{filename:path}")
|
||
def remove_background_sam(filename: str):
|
||
"""SAM2-based background removal.
|
||
|
||
Writes the transparent result as a sidecar <stem>.nobg.png alongside the
|
||
original, which is left untouched. Returns the sidecar URL so the UI can
|
||
switch the viewer without touching the source file.
|
||
Falls back to rembg when SAM2 is unavailable.
|
||
"""
|
||
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]
|
||
output_dir = os.path.dirname(path)
|
||
|
||
dir_part = os.path.dirname(filename)
|
||
basename = os.path.basename(filename)
|
||
stem = os.path.splitext(basename)[0]
|
||
nobg_basename = f"{stem}.nobg.png"
|
||
if dir_part:
|
||
nobg_filename = f"{dir_part}/{nobg_basename}"
|
||
else:
|
||
nobg_filename = nobg_basename
|
||
nobg_path = os.path.join(output_dir, nobg_basename)
|
||
|
||
with open(path, "rb") as f:
|
||
png_bytes = f.read()
|
||
transparent_png = _apply_transparency_sam2(png_bytes)
|
||
with open(nobg_path, "wb") as f:
|
||
f.write(transparent_png)
|
||
|
||
# Register sidecar in DB so it appears in the same group
|
||
group_id = person[1] if person and person[1] else naming.get_base_name(os.path.basename(filename))
|
||
tags_list = None
|
||
if person[2]:
|
||
try:
|
||
tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2]
|
||
except Exception:
|
||
tags_list = None
|
||
database.upsert_person(
|
||
nobg_filename, filepath=nobg_path, group_id=group_id,
|
||
name=person[0], tags=tags_list, embedding=person[3],
|
||
clip_description=person[4], prompt=person[6], pose=person[7],
|
||
group_name=person[9], hidden=person[10],
|
||
has_background=False,
|
||
has_clothing=person[13],
|
||
is_source=person[14],
|
||
pose_description=person[15],
|
||
pose_skeleton=person[16],
|
||
source_refs=json.dumps([filename]), # original is the reference
|
||
)
|
||
_update_cached_file_meta(nobg_filename, exists=True)
|
||
_invalidate_static()
|
||
used_sam2 = _sam2_predictor is not False and _sam2_predictor is not None
|
||
return {
|
||
"status": "success",
|
||
"filename": filename,
|
||
"nobg_filename": nobg_filename,
|
||
"nobg_url": f"/output/{nobg_filename}",
|
||
"used_sam2": used_sam2,
|
||
}
|
||
|
||
|
||
@app.post("/images/{filename:path}/autocrop")
|
||
def autocrop_image(filename: str):
|
||
"""Crop away transparent borders from an image in-place."""
|
||
import numpy as np
|
||
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]
|
||
img = Image.open(path).convert("RGBA")
|
||
arr = np.array(img)
|
||
alpha = arr[:, :, 3]
|
||
rows = np.any(alpha > 0, axis=1)
|
||
cols = np.any(alpha > 0, axis=0)
|
||
if not rows.any():
|
||
raise HTTPException(400, "Image is fully transparent")
|
||
rmin, rmax = np.where(rows)[0][[0, -1]]
|
||
cmin, cmax = np.where(cols)[0][[0, -1]]
|
||
cropped = img.crop((cmin, rmin, cmax + 1, rmax + 1))
|
||
cropped.save(path, format="PNG")
|
||
_update_cached_file_meta(filename, exists=True)
|
||
_invalidate_static()
|
||
_metadata_executor.submit(_process_image_for_metadata, filename)
|
||
return {"status": "success", "filename": filename, "box": [int(cmin), int(rmin), int(cmax+1), int(rmax+1)]}
|
||
|
||
|
||
class CropRequest(BaseModel):
|
||
x1: int
|
||
y1: int
|
||
x2: int
|
||
y2: int
|
||
as_copy: bool = False # True → crop a fresh copy, leaving the original untouched
|
||
|
||
|
||
@app.post("/images/{filename:path}/crop")
|
||
def manual_crop_image(filename: str, req: CropRequest):
|
||
"""Crop the image to the given pixel rectangle (in original image coordinates).
|
||
|
||
By default the crop is applied in-place. When ``as_copy`` is set, a new copy is
|
||
created first (referencing the original via ``source_refs``) and the crop is applied
|
||
to that copy, so the original is preserved.
|
||
"""
|
||
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")
|
||
src_path = person[5]
|
||
|
||
if req.as_copy:
|
||
# Mirror duplicate_image: copy file + register a DB row that points back to the original.
|
||
from datetime import datetime as _dt
|
||
if filename.startswith("_turntable/"):
|
||
dir_part = ""
|
||
output_dir = _load_output_dir()
|
||
else:
|
||
output_dir = os.path.dirname(src_path)
|
||
dir_part = os.path.dirname(filename)
|
||
basename = os.path.basename(filename)
|
||
stem, ext = os.path.splitext(basename)
|
||
if not ext:
|
||
ext = ".png"
|
||
ts = _dt.now().strftime("%Y%m%d_%H%M%S")
|
||
new_basename = f"{ts}_crop_{stem}{ext}"
|
||
if dir_part:
|
||
new_filename = f"{dir_part}/{new_basename}"
|
||
else:
|
||
new_filename = new_basename
|
||
path = os.path.join(output_dir, new_basename)
|
||
shutil.copy2(src_path, path)
|
||
group_id = person[1] if person and person[1] else naming.get_base_name(os.path.basename(filename))
|
||
tags_list = None
|
||
if person[2]:
|
||
try:
|
||
tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2]
|
||
except Exception:
|
||
tags_list = None
|
||
database.upsert_person(
|
||
new_filename, filepath=path, group_id=group_id,
|
||
name=person[0], tags=tags_list, embedding=person[3],
|
||
clip_description=person[4], prompt=person[6], pose=person[7],
|
||
group_name=person[9], hidden=person[10],
|
||
has_background=person[11],
|
||
has_clothing=person[13],
|
||
is_source=person[14],
|
||
pose_description=person[15],
|
||
pose_skeleton=person[16],
|
||
source_refs=json.dumps([filename]), # original is the reference
|
||
)
|
||
else:
|
||
new_filename = filename
|
||
path = src_path
|
||
|
||
img = Image.open(path)
|
||
w, h = img.size
|
||
x1 = max(0, min(req.x1, w))
|
||
y1 = max(0, min(req.y1, h))
|
||
x2 = max(0, min(req.x2, w))
|
||
y2 = max(0, min(req.y2, h))
|
||
if x2 <= x1 or y2 <= y1:
|
||
if req.as_copy:
|
||
# Roll back the copy we just made so a bad rect doesn't leave an orphan.
|
||
try:
|
||
database.delete_person(new_filename)
|
||
os.remove(path)
|
||
except Exception:
|
||
pass
|
||
raise HTTPException(400, "Invalid crop rectangle")
|
||
cropped = img.crop((x1, y1, x2, y2))
|
||
fmt = "PNG" if path.lower().endswith(".png") else "JPEG"
|
||
cropped.save(path, format=fmt)
|
||
_update_cached_file_meta(new_filename, exists=True)
|
||
_invalidate_static()
|
||
_metadata_executor.submit(_process_image_for_metadata, new_filename)
|
||
return {"status": "success", "filename": filename, "new_filename": new_filename,
|
||
"new_url": f"/output/{new_filename}", "as_copy": req.as_copy,
|
||
"box": [x1, y1, x2, y2]}
|
||
|
||
|
||
class PadRequest(BaseModel):
|
||
top: int | float | str = 0
|
||
right: int | float | str = 0
|
||
bottom: int | float | str = 0
|
||
left: int | float | str = 0
|
||
as_copy: bool = True
|
||
fill: str = "transparent" # "black", "white", "transparent"
|
||
outpaint: bool = False
|
||
prompt: str | None = None
|
||
|
||
|
||
def _apply_manual_pad(pil: Image.Image, top, right, bottom, left,
|
||
fill: str = "transparent") -> Image.Image:
|
||
"""Expand canvas by padding pixels on each side. Fill with black, white, or transparency.
|
||
Supports pixel values (int) or percentages (str like "10%" or float < 1.0).
|
||
"""
|
||
w, h = pil.size
|
||
|
||
def resolve(val, total):
|
||
if not val:
|
||
return 0
|
||
if isinstance(val, str):
|
||
if "%" in val:
|
||
try:
|
||
return int(float(val.replace("%", "")) * total / 100)
|
||
except:
|
||
return 0
|
||
if "px" in val:
|
||
try:
|
||
return int(float(val.replace("px", "")))
|
||
except:
|
||
return 0
|
||
try:
|
||
f = float(val)
|
||
if 0 < f < 1.0:
|
||
return int(f * total)
|
||
return int(f)
|
||
except:
|
||
return 0
|
||
|
||
top = max(0, resolve(top, h))
|
||
right = max(0, resolve(right, w))
|
||
bottom = max(0, resolve(bottom, h))
|
||
left = max(0, resolve(left, w))
|
||
|
||
if not any([top, right, bottom, left]):
|
||
return pil
|
||
new_w = w + left + right
|
||
new_h = h + top + bottom
|
||
if fill == "transparent":
|
||
canvas = Image.new("RGBA", (new_w, new_h), (0, 0, 0, 0))
|
||
pil = pil.convert("RGBA")
|
||
elif fill == "white":
|
||
canvas = Image.new("RGB", (new_w, new_h), (255, 255, 255))
|
||
pil = pil.convert("RGB")
|
||
else:
|
||
canvas = Image.new("RGB", (new_w, new_h), (0, 0, 0))
|
||
pil = pil.convert("RGB")
|
||
canvas.paste(pil, (left, top))
|
||
return canvas
|
||
|
||
|
||
@app.post("/images/{filename:path}/pad")
|
||
def pad_image(filename: str, req: PadRequest):
|
||
"""Expand the image canvas by adding blank padding on each side.
|
||
|
||
as_copy=True (default) creates a new image referencing the original;
|
||
as_copy=False modifies in-place.
|
||
"""
|
||
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")
|
||
src_path = person[5]
|
||
|
||
if req.as_copy:
|
||
from datetime import datetime as _dt
|
||
if filename.startswith("_turntable/"):
|
||
dir_part = ""
|
||
output_dir = _load_output_dir()
|
||
else:
|
||
output_dir = os.path.dirname(src_path)
|
||
dir_part = os.path.dirname(filename)
|
||
basename = os.path.basename(filename)
|
||
stem, ext = os.path.splitext(basename)
|
||
if not ext:
|
||
ext = ".png"
|
||
ts = _dt.now().strftime("%Y%m%d_%H%M%S")
|
||
new_basename = f"{ts}_pad_{stem}{ext}"
|
||
if dir_part:
|
||
new_filename = f"{dir_part}/{new_basename}"
|
||
else:
|
||
new_filename = new_basename
|
||
path = os.path.join(output_dir, new_basename)
|
||
shutil.copy2(src_path, path)
|
||
group_id = person[1] if person and person[1] else naming.get_base_name(os.path.basename(filename))
|
||
tags_list = None
|
||
if person[2]:
|
||
try:
|
||
tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2]
|
||
except Exception:
|
||
tags_list = None
|
||
database.upsert_person(
|
||
new_filename, filepath=path, group_id=group_id,
|
||
name=person[0], tags=tags_list, embedding=person[3],
|
||
clip_description=person[4], prompt=person[6], pose=person[7],
|
||
group_name=person[9], hidden=person[10],
|
||
has_background=person[11],
|
||
has_clothing=person[13],
|
||
is_source=person[14],
|
||
pose_description=person[15],
|
||
pose_skeleton=person[16],
|
||
source_refs=json.dumps([filename]),
|
||
)
|
||
else:
|
||
new_filename = filename
|
||
path = src_path
|
||
|
||
img = Image.open(path)
|
||
padded = _apply_manual_pad(img, req.top, req.right, req.bottom, req.left, req.fill)
|
||
if padded is img:
|
||
raise HTTPException(400, "No padding specified (all sides are 0)")
|
||
|
||
if req.outpaint:
|
||
# Trigger Qwen outpainting
|
||
# If the image is transparent, composite it onto a neutral background (black)
|
||
# so the AI can see the area to be filled.
|
||
qwen_input = padded
|
||
fill_desc = ""
|
||
if qwen_input.mode == "RGBA":
|
||
# Composite onto black for maximal contrast in outpainting
|
||
bg = Image.new("RGB", qwen_input.size, (0, 0, 0))
|
||
bg.paste(qwen_input, mask=qwen_input.split()[3])
|
||
qwen_input = bg
|
||
fill_desc = " replacing the black background areas"
|
||
elif req.fill in ["black", "white"]:
|
||
fill_desc = f" replacing the {req.fill} background areas"
|
||
|
||
out_instr = f"Naturally outpaint and extend the borders of the image to complete the scene{fill_desc}."
|
||
actual_prompt = req.prompt or out_instr
|
||
if req.prompt and out_instr.lower() not in actual_prompt.lower():
|
||
actual_prompt = f"{actual_prompt}. {out_instr}"
|
||
elif not req.prompt and person[6]: # original prompt
|
||
actual_prompt = f"{person[6]}. {out_instr}"
|
||
|
||
try:
|
||
png_bytes = _run_pipeline(qwen_input, actual_prompt)
|
||
padded = Image.open(io.BytesIO(png_bytes))
|
||
# Register the outpaint prompt and set has_background=True in DB
|
||
database.upsert_person(new_filename, prompt=actual_prompt, has_background=True)
|
||
except Exception as e:
|
||
print(f"Outpaint error: {e}")
|
||
raise HTTPException(500, f"Outpaint failed: {e}")
|
||
|
||
fmt = "PNG" if path.lower().endswith(".png") else "JPEG"
|
||
if req.fill == "transparent":
|
||
fmt = "PNG" # JPEG cannot store alpha
|
||
padded.save(path, format=fmt)
|
||
_update_cached_file_meta(new_filename, exists=True)
|
||
_invalidate_static()
|
||
_metadata_executor.submit(_process_image_for_metadata, new_filename)
|
||
return {
|
||
"status": "success", "filename": filename, "new_filename": new_filename,
|
||
"new_url": f"/output/{new_filename}", "as_copy": req.as_copy,
|
||
"size": list(padded.size),
|
||
}
|
||
|
||
|
||
class RotateRequest(BaseModel):
|
||
degrees: int = 90 # clockwise rotation; must be a multiple of 90
|
||
|
||
|
||
@app.post("/images/{filename:path}/rotate")
|
||
def rotate_image(filename: str, req: RotateRequest):
|
||
"""Rotate an image clockwise in 90° steps, in place (lossless transpose)."""
|
||
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]
|
||
deg = req.degrees % 360
|
||
if deg not in (0, 90, 180, 270):
|
||
raise HTTPException(400, "degrees must be a multiple of 90")
|
||
if deg:
|
||
# PIL transpose is defined counter-clockwise; map clockwise degrees onto it.
|
||
cw_to_transpose = {
|
||
90: Image.Transpose.ROTATE_270,
|
||
180: Image.Transpose.ROTATE_180,
|
||
270: Image.Transpose.ROTATE_90,
|
||
}
|
||
img = Image.open(path).transpose(cw_to_transpose[deg])
|
||
fmt = "PNG" if path.lower().endswith(".png") else "JPEG"
|
||
img.save(path, format=fmt)
|
||
_update_cached_file_meta(filename, exists=True)
|
||
_invalidate_static()
|
||
return {"status": "success", "filename": filename, "degrees": deg}
|
||
|
||
|
||
@app.post("/images/{filename:path}/duplicate")
|
||
def duplicate_image(filename: str):
|
||
"""Copy an image into the same group with a fresh timestamp-based filename."""
|
||
import shutil as _shutil
|
||
from datetime import datetime as _dt
|
||
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]
|
||
if filename.startswith("_turntable/"):
|
||
dir_part = ""
|
||
output_dir = _load_output_dir()
|
||
else:
|
||
output_dir = os.path.dirname(path)
|
||
dir_part = os.path.dirname(filename)
|
||
basename = os.path.basename(filename)
|
||
stem, ext = os.path.splitext(basename)
|
||
if not ext:
|
||
ext = ".png"
|
||
ts = _dt.now().strftime("%Y%m%d_%H%M%S")
|
||
new_basename = f"{ts}_dup_{stem}{ext}"
|
||
if dir_part:
|
||
new_filename = f"{dir_part}/{new_basename}"
|
||
else:
|
||
new_filename = new_basename
|
||
|
||
new_path = os.path.join(output_dir, new_basename)
|
||
_shutil.copy2(path, new_path)
|
||
group_id = person[1] if person and person[1] else naming.get_base_name(os.path.basename(filename))
|
||
|
||
# Parse tags list if present
|
||
tags_list = None
|
||
if person[2]:
|
||
try:
|
||
tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2]
|
||
except Exception:
|
||
tags_list = None
|
||
|
||
database.upsert_person(
|
||
new_filename, filepath=new_path, group_id=group_id,
|
||
name=person[0], tags=tags_list, embedding=person[3],
|
||
clip_description=person[4], prompt=person[6], pose=person[7],
|
||
group_name=person[9], hidden=person[10],
|
||
has_background=person[11],
|
||
has_clothing=person[13],
|
||
is_source=person[14],
|
||
pose_description=person[15],
|
||
pose_skeleton=person[16],
|
||
source_refs=json.dumps([filename]), # original is the reference
|
||
)
|
||
_update_cached_file_meta(new_filename, exists=True)
|
||
_invalidate_static()
|
||
return {"status": "success", "new_filename": new_filename, "new_url": f"/output/{new_filename}"}
|
||
|
||
|
||
@app.post("/restore-background/{filename:path}")
|
||
def restore_background(filename: str):
|
||
"""Flatten RGBA → RGB (white composite), making the image opaque again."""
|
||
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]
|
||
img = Image.open(path)
|
||
if img.mode == "RGBA":
|
||
bg = Image.new("RGB", img.size, (255, 255, 255))
|
||
bg.paste(img, mask=img.split()[3])
|
||
buf = io.BytesIO()
|
||
bg.save(buf, format="PNG")
|
||
with open(path, "wb") as f:
|
||
f.write(buf.getvalue())
|
||
_update_cached_file_meta(filename, exists=True)
|
||
_invalidate_static()
|
||
return {"status": "success", "filename": filename}
|
||
|
||
|
||
@app.get("/sam2/check")
|
||
def sam2_check():
|
||
"""Return whether SAM2 is available."""
|
||
predictor = _load_sam2()
|
||
return {"sam2": predictor is not False and predictor is not None}
|
||
|
||
|
||
# --- 2D body-pose preview -----------------------------------------------------
|
||
# Estimates COCO-17 keypoints from the model image so the UI can overlay a
|
||
# posenet-style skeleton. Estimator is feature-detected: rtmlib (ONNX, reuses the
|
||
# already-installed onnxruntime) is preferred, mediapipe is a fallback. If neither
|
||
# is installed the endpoints report unavailable instead of erroring the request.
|
||
|
||
_pose_estimator = None # cached (callable, backend_name) or False if unavailable
|
||
_pose_lock = threading.Lock()
|
||
|
||
# COCO-17 keypoint names (the order rtmlib's Body model returns).
|
||
POSE_KEYPOINT_NAMES = [
|
||
"nose", "left_eye", "right_eye", "left_ear", "right_ear",
|
||
"left_shoulder", "right_shoulder", "left_elbow", "right_elbow",
|
||
"left_wrist", "right_wrist", "left_hip", "right_hip",
|
||
"left_knee", "right_knee", "left_ankle", "right_ankle",
|
||
]
|
||
# Bone connections (index pairs into COCO-17) for drawing the skeleton.
|
||
POSE_SKELETON = [
|
||
(5, 7), (7, 9), (6, 8), (8, 10), # arms
|
||
(11, 13), (13, 15), (12, 14), (14, 16), # legs
|
||
(5, 6), (11, 12), (5, 11), (6, 12), # torso
|
||
(0, 1), (0, 2), (1, 3), (2, 4), (0, 5), (0, 6), # head/neck
|
||
]
|
||
# mediapipe Pose (33 landmarks) → COCO-17 index map.
|
||
_MP_TO_COCO = [0, 2, 5, 7, 8, 11, 12, 13, 14, 15, 16, 23, 24, 25, 26, 27, 28]
|
||
|
||
|
||
def _load_pose_estimator():
|
||
global _pose_estimator
|
||
if _pose_estimator is not None:
|
||
return _pose_estimator
|
||
with _pose_lock:
|
||
if _pose_estimator is not None:
|
||
return _pose_estimator
|
||
# Preferred: rtmlib (RTMPose, ONNX) — returns COCO-17 directly.
|
||
try:
|
||
from rtmlib import Body
|
||
import numpy as np
|
||
model = Body(mode="balanced", backend="onnxruntime", device="cpu")
|
||
|
||
def _infer_rtm(pil):
|
||
bgr = np.array(pil.convert("RGB"))[:, :, ::-1]
|
||
kpts, scores = model(bgr) # (N,17,2), (N,17)
|
||
people = []
|
||
for person_kpts, person_scores in zip(kpts, scores):
|
||
people.append([[float(x), float(y), float(s)]
|
||
for (x, y), s in zip(person_kpts, person_scores)])
|
||
return people
|
||
|
||
_pose_estimator = (_infer_rtm, "rtmlib")
|
||
print("[pose] using rtmlib (RTMPose)")
|
||
return _pose_estimator
|
||
except Exception as e:
|
||
print(f"[pose] rtmlib unavailable: {e}")
|
||
|
||
# Fallback: mediapipe Pose (single person, normalized landmarks).
|
||
try:
|
||
import mediapipe as mp
|
||
import numpy as np
|
||
mp_pose = mp.solutions.pose.Pose(static_image_mode=True, model_complexity=2)
|
||
|
||
def _infer_mp(pil):
|
||
rgb = np.array(pil.convert("RGB"))
|
||
h, w = rgb.shape[:2]
|
||
res = mp_pose.process(rgb)
|
||
if not res.pose_landmarks:
|
||
return []
|
||
lm = res.pose_landmarks.landmark
|
||
kpts = []
|
||
for mp_idx in _MP_TO_COCO:
|
||
p = lm[mp_idx]
|
||
kpts.append([float(p.x * w), float(p.y * h), float(p.visibility)])
|
||
return [kpts]
|
||
|
||
_pose_estimator = (_infer_mp, "mediapipe")
|
||
print("[pose] using mediapipe Pose")
|
||
return _pose_estimator
|
||
except Exception as e:
|
||
print(f"[pose] mediapipe unavailable: {e}")
|
||
|
||
_pose_estimator = False
|
||
return _pose_estimator
|
||
|
||
|
||
# --- pose similarity (descriptor + index) -------------------------------------
|
||
# Pose descriptors are normalized (translation + scale invariant) COCO-17 vectors,
|
||
# cached in <output>/_data/poses_index.json so we can rank library images by pose.
|
||
|
||
_POSE_MIN_SCORE = 0.3
|
||
# Left/right keypoint pairs for the mirror-invariant distance.
|
||
_POSE_MIRROR = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
|
||
_pose_index_status = {"running": False, "done": 0, "total": 0}
|
||
|
||
|
||
def _pose_descriptor(keypoints):
|
||
"""Normalize one person's COCO-17 keypoints into a translation/scale-invariant
|
||
descriptor: {"vec": [34 floats], "vis": [17 ints]}. Returns None if too sparse."""
|
||
vis = [1 if (kp[2] >= _POSE_MIN_SCORE) else 0 for kp in keypoints]
|
||
if sum(vis) < 6:
|
||
return None
|
||
|
||
def _mid(a, b):
|
||
if vis[a] and vis[b]:
|
||
return ((keypoints[a][0] + keypoints[b][0]) / 2.0,
|
||
(keypoints[a][1] + keypoints[b][1]) / 2.0)
|
||
return None
|
||
|
||
hip = _mid(11, 12)
|
||
sho = _mid(5, 6)
|
||
center = hip or sho
|
||
if center is None:
|
||
return None
|
||
# Scale by torso length; fall back to keypoint spread if torso isn't visible.
|
||
if hip and sho:
|
||
scale = ((hip[0] - sho[0]) ** 2 + (hip[1] - sho[1]) ** 2) ** 0.5
|
||
else:
|
||
xs = [keypoints[i][0] for i in range(17) if vis[i]]
|
||
ys = [keypoints[i][1] for i in range(17) if vis[i]]
|
||
scale = max(max(xs) - min(xs), max(ys) - min(ys))
|
||
if not scale or scale < 1e-3:
|
||
return None
|
||
|
||
vec = []
|
||
for i in range(17):
|
||
if vis[i]:
|
||
vec.append((keypoints[i][0] - center[0]) / scale)
|
||
vec.append((keypoints[i][1] - center[1]) / scale)
|
||
else:
|
||
vec.extend([0.0, 0.0])
|
||
return {"vec": vec, "vis": vis}
|
||
|
||
|
||
def _pose_distance(a, b):
|
||
"""Weighted L2 between two descriptors over jointly-visible joints, taking the
|
||
min of the direct and left-right-mirrored comparison. Lower = more similar."""
|
||
def _dist(av, avis, bv, bvis, mirror):
|
||
total, n = 0.0, 0
|
||
for i in range(17):
|
||
j = _POSE_MIRROR[i] if mirror else i
|
||
if not (avis[i] and bvis[j]):
|
||
continue
|
||
bx = bv[j * 2] * (-1 if mirror else 1) # flip x when mirrored
|
||
by = bv[j * 2 + 1]
|
||
dx = av[i * 2] - bx
|
||
dy = av[i * 2 + 1] - by
|
||
total += dx * dx + dy * dy
|
||
n += 1
|
||
return (total / n) ** 0.5 if n >= 4 else float("inf")
|
||
direct = _dist(a["vec"], a["vis"], b["vec"], b["vis"], False)
|
||
mirror = _dist(a["vec"], a["vis"], b["vec"], b["vis"], True)
|
||
return min(direct, mirror)
|
||
|
||
def _describe_pose(kpts):
|
||
"""Generate a highly detailed human-readable description of a COCO-17 pose,
|
||
including body orientation (standing, sitting, laying down), facing direction,
|
||
arms position, and legs posture.
|
||
"""
|
||
vis = [k[2] >= _POSE_MIN_SCORE for k in kpts]
|
||
if sum(vis) < 5: return "Indeterminate pose"
|
||
|
||
parts = []
|
||
|
||
# Extract some key positions
|
||
head_x = kpts[0][0] if vis[0] else ((kpts[5][0] + kpts[6][0])/2 if (vis[5] and vis[6]) else None)
|
||
head_y = kpts[0][1] if vis[0] else ((kpts[5][1] + kpts[6][1])/2 if (vis[5] and vis[6]) else None)
|
||
|
||
hip_x = (kpts[11][0] + kpts[12][0])/2 if (vis[11] and vis[12]) else (kpts[11][0] if vis[11] else (kpts[12][0] if vis[12] else None))
|
||
hip_y = (kpts[11][1] + kpts[12][1])/2 if (vis[11] and vis[12]) else (kpts[11][1] if vis[11] else (kpts[12][1] if vis[12] else None))
|
||
|
||
sh_y = (kpts[5][1] + kpts[6][1])/2 if (vis[5] and vis[6]) else (kpts[0][1] if vis[0] else None)
|
||
torso_h = abs(hip_y - sh_y) if (hip_y is not None and sh_y is not None) else 100.0
|
||
|
||
# 1. Posture (standing, sitting, kneeling/crouching, laying down/reclining)
|
||
posture = "upright"
|
||
if head_y is not None and hip_y is not None:
|
||
if head_y > hip_y + 30:
|
||
posture = "upside down"
|
||
else:
|
||
# Check horizontal vs vertical distance to identify lying down
|
||
dx = abs(head_x - hip_x) if head_x is not None and hip_x is not None else 0
|
||
dy = abs(head_y - hip_y)
|
||
if dx > 1.2 * dy:
|
||
posture = "lying down/reclining"
|
||
|
||
if posture == "upright":
|
||
has_hips = vis[11] or vis[12]
|
||
has_knees = vis[13] or vis[14]
|
||
has_ankles = vis[15] or vis[16]
|
||
if has_hips and has_knees:
|
||
h_y = hip_y
|
||
h_x = hip_x
|
||
k_y = (kpts[13][1] + kpts[14][1])/2 if (vis[13] and vis[14]) else (kpts[13][1] if vis[13] else kpts[14][1])
|
||
k_x = (kpts[13][0] + kpts[14][0])/2 if (vis[13] and vis[14]) else (kpts[13][0] if vis[13] else kpts[14][0])
|
||
|
||
thigh_dy = abs(k_y - h_y)
|
||
thigh_dx = abs(k_x - h_x)
|
||
|
||
if has_ankles:
|
||
a_y = (kpts[15][1] + kpts[16][1])/2 if (vis[15] and vis[16]) else (kpts[15][1] if vis[15] else kpts[16][1])
|
||
a_x = (kpts[15][0] + kpts[16][0])/2 if (vis[15] and vis[16]) else (kpts[15][0] if vis[15] else kpts[16][0])
|
||
shin_dy = abs(a_y - k_y)
|
||
shin_dx = abs(a_x - k_x)
|
||
|
||
# Sitting: thigh horizontal, shin vertical
|
||
if thigh_dy < 0.6 * thigh_dx and shin_dy > 1.2 * shin_dx:
|
||
posture = "sitting"
|
||
elif thigh_dy < 0.45 * torso_h and shin_dy > 0.5 * torso_h:
|
||
posture = "sitting"
|
||
# Crouching/Kneeling: hips close to ankles/ground
|
||
elif abs(h_y - a_y) < 0.85 * torso_h:
|
||
posture = "crouching/kneeling"
|
||
else:
|
||
posture = "standing"
|
||
else:
|
||
# Ankles not visible
|
||
if thigh_dy < 0.5 * torso_h:
|
||
posture = "sitting"
|
||
else:
|
||
posture = "standing"
|
||
|
||
parts.append(posture)
|
||
|
||
# 2. Body Orientation / Facing Direction
|
||
if vis[5] and vis[6]:
|
||
sh_dist = abs(kpts[5][0] - kpts[6][0])
|
||
# Profile view check (shoulders compressed horizontally)
|
||
if sh_dist < 0.25 * torso_h:
|
||
parts.append("turned sideways (profile view)")
|
||
# Back view check
|
||
elif kpts[5][0] < kpts[6][0]:
|
||
parts.append("facing away (back view)")
|
||
else:
|
||
parts.append("facing forward (front view)")
|
||
|
||
# 3. Arms Posture
|
||
if vis[9] and vis[10]: # wrists
|
||
sh_y_val = sh_y if sh_y is not None else kpts[0][1]
|
||
if kpts[9][1] < sh_y_val and kpts[10][1] < sh_y_val:
|
||
parts.append("arms raised")
|
||
elif kpts[9][1] > sh_y_val + torso_h * 0.8 and kpts[10][1] > sh_y_val + torso_h * 0.8:
|
||
# Check if close to hips (hands on hips)
|
||
if (vis[11] and abs(kpts[9][0] - kpts[11][0]) < torso_h * 0.2) or (vis[12] and abs(kpts[10][0] - kpts[12][0]) < torso_h * 0.2):
|
||
parts.append("hands on hips")
|
||
else:
|
||
parts.append("arms down")
|
||
else:
|
||
parts.append("arms at sides")
|
||
|
||
# 4. Legs Posture
|
||
if vis[15] and vis[16]: # ankles
|
||
ankle_dist = abs(kpts[15][0] - kpts[16][0])
|
||
if ankle_dist > torso_h * 0.6:
|
||
parts.append("legs spread")
|
||
else:
|
||
parts.append("legs together")
|
||
|
||
if not parts: return "Generic pose"
|
||
return ", ".join(parts)
|
||
|
||
|
||
def _best_person(people):
|
||
"""Pick the largest-bbox person from an estimator result (most prominent subject)."""
|
||
best, best_area = None, -1.0
|
||
for kpts in people:
|
||
xs = [k[0] for k in kpts if k[2] >= _POSE_MIN_SCORE]
|
||
ys = [k[1] for k in kpts if k[2] >= _POSE_MIN_SCORE]
|
||
if len(xs) < 2:
|
||
continue
|
||
area = (max(xs) - min(xs)) * (max(ys) - min(ys))
|
||
if area > best_area:
|
||
best, best_area = kpts, area
|
||
return best
|
||
|
||
|
||
def _pose_index_path():
|
||
return os.path.join(_load_output_dir(), "_data", "poses_index.json")
|
||
|
||
|
||
def _load_pose_index():
|
||
try:
|
||
with open(_pose_index_path(), "r") as f:
|
||
return json.load(f)
|
||
except Exception:
|
||
return {}
|
||
|
||
|
||
_pose_index_lock = threading.Lock()
|
||
|
||
|
||
def _save_pose_index_entry(filename, desc):
|
||
with _pose_index_lock:
|
||
idx = _load_pose_index()
|
||
idx[filename] = desc
|
||
os.makedirs(os.path.dirname(_pose_index_path()), exist_ok=True)
|
||
_write_json(_pose_index_path(), idx)
|
||
|
||
|
||
@app.get("/pose/check")
|
||
def pose_check():
|
||
"""Report whether a body-pose estimator is available (and which backend)."""
|
||
est = _load_pose_estimator()
|
||
if not est:
|
||
return {"available": False,
|
||
"hint": "pip install rtmlib onnxruntime (or: pip install mediapipe)"}
|
||
return {"available": True, "backend": est[1]}
|
||
|
||
|
||
@app.post("/images/{filename:path}/pose")
|
||
def estimate_pose(filename: str):
|
||
"""Estimate COCO-17 body keypoints for an image. Returns pixel-space keypoints
|
||
plus the skeleton edge list so the frontend can overlay a posenet-style preview."""
|
||
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")
|
||
est = _load_pose_estimator()
|
||
if not est:
|
||
raise HTTPException(501, "No pose estimator installed. Try: pip install rtmlib onnxruntime")
|
||
infer, backend = est
|
||
pil = Image.open(person[5]).convert("RGB")
|
||
try:
|
||
people = infer(pil)
|
||
except Exception as e:
|
||
raise HTTPException(500, f"Pose estimation failed: {e}")
|
||
# Cache the descriptor so "find similar pose" can rank this image later.
|
||
best = _best_person(people)
|
||
pose_desc = None
|
||
pose_skeleton_json = None
|
||
if best is not None:
|
||
pose_desc = _describe_pose(best)
|
||
pose_skeleton_json = json.dumps(best)
|
||
desc = _pose_descriptor(best)
|
||
if desc is not None:
|
||
try:
|
||
_save_pose_index_entry(filename, desc)
|
||
except Exception as e:
|
||
print(f"[pose] index save failed for {filename}: {e}")
|
||
# Save to DB
|
||
database.upsert_person(filename, pose_description=pose_desc, pose_skeleton=pose_skeleton_json)
|
||
_update_cached_file_meta(filename, exists=True)
|
||
_invalidate_static()
|
||
|
||
return {
|
||
"status": "success",
|
||
"backend": backend,
|
||
"width": pil.width,
|
||
"height": pil.height,
|
||
"names": POSE_KEYPOINT_NAMES,
|
||
"skeleton": POSE_SKELETON,
|
||
"people": people,
|
||
"pose_description": pose_desc,
|
||
"pose_skeleton": pose_skeleton_json,
|
||
}
|
||
|
||
|
||
def _build_pose_index_task():
|
||
try:
|
||
est = _load_pose_estimator()
|
||
if not est:
|
||
return
|
||
infer, _ = est
|
||
output_dir = _load_output_dir()
|
||
with _pose_index_lock:
|
||
idx = _load_pose_index()
|
||
persons = database.list_persons()
|
||
|
||
# p[17] is pose_description, p[18] is pose_skeleton
|
||
todo = [p for p in persons
|
||
if (p[0] not in idx or p[17] is None or p[18] is None)
|
||
and (p[12] or "image") != "video"
|
||
and p[0].lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
|
||
|
||
_pose_index_status.update(running=True, done=0, total=len(todo))
|
||
print(f"[pose] index build/backfill: {len(todo)} images to process")
|
||
dirty = 0
|
||
for p in todo:
|
||
fn = p[0]
|
||
try:
|
||
fpath = os.path.join(output_dir, fn)
|
||
if os.path.exists(fpath):
|
||
people = infer(Image.open(fpath).convert("RGB"))
|
||
best = _best_person(people)
|
||
|
||
pose_desc = None
|
||
pose_skel = None
|
||
if best is not None:
|
||
pose_desc = _describe_pose(best)
|
||
pose_skel = json.dumps(best)
|
||
desc = _pose_descriptor(best)
|
||
if desc is not None:
|
||
idx[fn] = desc
|
||
dirty += 1
|
||
|
||
# Update DB if missing
|
||
if p[17] is None or p[18] is None:
|
||
database.upsert_person(fn, pose_description=pose_desc, pose_skeleton=pose_skel)
|
||
except Exception as e:
|
||
print(f"[pose] index error for {fn}: {e}")
|
||
_pose_index_status["done"] += 1
|
||
# Batch-flush every 50 to avoid O(n^2) full-file rewrites.
|
||
if dirty >= 50:
|
||
with _pose_index_lock:
|
||
_write_json(_pose_index_path(), idx)
|
||
dirty = 0
|
||
print(f"[pose] index progress: {_pose_index_status['done']}/{len(todo)}")
|
||
with _pose_index_lock:
|
||
_write_json(_pose_index_path(), idx)
|
||
print(f"[pose] index build complete: {len(idx)} entries")
|
||
_invalidate_static()
|
||
except Exception as e:
|
||
print(f"[pose] index build failed: {e}")
|
||
finally:
|
||
_pose_index_status["running"] = False
|
||
|
||
|
||
@app.post("/pose/index")
|
||
def build_pose_index():
|
||
"""Compute pose descriptors for all library images lacking one (daemon thread)."""
|
||
if not _load_pose_estimator():
|
||
raise HTTPException(501, "No pose estimator installed. Try: pip install rtmlib onnxruntime")
|
||
if _pose_index_status.get("running"):
|
||
return {"status": "already_running", **_pose_index_status}
|
||
threading.Thread(target=_build_pose_index_task, daemon=True).start()
|
||
return {"status": "started"}
|
||
|
||
|
||
@app.get("/pose/index/status")
|
||
def pose_index_status():
|
||
idx = _load_pose_index()
|
||
return {**_pose_index_status, "indexed": len(idx)}
|
||
|
||
|
||
def _rank_similar_poses(query_desc, limit, exclude=None):
|
||
idx = _load_pose_index()
|
||
scored = []
|
||
for fn, desc in idx.items():
|
||
if fn == exclude or not desc or "vec" not in desc:
|
||
continue
|
||
d = _pose_distance(query_desc, desc)
|
||
if d != float("inf"):
|
||
scored.append((d, fn))
|
||
scored.sort(key=lambda x: x[0])
|
||
groups = get_groups() if scored else {}
|
||
return [{"filename": fn, "group_id": groups.get(fn), "distance": round(d, 4)}
|
||
for d, fn in scored[:limit]]
|
||
|
||
|
||
@app.get("/pose/similar/{filename:path}")
|
||
def similar_pose(filename: str, limit: int = 12):
|
||
"""Rank library images by pose similarity to the given image."""
|
||
idx = _load_pose_index()
|
||
query = idx.get(filename)
|
||
if query is None:
|
||
# Compute on demand (also caches it).
|
||
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")
|
||
est = _load_pose_estimator()
|
||
if not est:
|
||
raise HTTPException(501, "No pose estimator installed.")
|
||
best = _best_person(est[0](Image.open(person[5]).convert("RGB")))
|
||
query = _pose_descriptor(best) if best is not None else None
|
||
if query is None:
|
||
raise HTTPException(404, "No detectable pose in this image")
|
||
try:
|
||
_save_pose_index_entry(filename, query)
|
||
except Exception:
|
||
pass
|
||
return {"filename": filename, "similar": _rank_similar_poses(query, limit, exclude=filename)}
|
||
|
||
|
||
class PoseSimilarRequest(BaseModel):
|
||
keypoints: list[list[float]] # [[x,y,score], ...17] in image pixels
|
||
width: int = 0
|
||
height: int = 0
|
||
limit: int = 12
|
||
|
||
|
||
@app.post("/pose/similar")
|
||
def similar_pose_from_keypoints(req: PoseSimilarRequest):
|
||
"""Rank library images by similarity to a supplied (e.g. hand-edited) skeleton."""
|
||
query = _pose_descriptor(req.keypoints)
|
||
if query is None:
|
||
raise HTTPException(400, "Supplied pose is too sparse to match")
|
||
return {"similar": _rank_similar_poses(query, req.limit)}
|
||
|
||
|
||
# --- pose-guided generation ---------------------------------------------------
|
||
# Renders COCO-17 keypoints as an OpenPose-style skeleton image and passes it as
|
||
# image2 to Qwen so the VLM can match the body pose while preserving appearance.
|
||
|
||
# Color per POSE_SKELETON entry (18 bones in POSE_SKELETON order).
|
||
_OPENPOSE_BONE_COLORS = [
|
||
(0, 80, 255), # 5→7 L upper arm
|
||
(0, 130, 255), # 7→9 L lower arm
|
||
(255, 50, 0), # 6→8 R upper arm
|
||
(255, 110, 0), # 8→10 R lower arm
|
||
(255, 200, 0), # 11→13 L upper leg
|
||
(240, 240, 0), # 13→15 L lower leg
|
||
(100, 255, 0), # 12→14 R upper leg
|
||
( 50, 220, 0), # 14→16 R lower leg
|
||
(200, 0, 220), # 5→6 shoulders
|
||
(200, 0, 220), # 11→12 hips
|
||
(180, 60, 180), # 5→11 L torso side
|
||
( 60, 180, 180), # 6→12 R torso side
|
||
( 0, 230, 230), # 0→1 nose–L eye
|
||
( 0, 230, 230), # 0→2 nose–R eye
|
||
( 0, 180, 180), # 1→3 L eye–ear
|
||
( 0, 180, 180), # 2→4 R eye–ear
|
||
( 80, 80, 255), # 0→5 head–L shoulder
|
||
(255, 80, 80), # 0→6 head–R shoulder
|
||
]
|
||
|
||
# Per-joint hue (HSV-like, converted to RGB for PIL):
|
||
def _joint_color(idx: int) -> tuple[int, int, int]:
|
||
h = idx / 17.0 * 6.0
|
||
x = 1 - abs(h % 2 - 1)
|
||
if h < 1: r, g, b = 1, x, 0
|
||
elif h < 2: r, g, b = x, 1, 0
|
||
elif h < 3: r, g, b = 0, 1, x
|
||
elif h < 4: r, g, b = 0, x, 1
|
||
elif h < 5: r, g, b = x, 0, 1
|
||
else: r, g, b = 1, 0, x
|
||
return (int(r * 255), int(g * 255), int(b * 255))
|
||
|
||
|
||
def _render_openpose_image(keypoints: list, width: int, height: int) -> Image.Image:
|
||
"""Draw COCO-17 keypoints as a white stick-figure on a vivid magenta background.
|
||
|
||
Vivid background (not black/dark) prevents Qwen from treating the skeleton as
|
||
a real scene element and from blending the bone colors into the generated output.
|
||
White lines on magenta read unambiguously as a technical reference diagram.
|
||
"""
|
||
from PIL import ImageDraw
|
||
MIN_CONF = 0.25
|
||
cap_w = min(width, 896)
|
||
cap_h = min(height, 1152)
|
||
scale = min(cap_w / max(width, 1), cap_h / max(height, 1))
|
||
out_w, out_h = int(width * scale), int(height * scale)
|
||
|
||
img = Image.new("RGB", (out_w, out_h), color=(200, 0, 180)) # vivid magenta
|
||
draw = ImageDraw.Draw(img)
|
||
|
||
lw = max(4, out_w // 80) # line width scales with image
|
||
jr = max(5, out_w // 70) # joint radius
|
||
|
||
kpts = keypoints # [[x, y, conf], ...]
|
||
|
||
# bones — white so they stand out on the magenta background without
|
||
# adding any colors that Qwen might reproduce in the generated image
|
||
for a, b in POSE_SKELETON:
|
||
if a >= len(kpts) or b >= len(kpts):
|
||
continue
|
||
pa, pb = kpts[a], kpts[b]
|
||
if pa[2] < MIN_CONF or pb[2] < MIN_CONF:
|
||
continue
|
||
ax, ay = int(pa[0] * scale), int(pa[1] * scale)
|
||
bx, by = int(pb[0] * scale), int(pb[1] * scale)
|
||
draw.line([(ax, ay), (bx, by)], fill=(255, 255, 255), width=lw)
|
||
|
||
# joints — light yellow circles with dark outline for visibility
|
||
for kp in kpts:
|
||
if kp[2] < MIN_CONF:
|
||
continue
|
||
x, y = int(kp[0] * scale), int(kp[1] * scale)
|
||
draw.ellipse([(x - jr, y - jr), (x + jr, y + jr)],
|
||
fill=(255, 240, 100), outline=(60, 0, 60), width=max(1, lw // 2))
|
||
|
||
return img
|
||
|
||
|
||
class PoseFromWireframeRequest(BaseModel):
|
||
video: str # filename of a wireframe video in the output dir
|
||
time_s: float = 0.0 # timestamp in seconds to extract the frame from
|
||
|
||
|
||
@app.post("/pose/from-wireframe")
|
||
def pose_from_wireframe(req: PoseFromWireframeRequest):
|
||
"""Extract a frame from a wireframe video, run pose estimation, return COCO-17 keypoints.
|
||
|
||
This is the pose-discovery tool: scrub through a wireframe video to find a frame
|
||
that shows the desired pose, then load those keypoints into the pose overlay without
|
||
ever passing the wireframe image to Qwen.
|
||
"""
|
||
import tempfile
|
||
import subprocess as sp
|
||
|
||
output_dir = _load_output_dir()
|
||
video_path = os.path.join(output_dir, req.video)
|
||
if not os.path.exists(video_path):
|
||
raise HTTPException(404, f"Video not found: {req.video}")
|
||
|
||
est = _load_pose_estimator()
|
||
if not est:
|
||
raise HTTPException(501, "No pose estimator installed. Try: pip install rtmlib onnxruntime")
|
||
infer, backend = est
|
||
|
||
# Extract one frame at the requested timestamp with ffmpeg.
|
||
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tf:
|
||
frame_path = tf.name
|
||
try:
|
||
result = sp.run(
|
||
[
|
||
"/usr/bin/ffmpeg", "-y",
|
||
"-ss", str(req.time_s),
|
||
"-i", video_path,
|
||
"-frames:v", "1",
|
||
"-q:v", "2",
|
||
frame_path,
|
||
],
|
||
capture_output=True, timeout=30,
|
||
)
|
||
if result.returncode != 0:
|
||
raise HTTPException(500, f"ffmpeg frame extract failed: {result.stderr.decode()[:300]}")
|
||
|
||
pil = Image.open(frame_path).convert("RGB")
|
||
finally:
|
||
try:
|
||
os.unlink(frame_path)
|
||
except OSError:
|
||
pass
|
||
|
||
try:
|
||
people = infer(pil)
|
||
except Exception as e:
|
||
raise HTTPException(500, f"Pose estimation failed: {e}")
|
||
|
||
best = _best_person(people)
|
||
if best is None:
|
||
raise HTTPException(422, "No person detected in the extracted frame")
|
||
|
||
return {
|
||
"status": "success",
|
||
"backend": backend,
|
||
"width": pil.width,
|
||
"height": pil.height,
|
||
"names": POSE_KEYPOINT_NAMES,
|
||
"skeleton": POSE_SKELETON,
|
||
"keypoints": best,
|
||
"time_s": req.time_s,
|
||
}
|
||
|
||
|
||
class PoseGenRequest(BaseModel):
|
||
model_filename: str
|
||
keypoints: list[list[float]] # 17 × [x, y, conf] in original image pixels
|
||
width: int # original image dimensions the keypoints came from
|
||
height: int
|
||
gesture_name: str | None = None # preset name for prompt enrichment and tagging
|
||
prompt: str | None = None # full override; auto-built if None
|
||
seed: int = -1
|
||
extra_filename: str | None = None # optional image3 extra reference
|
||
|
||
|
||
def _keypoints_to_pose_text(kpts: list, width: int, height: int) -> str:
|
||
"""Convert COCO-17 keypoints to a natural-language pose description for Qwen."""
|
||
MIN_CONF = 0.25
|
||
W, H = max(width, 1), max(height, 1)
|
||
|
||
def vis(idx):
|
||
return idx < len(kpts) and kpts[idx][2] >= MIN_CONF
|
||
|
||
def pt(idx):
|
||
return (kpts[idx][0] / W, kpts[idx][1] / H) if vis(idx) else None
|
||
|
||
parts = []
|
||
|
||
# --- head orientation ---
|
||
nose, l_ear, r_ear = pt(0), pt(3), pt(4)
|
||
if l_ear and r_ear:
|
||
ear_mid_x = (l_ear[0] + r_ear[0]) / 2
|
||
if nose:
|
||
dx = nose[0] - ear_mid_x
|
||
if dx < -0.05:
|
||
parts.append("head turned to the left")
|
||
elif dx > 0.05:
|
||
parts.append("head turned to the right")
|
||
else:
|
||
parts.append("head facing forward")
|
||
elif l_ear and not r_ear:
|
||
parts.append("head turned strongly to the right")
|
||
elif r_ear and not l_ear:
|
||
parts.append("head turned strongly to the left")
|
||
|
||
# --- body rotation (shoulders vs hips) ---
|
||
l_sh, r_sh = pt(5), pt(6)
|
||
l_hip, r_hip = pt(11), pt(12)
|
||
if l_sh and r_sh:
|
||
sh_w = abs(r_sh[0] - l_sh[0])
|
||
if l_hip and r_hip:
|
||
hip_w = abs(r_hip[0] - l_hip[0])
|
||
ratio = sh_w / max(hip_w, 0.01)
|
||
if ratio < 0.5:
|
||
parts.append("body rotated ~90° (side profile)")
|
||
elif ratio < 0.75:
|
||
parts.append("body rotated ~45° (three-quarter view)")
|
||
else:
|
||
parts.append("body facing forward")
|
||
|
||
# --- arm positions ---
|
||
for side, sh_i, el_i, wr_i in [("left", 5, 7, 9), ("right", 6, 8, 10)]:
|
||
sh, el, wr = pt(sh_i), pt(el_i), pt(wr_i)
|
||
if not sh:
|
||
continue
|
||
if wr:
|
||
dy = sh[1] - wr[1] # positive = wrist above shoulder
|
||
dx = abs(wr[0] - sh[0])
|
||
if dy > 0.15:
|
||
if el and pt(el_i):
|
||
el_dy = sh[1] - el[1]
|
||
if el_dy > 0.05:
|
||
parts.append(f"{side} arm raised straight up")
|
||
else:
|
||
parts.append(f"{side} arm raised with elbow bent")
|
||
else:
|
||
parts.append(f"{side} arm raised above shoulder")
|
||
elif dx > 0.25:
|
||
parts.append(f"{side} arm extended sideways")
|
||
elif el:
|
||
el_dx = abs(el[0] - sh[0])
|
||
if el_dx > 0.15 and abs(wr[1] - sh[1]) < 0.1:
|
||
parts.append(f"{side} arm extended forward")
|
||
else:
|
||
parts.append(f"{side} arm at side")
|
||
else:
|
||
parts.append(f"{side} arm at side")
|
||
elif el:
|
||
el_dy = sh[1] - el[1]
|
||
if el_dy > 0.1:
|
||
parts.append(f"{side} elbow raised")
|
||
|
||
# --- leg positions ---
|
||
for side, hip_i, kn_i, an_i in [("left", 11, 13, 15), ("right", 12, 14, 16)]:
|
||
hip, kn, an = pt(hip_i), pt(kn_i), pt(an_i)
|
||
if not hip:
|
||
continue
|
||
if an:
|
||
dx = abs(an[0] - hip[0])
|
||
dy = an[1] - hip[1] # positive = ankle below hip (normal standing)
|
||
if dy < 0.1:
|
||
parts.append(f"{side} leg raised / knee up")
|
||
elif dx > 0.2:
|
||
parts.append(f"{side} leg stepped out to the side")
|
||
elif kn:
|
||
kn_dx = abs(kn[0] - hip[0])
|
||
if kn_dx > 0.15:
|
||
parts.append(f"{side} leg bent / lunging")
|
||
else:
|
||
parts.append(f"{side} leg straight, standing")
|
||
else:
|
||
parts.append(f"{side} leg straight")
|
||
|
||
if not parts:
|
||
return "standing neutral pose"
|
||
return ", ".join(parts)
|
||
|
||
|
||
def _pad_for_pose(pil: Image.Image, keypoints: list, margin: float = 0.15) -> tuple[Image.Image, tuple[int,int,int,int]]:
|
||
"""Expand the image so target keypoints fit within the canvas.
|
||
|
||
Only expands when keypoints are actually outside [0,W]×[0,H]. Expansion is
|
||
symmetric on each axis (equal left/right, equal top/bottom) so the person
|
||
stays centred on the new canvas. Fill is black to match PIL's RGBA→RGB fill.
|
||
"""
|
||
W, H = pil.size
|
||
all_kps = [kp for kp in keypoints if len(kp) >= 2]
|
||
if not all_kps:
|
||
return pil, (0, 0, 0, 0)
|
||
|
||
xs = [kp[0] for kp in all_kps]
|
||
ys = [kp[1] for kp in all_kps]
|
||
|
||
# How far each side actually exceeds the image bounds (0 if within bounds).
|
||
over_l = max(0.0, -min(xs))
|
||
over_r = max(0.0, max(xs) - W)
|
||
over_t = max(0.0, -min(ys))
|
||
over_b = max(0.0, max(ys) - H)
|
||
|
||
# If nothing is outside, no padding needed at all.
|
||
if over_l == over_r == over_t == over_b == 0:
|
||
return pil, (0, 0, 0, 0)
|
||
|
||
# Symmetric expansion per axis: take the larger overhang and add margin.
|
||
px = int(max(over_l, over_r) + margin * W)
|
||
py = int(max(over_t, over_b) + margin * H)
|
||
|
||
new_w, new_h = W + 2 * px, H + 2 * py
|
||
canvas = Image.new("RGB", (new_w, new_h), color=(0, 0, 0))
|
||
canvas.paste(pil, (px, py))
|
||
print(f"[pose-gen] padded {W}×{H} → {new_w}×{new_h} (±{px}px x, ±{py}px y)")
|
||
return canvas, (px, py, px, py)
|
||
|
||
|
||
def _pose_gen_worker(job_id: str, model_filename: str, prompt: str, seed: int,
|
||
keypoints: list | None = None,
|
||
gesture_name: str | None = None,
|
||
extra_filename: str | None = None):
|
||
output_dir = _load_output_dir()
|
||
try:
|
||
model_path = os.path.join(output_dir, model_filename)
|
||
model_pil = Image.open(model_path).convert("RGB")
|
||
|
||
# Auto-pad if the target pose would extend outside the current image bounds.
|
||
if keypoints:
|
||
model_pil, _padding = _pad_for_pose(model_pil, keypoints)
|
||
|
||
# No skeleton image — pose is encoded entirely in the text prompt.
|
||
extra_images = None
|
||
if extra_filename:
|
||
ep = os.path.join(output_dir, extra_filename)
|
||
extra_images = [Image.open(ep).convert("RGB")]
|
||
|
||
png_bytes = _run_pipeline(model_pil, prompt, seed, MAX_AREA, extra_images=extra_images)
|
||
ts = time.strftime("%Y%m%d_%H%M%S")
|
||
dir_part = "" if model_filename.startswith("_turntable/") else os.path.dirname(model_filename)
|
||
basename = os.path.basename(model_filename)
|
||
clean_basename = naming.get_base_name(basename)
|
||
if not clean_basename.lower().endswith(".png"):
|
||
clean_basename = os.path.splitext(clean_basename)[0] + ".png"
|
||
new_basename = f"{ts}_pose_{clean_basename}"
|
||
if dir_part:
|
||
out_name = f"{dir_part}/{new_basename}"
|
||
else:
|
||
out_name = new_basename
|
||
out_path = os.path.join(output_dir, out_name)
|
||
with open(out_path, "wb") as f:
|
||
f.write(png_bytes)
|
||
|
||
person = database.get_person(model_filename)
|
||
group_id = person[1] if person and person[1] else naming.get_base_name(os.path.basename(model_filename))
|
||
|
||
embedding = embeddings.generate_embedding(out_path)
|
||
next_order = database.get_next_sort_order(group_id)
|
||
refs = [model_filename]
|
||
if gesture_name:
|
||
refs.append(f"gesture:{gesture_name}")
|
||
database.upsert_person(
|
||
out_name, filepath=out_path, embedding=embedding,
|
||
group_id=group_id, prompt=prompt,
|
||
pose=gesture_name or "custom",
|
||
sort_order=next_order,
|
||
source_refs=json.dumps(refs),
|
||
)
|
||
_update_cached_file_meta(out_name, exists=True)
|
||
|
||
jobs[job_id]["status"] = "done"
|
||
jobs[job_id]["output"] = out_name
|
||
_invalidate_static()
|
||
except Exception as e:
|
||
print(f"[pose-gen] error: {e}")
|
||
jobs[job_id]["status"] = "error"
|
||
jobs[job_id]["error"] = str(e)
|
||
|
||
|
||
@app.post("/pose/render")
|
||
def render_pose_image(req: PoseSimilarRequest):
|
||
"""Render the supplied keypoints as an OpenPose-style skeleton PNG (base64) — for preview only."""
|
||
import base64
|
||
skeleton = _render_openpose_image(req.keypoints, req.width or 512, req.height or 768)
|
||
buf = io.BytesIO()
|
||
skeleton.save(buf, format="PNG")
|
||
return {"image": base64.b64encode(buf.getvalue()).decode()}
|
||
|
||
|
||
@app.post("/generate-with-pose")
|
||
def generate_with_pose(req: PoseGenRequest):
|
||
"""Generate the person in a specific body pose. Pose is encoded as text — no skeleton image
|
||
is passed to Qwen to avoid wireframe bleed-through in the output."""
|
||
output_dir = _load_output_dir()
|
||
model_path = os.path.join(output_dir, req.model_filename)
|
||
if not os.path.exists(model_path):
|
||
raise HTTPException(404, f"Model image not found: {req.model_filename}")
|
||
if req.extra_filename:
|
||
ep = os.path.join(output_dir, req.extra_filename)
|
||
if not os.path.exists(ep):
|
||
raise HTTPException(404, f"Extra reference not found: {req.extra_filename}")
|
||
|
||
if req.prompt:
|
||
prompt = req.prompt
|
||
else:
|
||
# Convert skeleton keypoints to natural-language description
|
||
pose_text = _keypoints_to_pose_text(req.keypoints, req.width, req.height)
|
||
if req.gesture_name:
|
||
pose_text = f'{req.gesture_name} gesture: {pose_text}'
|
||
prompt = (
|
||
f"Change the body pose of the person in image 1 to: {pose_text}. "
|
||
"Keep the person's face, hair, skin tone, clothing, and body proportions "
|
||
"exactly as they are — only the limb positions and body orientation change. "
|
||
"Transparent background."
|
||
)
|
||
|
||
job_id = uuid.uuid4().hex[:8]
|
||
jobs[job_id] = {"status": "running", "type": "pose_gen", "total": 1, "done": 0, "failed": 0}
|
||
threading.Thread(
|
||
target=_pose_gen_worker,
|
||
args=(job_id, req.model_filename, prompt, req.seed),
|
||
kwargs={
|
||
"keypoints": req.keypoints,
|
||
"gesture_name": req.gesture_name,
|
||
"extra_filename": req.extra_filename,
|
||
},
|
||
daemon=True,
|
||
).start()
|
||
return {"job_id": job_id, "model": req.model_filename}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Orbit preview — depth-card (fast) or Qwen turntable (quality)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class OrbitRequest(BaseModel):
|
||
filename: str # single image filename relative to output_dir
|
||
engine: str = "depth" # "depth" (fast depth-card) or "qwen" (near-real)
|
||
# depth-card params
|
||
n_frames: int = 36
|
||
parallax: float = 0.08
|
||
mode: str = "swing" # "swing" or "orbit"
|
||
fps: int = 24
|
||
max_angle_deg: float = 35.0
|
||
# qwen turntable params
|
||
n_views: int = 24
|
||
steps: int = 8
|
||
seed: int = 42
|
||
|
||
|
||
@app.post("/orbit")
|
||
def create_orbit(req: OrbitRequest):
|
||
"""
|
||
Build an orbit preview for one gallery image.
|
||
|
||
engine='depth': fast 2.5D depth-card parallax (seconds)
|
||
engine='qwen': near-real Qwen turntable — checks turntable cache first;
|
||
if a complete turntable exists, returns cached video immediately.
|
||
Otherwise queues generation (~9 min for 24 views).
|
||
|
||
Returns {job_id} for polling via GET /batch/{job_id}.
|
||
"""
|
||
output_dir = _load_output_dir()
|
||
img_path = os.path.join(output_dir, req.filename)
|
||
if not os.path.exists(img_path):
|
||
raise HTTPException(status_code=400, detail=f"Image not found: {req.filename}")
|
||
|
||
if req.engine == "qwen":
|
||
return _create_qwen_orbit(req, output_dir, img_path)
|
||
|
||
# --- depth-card ---
|
||
job_id = uuid.uuid4().hex[:8]
|
||
jobs[job_id] = {"status": "running", "type": "orbit", "total": 1, "done": 0, "failed": 0}
|
||
orbit_out = os.path.join(output_dir, f"orbit_{job_id}")
|
||
|
||
def _worker():
|
||
try:
|
||
try:
|
||
from orbit_module import run_orbit_pipeline
|
||
except ImportError:
|
||
sys.path.insert(0, _HERE)
|
||
from orbit_module import run_orbit_pipeline
|
||
|
||
result = run_orbit_pipeline(
|
||
image_path=img_path,
|
||
output_dir=orbit_out,
|
||
n_frames=req.n_frames,
|
||
parallax_strength=req.parallax,
|
||
mode=req.mode,
|
||
fps=req.fps,
|
||
max_angle_deg=req.max_angle_deg,
|
||
debug=True,
|
||
)
|
||
mp4_dst_name = f"orbit_{job_id}.mp4"
|
||
shutil.copy2(result["video_path"], os.path.join(output_dir, mp4_dst_name))
|
||
jobs[job_id]["status"] = "done"
|
||
jobs[job_id]["done"] = 1
|
||
jobs[job_id]["video_filename"] = mp4_dst_name
|
||
jobs[job_id]["has_alpha"] = bool(result.get("has_alpha", False))
|
||
jobs[job_id]["orbit_dir"] = orbit_out
|
||
except Exception as e:
|
||
import traceback
|
||
print(f"[orbit] error: {e}\n{traceback.format_exc()}")
|
||
jobs[job_id]["status"] = "error"
|
||
jobs[job_id]["error"] = str(e)
|
||
|
||
threading.Thread(target=_worker, daemon=True).start()
|
||
return {"job_id": job_id}
|
||
|
||
|
||
def _create_qwen_orbit(req: OrbitRequest, output_dir: str, img_path: str) -> dict:
|
||
"""
|
||
Qwen turntable orbit: check cache first, else generate in background.
|
||
The group_id is derived from the DB row for the requested image.
|
||
If no group, falls back to a per-file cache keyed on the bare filename.
|
||
"""
|
||
import turntable_cache as tc
|
||
|
||
# Resolve group_id for this image
|
||
group_id = None
|
||
try:
|
||
persons = database.list_persons()
|
||
for row in persons:
|
||
if row[0] == req.filename:
|
||
group_id = row[2]
|
||
break
|
||
except Exception:
|
||
pass
|
||
cache_key = str(group_id) if group_id else f"file_{req.filename}"
|
||
|
||
# Serve cached turntable immediately if already complete
|
||
cached_video = tc.get_group_video(output_dir, cache_key)
|
||
if cached_video and os.path.exists(cached_video):
|
||
job_id = uuid.uuid4().hex[:8]
|
||
rel = os.path.relpath(cached_video, output_dir)
|
||
# Build ordered frame list for the frame-flipper player
|
||
state = tc.load_state(output_dir, cache_key)
|
||
frames = []
|
||
if state:
|
||
for deg in state.get("angles", []):
|
||
dk = tc.deg_key(deg)
|
||
p = state.get("views", {}).get(dk)
|
||
if p and os.path.exists(p):
|
||
frames.append(os.path.relpath(p, output_dir).replace("\\", "/"))
|
||
jobs[job_id] = {
|
||
"status": "done", "type": "orbit_qwen",
|
||
"total": 1, "done": 1, "failed": 0,
|
||
"video_filename": rel,
|
||
"frames": frames,
|
||
"cached": True,
|
||
}
|
||
return {"job_id": job_id, "cached": True}
|
||
|
||
# Otherwise generate in a background thread
|
||
job_id = uuid.uuid4().hex[:8]
|
||
jobs[job_id] = {
|
||
"status": "running", "type": "orbit_qwen",
|
||
"total": req.n_views, "done": 0, "failed": 0,
|
||
}
|
||
|
||
def _qwen_worker():
|
||
try:
|
||
import sys as _s
|
||
if _HERE not in _s.path:
|
||
_s.path.insert(0, _HERE)
|
||
from orbit_qwen import run_qwen_orbit, yaw_prompt
|
||
|
||
qwen_out = os.path.join(tc.cache_dir(output_dir, cache_key))
|
||
os.makedirs(qwen_out, exist_ok=True)
|
||
|
||
def _prog(i, n, deg):
|
||
jobs[job_id]["done"] = i
|
||
jobs[job_id]["status_detail"] = f"{i}/{n} views ({int(deg)}°)"
|
||
|
||
result = run_qwen_orbit(
|
||
image_path=img_path,
|
||
output_dir=qwen_out,
|
||
n_views=req.n_views,
|
||
seed=req.seed,
|
||
mode="turntable",
|
||
steps=req.steps,
|
||
on_progress=_prog,
|
||
)
|
||
|
||
# Also update the persistent cache state so idle daemon knows it's done
|
||
state = tc.load_state(output_dir, cache_key) or tc.init_state(
|
||
output_dir, cache_key, img_path, req.filename,
|
||
n_views=req.n_views, seed=req.seed, steps=req.steps,
|
||
)
|
||
for v in result["views"]:
|
||
tc.mark_view_done(output_dir, cache_key, state, v["deg"], v["path"])
|
||
# Register frame in DB for manual job too
|
||
try:
|
||
vname = os.path.relpath(v["path"], output_dir).replace("\\", "/")
|
||
# Extract angle index from path if possible, or just use degree
|
||
angle_idx = int(os.path.basename(v["path"]).split("_")[1])
|
||
database.upsert_person(
|
||
vname,
|
||
filepath=v["path"],
|
||
group_id=cache_key,
|
||
prompt=yaw_prompt(v["deg"]),
|
||
source_refs=json.dumps([req.filename]),
|
||
sort_order=200 + angle_idx,
|
||
pose=f"Orbit {int(v['deg'])}°",
|
||
tags=["ORBIT"]
|
||
)
|
||
except Exception as db_err:
|
||
print(f"[orbit-qwen] DB frame error: {db_err}")
|
||
|
||
tc.mark_completed(output_dir, cache_key, state, "")
|
||
_write_turntable_static()
|
||
|
||
# Include frame URLs in job result for frame-flipper player
|
||
frames = []
|
||
for v in result["views"]:
|
||
frames.append(os.path.relpath(v["path"], output_dir).replace("\\", "/"))
|
||
jobs[job_id]["status"] = "done"
|
||
jobs[job_id]["done"] = req.n_views
|
||
jobs[job_id]["video_filename"] = None
|
||
jobs[job_id]["frames"] = frames
|
||
except Exception as e:
|
||
import traceback
|
||
print(f"[orbit-qwen] error: {e}\n{traceback.format_exc()}")
|
||
jobs[job_id]["status"] = "error"
|
||
jobs[job_id]["error"] = str(e)
|
||
|
||
threading.Thread(target=_qwen_worker, daemon=True).start()
|
||
return {"job_id": job_id, "cached": False}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Turntable status / control endpoints
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@app.get("/turntable/status")
|
||
def get_turntable_status():
|
||
"""
|
||
Return background turntable generation state for all groups.
|
||
Used by the frontend to show a global progress badge.
|
||
"""
|
||
import turntable_cache as tc
|
||
output_dir = _load_output_dir()
|
||
summary = tc.get_status_summary(output_dir)
|
||
n_complete = sum(1 for v in summary.values() if v["completed"])
|
||
return {
|
||
"groups": summary,
|
||
"n_complete": n_complete,
|
||
"n_total": len(summary),
|
||
"is_generating": _idle_turntable_busy,
|
||
"is_paused": _idle_turntable_paused,
|
||
"idle_seconds": round(time.time() - _last_user_generation_time, 1),
|
||
}
|
||
|
||
|
||
@app.get("/turntable/status/{group_id:path}")
|
||
def get_turntable_status_group(group_id: str):
|
||
"""Return turntable state for a specific group."""
|
||
import turntable_cache as tc
|
||
output_dir = _load_output_dir()
|
||
state = tc.load_state(output_dir, group_id)
|
||
if not state:
|
||
return {"exists": False, "group_id": group_id}
|
||
return {
|
||
"exists": True,
|
||
"completed": state.get("completed", False),
|
||
"n_done": len(state.get("views", {})),
|
||
"n_total": state.get("n_views", 24),
|
||
"video_path": state.get("video_path"),
|
||
"preferred_filename": state.get("preferred_filename"),
|
||
}
|
||
|
||
|
||
@app.post("/turntable/pause")
|
||
def pause_turntable():
|
||
"""Pause idle background generation."""
|
||
global _idle_turntable_paused
|
||
_idle_turntable_paused = True
|
||
_write_turntable_static()
|
||
return {"paused": True}
|
||
|
||
|
||
@app.post("/turntable/resume")
|
||
def resume_turntable():
|
||
"""Resume idle background generation."""
|
||
global _idle_turntable_paused
|
||
_idle_turntable_paused = False
|
||
_write_turntable_static()
|
||
return {"paused": False}
|
||
|
||
|
||
@app.delete("/turntable/{group_id:path}")
|
||
def reset_turntable(group_id: str):
|
||
"""Wipe all cached orbit data for a group (files + DB records) to force re-generation."""
|
||
import turntable_cache as tc
|
||
output_dir = _load_output_dir()
|
||
tc.delete_state(output_dir, group_id)
|
||
|
||
# Remove from DB
|
||
try:
|
||
persons = database.list_persons(include_archived=True)
|
||
prefix = f"_turntable/{group_id}/"
|
||
for p in persons:
|
||
fname = p[0]
|
||
if fname.startswith(prefix):
|
||
database.delete_person(fname)
|
||
except Exception as e:
|
||
print(f"[turntable] DB cleanup error: {e}")
|
||
|
||
_invalidate_static()
|
||
return {"status": "reset", "group_id": group_id}
|
||
|
||
|
||
def _detect_people_count(keypoints: list) -> int:
|
||
"""Detect the number of people in an image from keypoints.
|
||
|
||
For now, we assume only one person is detected by the pose estimator.
|
||
This could be expanded to detect multiple people if needed.
|
||
"""
|
||
return 1 if keypoints else 0
|
||
|
||
|
||
def _detect_anatomical_completeness(keypoints: list, width: int = None, height: int = None) -> bool:
|
||
"""Detect if the person has complete anatomical structure.
|
||
|
||
Returns True if all major body parts are visible (head, torso, arms, legs)
|
||
and are fully contained within the frame (not cropped at boundaries).
|
||
"""
|
||
if not keypoints or len(keypoints) < 17:
|
||
return False
|
||
|
||
MIN_VISIBILITY = 0.3
|
||
|
||
# 1. Presence of major body segments (requires visibility >= 0.3)
|
||
has_head = any(idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY for idx in [0, 1, 2, 3, 4])
|
||
has_shoulders = any(idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY for idx in [5, 6])
|
||
has_hips = any(idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY for idx in [11, 12])
|
||
has_knees = any(idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY for idx in [13, 14])
|
||
has_ankles = any(idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY for idx in [15, 16])
|
||
|
||
# If any major body segment is completely missing, the anatomy is not complete (e.g., cropped above knees/hips/chest)
|
||
if not (has_head and has_shoulders and has_hips and has_knees and has_ankles):
|
||
return False
|
||
|
||
# 2. Boundary cropping check (if image dimensions are provided)
|
||
if width and height:
|
||
# Check if head is too close to top edge
|
||
head_kpts = [keypoints[idx] for idx in [0, 1, 2, 3, 4] if idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY]
|
||
if head_kpts:
|
||
min_head_y = min(kp[1] for kp in head_kpts)
|
||
if min_head_y < height * 0.02: # head is cropped at top
|
||
return False
|
||
|
||
# Check if ankles (feet) are too close to bottom edge
|
||
ankle_kpts = [keypoints[idx] for idx in [15, 16] if idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY]
|
||
if ankle_kpts:
|
||
max_ankle_y = max(kp[1] for kp in ankle_kpts)
|
||
if max_ankle_y > height * 0.98: # ankles/feet are cropped at bottom
|
||
return False
|
||
|
||
# Check if wrists (hands) are too close to left/right edge
|
||
wrist_kpts = [keypoints[idx] for idx in [9, 10] if idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY]
|
||
if wrist_kpts:
|
||
min_wrist_x = min(kp[0] for kp in wrist_kpts)
|
||
max_wrist_x = max(kp[0] for kp in wrist_kpts)
|
||
if min_wrist_x < width * 0.02 or max_wrist_x > width * 0.98: # wrists/hands are cropped at sides
|
||
return False
|
||
|
||
return True
|
||
|
||
|
||
def _detect_facial_direction(keypoints: list) -> str:
|
||
"""Detect the facial direction from keypoints, including fine-grained gaze details (up/down/left/right).
|
||
|
||
Returns a string describing the head/gaze orientation.
|
||
"""
|
||
if not keypoints or len(keypoints) < 17:
|
||
return "unknown"
|
||
|
||
# Key points for face direction detection
|
||
# Nose (0), left eye (1), right eye (2), left ear (3), right ear (4)
|
||
nose = keypoints[0] if len(keypoints) > 0 and keypoints[0][2] >= 0.3 else None
|
||
l_eye = keypoints[1] if len(keypoints) > 1 and keypoints[1][2] >= 0.3 else None
|
||
r_eye = keypoints[2] if len(keypoints) > 2 and keypoints[2][2] >= 0.3 else None
|
||
l_ear = keypoints[3] if len(keypoints) > 3 and keypoints[3][2] >= 0.3 else None
|
||
r_ear = keypoints[4] if len(keypoints) > 4 and keypoints[4][2] >= 0.3 else None
|
||
|
||
if not nose:
|
||
return "unknown"
|
||
|
||
# 1. Determine horizontal direction
|
||
horiz = "forward"
|
||
if l_ear and r_ear:
|
||
ear_mid_x = (l_ear[0] + r_ear[0]) / 2
|
||
dx = nose[0] - ear_mid_x
|
||
# Normalize dx by distance between ears to make it scale invariant
|
||
ear_dist = abs(l_ear[0] - r_ear[0])
|
||
if ear_dist > 0:
|
||
norm_dx = dx / ear_dist
|
||
if norm_dx < -0.06:
|
||
horiz = "left"
|
||
elif norm_dx > 0.06:
|
||
horiz = "right"
|
||
elif l_ear and not r_ear:
|
||
horiz = "strongly right"
|
||
elif r_ear and not l_ear:
|
||
horiz = "strongly left"
|
||
|
||
# 2. Determine vertical direction (up/down)
|
||
vert = "level"
|
||
# Try using eyes first as they are closer to the nose
|
||
if l_eye and r_eye:
|
||
eye_y = (l_eye[1] + r_eye[1]) / 2
|
||
eye_dist = abs(l_eye[0] - r_eye[0])
|
||
if eye_dist > 0:
|
||
v_ratio = (nose[1] - eye_y) / eye_dist
|
||
if v_ratio < 0.15:
|
||
vert = "up"
|
||
elif v_ratio > 0.65:
|
||
vert = "down"
|
||
elif l_ear and r_ear:
|
||
ear_y = (l_ear[1] + r_ear[1]) / 2
|
||
ear_dist = abs(l_ear[0] - r_ear[0])
|
||
if ear_dist > 0:
|
||
v_ratio = (nose[1] - ear_y) / ear_dist
|
||
if v_ratio < -0.1:
|
||
vert = "up"
|
||
elif v_ratio > 0.3:
|
||
vert = "down"
|
||
|
||
# 3. Combine horizontal and vertical
|
||
if horiz == "forward":
|
||
if vert == "level":
|
||
return "looking forward"
|
||
elif vert == "up":
|
||
return "looking forward and up"
|
||
elif vert == "down":
|
||
return "looking forward and down"
|
||
elif horiz == "left":
|
||
if vert == "level":
|
||
return "looking left"
|
||
elif vert == "up":
|
||
return "looking left and up"
|
||
elif vert == "down":
|
||
return "looking left and down"
|
||
elif horiz == "right":
|
||
if vert == "level":
|
||
return "looking right"
|
||
elif vert == "up":
|
||
return "looking right and up"
|
||
elif vert == "down":
|
||
return "looking right and down"
|
||
elif horiz == "strongly left":
|
||
if vert == "level":
|
||
return "looking strongly left"
|
||
elif vert == "up":
|
||
return "looking strongly left and up"
|
||
elif vert == "down":
|
||
return "looking strongly left and down"
|
||
elif horiz == "strongly right":
|
||
if vert == "level":
|
||
return "looking strongly right"
|
||
elif vert == "up":
|
||
return "looking strongly right and up"
|
||
elif vert == "down":
|
||
return "looking strongly right and down"
|
||
|
||
return "looking forward"
|
||
|
||
|
||
def _estimate_bbox_for_tag(tag: str, keypoints: list, width: int, height: int, alpha_bbox: list = None) -> list:
|
||
"""Estimate a bounding box for a given tag using pose keypoints or alpha bounding box.
|
||
|
||
Returns a list [x1, y1, x2, y2] of pixel coordinates, or None.
|
||
"""
|
||
import math
|
||
if not keypoints or len(keypoints) < 17:
|
||
if alpha_bbox:
|
||
return [int(v) for v in alpha_bbox]
|
||
return [0, 0, width, height]
|
||
|
||
tag_lower = tag.lower().replace("_", " ")
|
||
|
||
# Define terms lists mapped to anatomical structures
|
||
head_terms = ["hair", "head", "face", "eye", "eyes", "nose", "ear", "ears", "mouth", "makeup", "eyebrow", "eyebrows", "glasses", "sunglasses", "earrings", "jewelry", "blush", "necklace", "collar", "hat", "cap", "crown", "smile", "gaze", "cheek", "teeth", "lips"]
|
||
chest_terms = ["breast", "nipple", "nipples", "breasts", "chest", "cleavage", "bra", "bikini top", "top", "shirt", "collarbone", "pendant"]
|
||
stomach_terms = ["navel", "stomach", "belly", "abs", "midriff", "waist", "panties", "underwear", "pelvis", "bikini bottom", "hips", "hip"]
|
||
arm_terms = ["arm", "arms", "hand", "hands", "wrist", "wrists", "elbow", "elbows", "finger", "fingers", "sleeve", "sleeves", "glove", "gloves"]
|
||
leg_terms = ["leg", "legs", "thigh", "thighs", "knee", "knees", "calf", "calves", "foot", "feet", "ankle", "ankles", "shoe", "shoes", "socks", "sock", "boots", "boot"]
|
||
|
||
bbox = None
|
||
|
||
# 1. Head/Face
|
||
if any(term in tag_lower for term in head_terms):
|
||
head_kpts = [keypoints[i] for i in range(5) if i < len(keypoints) and keypoints[i][2] >= 0.3]
|
||
if head_kpts:
|
||
xs = [kp[0] for kp in head_kpts]
|
||
ys = [kp[1] for kp in head_kpts]
|
||
min_x, max_x = min(xs), max(xs)
|
||
min_y, max_y = min(ys), max(ys)
|
||
|
||
# Determine padding based on shoulder distance if available
|
||
if len(keypoints) > 6 and keypoints[5][2] >= 0.3 and keypoints[6][2] >= 0.3:
|
||
shoulder_dist = math.dist(keypoints[5][:2], keypoints[6][:2])
|
||
pad_x = shoulder_dist * 0.35
|
||
pad_y = shoulder_dist * 0.45
|
||
else:
|
||
pad_x = max(max_x - min_x, width * 0.08)
|
||
pad_y = max(max_y - min_y, height * 0.08)
|
||
|
||
bbox = [
|
||
min_x - pad_x,
|
||
min_y - pad_y * 1.3, # pull top higher to cover hair/hats
|
||
max_x + pad_x,
|
||
max_y + pad_y * 0.7
|
||
]
|
||
|
||
# 2. Chest/Breasts
|
||
elif any(term in tag_lower for term in chest_terms):
|
||
if len(keypoints) > 6 and keypoints[5][2] >= 0.3 and keypoints[6][2] >= 0.3:
|
||
sh_mid_x = (keypoints[5][0] + keypoints[6][0]) / 2
|
||
sh_mid_y = (keypoints[5][1] + keypoints[6][1]) / 2
|
||
sh_dist = math.dist(keypoints[5][:2], keypoints[6][:2])
|
||
|
||
if len(keypoints) > 12 and keypoints[11][2] >= 0.3 and keypoints[12][2] >= 0.3:
|
||
hip_mid_y = (keypoints[11][1] + keypoints[12][1]) / 2
|
||
torso_h = hip_mid_y - sh_mid_y
|
||
else:
|
||
torso_h = sh_dist * 1.2
|
||
|
||
chest_center_y = sh_mid_y + torso_h * 0.28
|
||
chest_w = sh_dist * 0.95
|
||
chest_h = torso_h * 0.42
|
||
bbox = [
|
||
sh_mid_x - chest_w / 2,
|
||
chest_center_y - chest_h / 2,
|
||
sh_mid_x + chest_w / 2,
|
||
chest_center_y + chest_h / 2
|
||
]
|
||
|
||
# 3. Midriff/Pelvis/Hips/Underwear
|
||
elif any(term in tag_lower for term in stomach_terms):
|
||
if len(keypoints) > 12 and keypoints[11][2] >= 0.3 and keypoints[12][2] >= 0.3:
|
||
hip_mid_x = (keypoints[11][0] + keypoints[12][0]) / 2
|
||
hip_mid_y = (keypoints[11][1] + keypoints[12][1]) / 2
|
||
hip_dist = math.dist(keypoints[11][:2], keypoints[12][:2])
|
||
|
||
if len(keypoints) > 6 and keypoints[5][2] >= 0.3 and keypoints[6][2] >= 0.3:
|
||
sh_mid_y = (keypoints[5][1] + keypoints[6][1]) / 2
|
||
torso_h = hip_mid_y - sh_mid_y
|
||
else:
|
||
torso_h = hip_dist * 1.5
|
||
|
||
if any(term in tag_lower for term in ["navel", "stomach", "belly", "abs", "midriff", "waist"]):
|
||
center_y = hip_mid_y - torso_h * 0.22
|
||
box_h = torso_h * 0.32
|
||
box_w = hip_dist * 1.15
|
||
else: # panties, underwear, pelvis, bikini bottom, hips
|
||
center_y = hip_mid_y + torso_h * 0.05
|
||
box_h = torso_h * 0.42
|
||
box_w = hip_dist * 1.25
|
||
|
||
bbox = [
|
||
hip_mid_x - box_w / 2,
|
||
center_y - box_h / 2,
|
||
hip_mid_x + box_w / 2,
|
||
center_y + box_h / 2
|
||
]
|
||
|
||
# 4. Arms
|
||
elif any(term in tag_lower for term in arm_terms):
|
||
arm_indices = [5, 6, 7, 8, 9, 10]
|
||
xs = [keypoints[i][0] for i in arm_indices if i < len(keypoints) and keypoints[i][2] >= 0.3]
|
||
ys = [keypoints[i][1] for i in arm_indices if i < len(keypoints) and keypoints[i][2] >= 0.3]
|
||
if xs:
|
||
min_x, max_x = min(xs), max(xs)
|
||
min_y, max_y = min(ys), max(ys)
|
||
pad_x = (max_x - min_x) * 0.12 + 15
|
||
pad_y = (max_y - min_y) * 0.12 + 15
|
||
bbox = [min_x - pad_x, min_y - pad_y, max_x + pad_x, max_y + pad_y]
|
||
|
||
# 5. Legs
|
||
elif any(term in tag_lower for term in leg_terms):
|
||
leg_indices = [11, 12, 13, 14, 15, 16]
|
||
xs = [keypoints[i][0] for i in leg_indices if i < len(keypoints) and keypoints[i][2] >= 0.3]
|
||
ys = [keypoints[i][1] for i in leg_indices if i < len(keypoints) and keypoints[i][2] >= 0.3]
|
||
if xs:
|
||
min_x, max_x = min(xs), max(xs)
|
||
min_y, max_y = min(ys), max(ys)
|
||
pad_x = (max_x - min_x) * 0.12 + 15
|
||
pad_y = (max_y - min_y) * 0.12 + 15
|
||
bbox = [min_x - pad_x, min_y - pad_y, max_x + pad_x, max_y + pad_y]
|
||
|
||
# 6. Fallback or general full-body
|
||
if bbox is None:
|
||
visible_kpts = [k for k in keypoints if k[2] >= 0.3]
|
||
if visible_kpts:
|
||
xs = [k[0] for k in visible_kpts]
|
||
ys = [k[1] for k in visible_kpts]
|
||
min_x, max_x = min(xs), max(xs)
|
||
min_y, max_y = min(ys), max(ys)
|
||
pad_x = (max_x - min_x) * 0.15 + 20
|
||
pad_y = (max_y - min_y) * 0.12 + 20
|
||
bbox = [min_x - pad_x, min_y - pad_y, max_x + pad_x, max_y + pad_y]
|
||
elif alpha_bbox:
|
||
bbox = list(alpha_bbox)
|
||
else:
|
||
bbox = [0, 0, width, height]
|
||
|
||
# Clip to image bounds and format
|
||
x1 = max(0, min(int(bbox[0]), width))
|
||
y1 = max(0, min(int(bbox[1]), height))
|
||
x2 = max(0, min(int(bbox[2]), width))
|
||
y2 = max(0, min(int(bbox[3]), height))
|
||
|
||
# Ensure some valid size
|
||
if x2 <= x1:
|
||
x2 = min(x1 + 10, width)
|
||
if y2 <= y1:
|
||
y2 = min(y1 + 10, height)
|
||
|
||
return [x1, y1, x2, y2]
|
||
|
||
|
||
def _detect_objects(pil_img: Image.Image, keypoints: list = None) -> list:
|
||
"""Detect objects in the image using WD tagger.
|
||
|
||
Returns a list of detected objects with bounding box coordinates.
|
||
"""
|
||
try:
|
||
# Run tagger with lower threshold to capture more objects
|
||
tags = _run_tagger(pil_img, threshold=0.2)
|
||
|
||
# Filter for object-related tags (general and character categories)
|
||
objects = []
|
||
|
||
width, height = pil_img.size
|
||
alpha_bbox = None
|
||
if pil_img.mode == 'RGBA':
|
||
alpha = pil_img.split()[-1]
|
||
alpha_bbox = alpha.getbbox()
|
||
|
||
for t in tags:
|
||
if t["cat"] in (0, 4): # general and character categories
|
||
bbox = _estimate_bbox_for_tag(t["tag"], keypoints, width, height, alpha_bbox)
|
||
objects.append({
|
||
"tag": t["tag"],
|
||
"score": t["score"],
|
||
"bbox": bbox
|
||
})
|
||
return objects
|
||
except Exception as e:
|
||
print(f"[object-detection] Error: {e}")
|
||
return []
|
||
|
||
|
||
def _process_image_for_metadata(filename: str):
|
||
"""Process a single image to extract metadata for the knowledge base.
|
||
|
||
This function extracts people count, anatomical completeness, facial direction,
|
||
and objects from an image using pose estimation and WD tagger, as well as
|
||
the pose description and COCO-17 skeleton coordinates.
|
||
"""
|
||
if not filename.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
|
||
print(f"[metadata] Skipping non-image file: {filename}")
|
||
return None
|
||
|
||
try:
|
||
output_dir = _load_output_dir()
|
||
fpath = os.path.join(output_dir, filename)
|
||
|
||
if not os.path.exists(fpath):
|
||
return None
|
||
|
||
pil_img = Image.open(fpath)
|
||
|
||
# Get pose estimation
|
||
est = _load_pose_estimator()
|
||
if not est:
|
||
print("[metadata] No pose estimator available")
|
||
return None
|
||
|
||
infer, _ = est
|
||
people = infer(pil_img)
|
||
best_person = _best_person(people)
|
||
|
||
# Extract metadata
|
||
width, height = pil_img.size
|
||
people_count = _detect_people_count(best_person)
|
||
anatomical_completeness = _detect_anatomical_completeness(best_person, width, height)
|
||
facial_direction = _detect_facial_direction(best_person)
|
||
|
||
# ALSO extract pose description and pose skeleton
|
||
pose_desc = None
|
||
pose_skel_json = None
|
||
if best_person is not None:
|
||
pose_desc = _describe_pose(best_person)
|
||
pose_skel_json = json.dumps(best_person)
|
||
desc = _pose_descriptor(best_person)
|
||
if desc is not None:
|
||
try:
|
||
_save_pose_index_entry(filename, desc)
|
||
except Exception as e:
|
||
print(f"[pose] index save failed for {filename}: {e}")
|
||
|
||
# Detect objects
|
||
objects = _detect_objects(pil_img, keypoints=best_person)
|
||
|
||
# Update database with new metadata
|
||
database.upsert_person(
|
||
filename,
|
||
people_count=people_count,
|
||
anatomical_completeness=anatomical_completeness,
|
||
facial_direction=facial_direction,
|
||
objects=objects if objects else None,
|
||
pose_description=pose_desc,
|
||
pose_skeleton=pose_skel_json
|
||
)
|
||
|
||
_update_cached_file_meta(filename, exists=True)
|
||
_invalidate_static()
|
||
|
||
return {
|
||
"filename": filename,
|
||
"people_count": people_count,
|
||
"anatomical_completeness": anatomical_completeness,
|
||
"facial_direction": facial_direction,
|
||
"objects": objects,
|
||
"pose_description": pose_desc,
|
||
"pose_skeleton": pose_skel_json
|
||
}
|
||
except Exception as e:
|
||
print(f"[metadata] Error processing {filename}: {e}")
|
||
return None
|
||
|
||
|
||
class BackfillMetadataRequest(BaseModel):
|
||
filenames: list[str] | None = None # If None, process all images in DB
|
||
force: bool = False
|
||
|
||
|
||
import asyncio
|
||
from concurrent.futures import ThreadPoolExecutor as _ThreadPoolExecutor
|
||
_metadata_executor = _ThreadPoolExecutor(max_workers=1, thread_name_prefix="metadata")
|
||
|
||
|
||
@app.post("/images/invalidate-metadata")
|
||
def invalidate_metadata():
|
||
"""Invalidate all metadata records by setting their columns to NULL.
|
||
This enables full reprocessing via backfill or the background idle loop.
|
||
"""
|
||
try:
|
||
database.invalidate_all_metadata()
|
||
_failed_backfill_filenames.clear()
|
||
_invalidate_static()
|
||
return {"status": "success", "message": "All earlier metadata has been invalidated successfully."}
|
||
except Exception as e:
|
||
print(f"[metadata] Invalidation error: {e}")
|
||
raise HTTPException(500, f"Invalidation failed: {str(e)}")
|
||
|
||
|
||
@app.post("/images/backfill-metadata")
|
||
async def backfill_metadata(req: BackfillMetadataRequest):
|
||
"""Backfill metadata for existing images in the database.
|
||
|
||
This endpoint processes all existing images to extract and store new metadata:
|
||
people count, anatomical completeness, facial direction, and objects.
|
||
"""
|
||
try:
|
||
if req.force:
|
||
_failed_backfill_filenames.clear()
|
||
|
||
# Get list of all image files
|
||
if req.filenames is not None:
|
||
filenames = req.filenames
|
||
else:
|
||
persons = database.list_persons()
|
||
filenames = [p[0] for p in persons if p[0].lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
|
||
|
||
# Process each image
|
||
processed_count = 0
|
||
failed_count = 0
|
||
results = []
|
||
|
||
loop = asyncio.get_running_loop()
|
||
for filename in filenames:
|
||
try:
|
||
result = await loop.run_in_executor(_metadata_executor, _process_image_for_metadata, filename)
|
||
if result:
|
||
results.append(result)
|
||
processed_count += 1
|
||
else:
|
||
failed_count += 1
|
||
except Exception as e:
|
||
print(f"[metadata] Failed to process {filename}: {e}")
|
||
failed_count += 1
|
||
|
||
return {
|
||
"status": "completed",
|
||
"processed": processed_count,
|
||
"failed": failed_count,
|
||
"total": len(filenames),
|
||
"results": results[:10] # Return first 10 results for preview
|
||
}
|
||
except Exception as e:
|
||
print(f"[metadata] Backfill error: {e}")
|
||
raise HTTPException(500, f"Backfill failed: {str(e)}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import uvicorn
|
||
uvicorn.run(app, host=os.environ.get("HOST", "0.0.0.0"),
|
||
port=int(os.environ.get("PORT", "8500")))
|