150ef6dab0
• Implemented global offline mode for all HuggingFace-dependent modules to eliminate unauthenticated request warnings and prevent external connection attempts during startup and runtime. Changes • Global Offline Configuration: Added HF_HUB_OFFLINE=1 and TRANSFORMERS_OFFLINE=1 to the environment variables at the top of edit_api.py, embeddings.py, watcher.py, and pose_llm/pose_llm_api.py. • Explicit Local Files Only: Updated all huggingface_hub.hf_hub_download calls in edit_api.py to include local_files_only=True, ensuring they never attempt to reach the Hub if models are already cached. • LLM Service Optimization: Updated AutoTokenizer and AutoModelForCausalLM calls in pose_llm/pose_llm_api.py to use local_files_only=True. • Consistency: Ensured that shared modules like embeddings.py which uses open_clip are also locked to offline mode to prevent background downloads.
6263 lines
238 KiB
Python
6263 lines
238 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. The MI50 is not fast, so we
|
||
# cap at ~1MP by default; raise via MAX_AREA env if you want bigger output.
|
||
MAX_AREA = int(os.environ.get("MAX_AREA", str(1024 * 1024)))
|
||
GEN_TIMEOUT = int(os.environ.get("GEN_TIMEOUT", "600")) # seconds per request
|
||
|
||
# Node ids in workflow_qwen_edit.json (kept stable on purpose).
|
||
NODE_LOADIMAGE = "4"
|
||
NODE_POSITIVE = "5"
|
||
NODE_NEGATIVE = "6"
|
||
NODE_LATENT = "7"
|
||
NODE_KSAMPLER = "8"
|
||
NODE_SAVE = "10"
|
||
|
||
MAX_SEED = 2**32 - 1
|
||
|
||
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) ---
|
||
REFINEMENT_SYSTEM = (
|
||
"You write explicit photographic — black void sceneries — restraining 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), "
|
||
"ending with a realism tag like 'Perfect anatomy, realistic' or 'Anatomically precise, hyperrealistic, keep the characteristics of the reference image'. "
|
||
"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)}°"
|
||
)
|
||
_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()
|
||
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}")
|
||
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
|
||
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:
|
||
result = swapper.get(result, best_face, src_face, paste_back=True)
|
||
except Exception:
|
||
pass
|
||
if gfpgan_restorer is not None:
|
||
try:
|
||
_, _, 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
|
||
|
||
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],
|
||
})
|
||
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()
|
||
|
||
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:
|
||
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:
|
||
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 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
|
||
|
||
|
||
@app.post("/refine-prompt")
|
||
def refine_prompt(req: RefineRequest):
|
||
"""Refine a prompt using the external uncensored chat LLM."""
|
||
if not req.prompt:
|
||
raise HTTPException(400, "Prompt is required")
|
||
|
||
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}")
|
||
|
||
user_content = f"Refine this pose: {req.prompt}"
|
||
if context_str:
|
||
user_content += f"\n\nUse the following image context details to ensure the refined prompt 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()
|
||
return {"refined": refined}
|
||
except Exception as e:
|
||
print(f"Refinement error: {e}")
|
||
raise HTTPException(500, f"LLM refinement failed: {str(e)}")
|
||
|
||
|
||
class DesignerGenerateRequest(BaseModel):
|
||
n: int = 3
|
||
context: str | None = None
|
||
filename: str | None = None
|
||
beta: bool = False
|
||
|
||
|
||
@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 prompt instructions
|
||
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."
|
||
)
|
||
|
||
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_prompt}
|
||
],
|
||
"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
|
||
new_poses = {}
|
||
for name, body in generated.items():
|
||
if not name or not body:
|
||
continue
|
||
if name.lower() in existing_lower or name.lower() in (k.lower() for k in new_poses):
|
||
print(f"[designer] skip duplicate: {name}")
|
||
continue
|
||
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,
|
||
})
|
||
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"]
|
||
})
|
||
|
||
return {
|
||
"videos": videos,
|
||
"groups": groups,
|
||
"standalone_images": standalone_images,
|
||
"wireframe_dir": wireframe_dir
|
||
}
|
||
|
||
|
||
@app.get("/videos")
|
||
def list_videos():
|
||
"""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 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}")
|
||
wireframe_dir = _load_wireframe_dir()
|
||
return _get_grouped_wireframes(wireframe_dir)
|
||
|
||
|
||
@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]}")
|
||
|
||
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():
|
||
output_dir = _load_output_dir()
|
||
static_file = os.path.join(output_dir, "_data", "names.json")
|
||
if 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 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():
|
||
output_dir = _load_output_dir()
|
||
static_file = os.path.join(output_dir, "_data", "groups.json")
|
||
if 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 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():
|
||
output_dir = _load_output_dir()
|
||
static_file = os.path.join(output_dir, "_data", "group-names.json")
|
||
if 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 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("/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
|
||
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")
|
||
|
||
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")
|
||
|
||
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
|
||
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)
|
||
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."
|
||
)
|
||
|
||
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."""
|
||
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:
|
||
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 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 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 simple human-readable description of a COCO-17 pose."""
|
||
vis = [k[2] >= _POSE_MIN_SCORE for k in kpts]
|
||
if sum(vis) < 5: return "Indeterminate pose"
|
||
|
||
parts = []
|
||
|
||
# Vertical orientation
|
||
if vis[0] and vis[11] and vis[12]: # nose and hips
|
||
hip_y = (kpts[11][1] + kpts[12][1]) / 2
|
||
head_y = kpts[0][1]
|
||
if head_y > hip_y + 20: parts.append("upside down")
|
||
elif head_y > hip_y - 20: parts.append("reclining/prone")
|
||
else: parts.append("upright")
|
||
|
||
# Arms
|
||
if vis[9] and vis[10]: # wrists
|
||
sh_y = (kpts[5][1] + kpts[6][1]) / 2 if (vis[5] and vis[6]) else kpts[0][1]
|
||
if kpts[9][1] < sh_y and kpts[10][1] < sh_y: parts.append("arms raised")
|
||
elif kpts[9][1] > sh_y + 100 and kpts[10][1] > sh_y + 100: parts.append("arms down")
|
||
else: parts.append("arms at sides")
|
||
|
||
# Legs
|
||
if vis[15] and vis[16]: # ankles
|
||
dist = abs(kpts[15][0] - kpts[16][0])
|
||
if dist > 150: 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'])}°"
|
||
)
|
||
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) -> bool:
|
||
"""Detect if the person has complete anatomical structure.
|
||
|
||
Returns True if all major body parts are visible (head, torso, arms, legs).
|
||
Uses pose keypoint visibility to determine completeness.
|
||
"""
|
||
if not keypoints or len(keypoints) < 17:
|
||
return False
|
||
|
||
# Minimum visibility threshold for each keypoint
|
||
MIN_VISIBILITY = 0.3
|
||
|
||
# Key keypoints that indicate anatomical completeness
|
||
# Head (0), shoulders (5,6), hips (11,12), elbows (7,8), wrists (9,10), knees (13,14), ankles (15,16)
|
||
keypoint_indices = [0, 5, 6, 11, 12, 7, 8, 9, 10, 13, 14, 15, 16]
|
||
|
||
visible_count = 0
|
||
for idx in keypoint_indices:
|
||
if idx < len(keypoints) and keypoints[idx][2] >= MIN_VISIBILITY:
|
||
visible_count += 1
|
||
|
||
# If more than half of the key keypoints are visible, consider it complete
|
||
return visible_count > len(keypoint_indices) * 0.5
|
||
|
||
|
||
def _detect_facial_direction(keypoints: list) -> str:
|
||
"""Detect the facial direction from keypoints.
|
||
|
||
Returns a string describing the head orientation.
|
||
"""
|
||
if not keypoints or len(keypoints) < 17:
|
||
return "unknown"
|
||
|
||
# Key points for face direction detection
|
||
# Nose (0), left ear (3), right ear (4)
|
||
nose = keypoints[0] if len(keypoints) > 0 and keypoints[0][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"
|
||
|
||
# Determine face direction based on ear positions
|
||
if l_ear and r_ear:
|
||
ear_mid_x = (l_ear[0] + r_ear[0]) / 2
|
||
dx = nose[0] - ear_mid_x
|
||
if dx < -0.05:
|
||
return "looking left"
|
||
elif dx > 0.05:
|
||
return "looking right"
|
||
else:
|
||
return "looking forward"
|
||
elif l_ear and not r_ear:
|
||
return "looking strongly right"
|
||
elif r_ear and not l_ear:
|
||
return "looking strongly left"
|
||
else:
|
||
return "looking forward"
|
||
|
||
|
||
def _detect_objects(pil_img: Image.Image) -> 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 = []
|
||
for t in tags:
|
||
if t["cat"] in (0, 4): # general and character categories
|
||
# For simplicity, we'll return just the tag name with confidence
|
||
# In a more advanced implementation, we could extract bounding boxes from the model
|
||
objects.append({
|
||
"tag": t["tag"],
|
||
"score": t["score"],
|
||
"bbox": None # No bounding box available from WD tagger
|
||
})
|
||
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
|
||
people_count = _detect_people_count(best_person)
|
||
anatomical_completeness = _detect_anatomical_completeness(best_person)
|
||
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)
|
||
|
||
# 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
|
||
|
||
|
||
import asyncio
|
||
from concurrent.futures import ThreadPoolExecutor as _ThreadPoolExecutor
|
||
_metadata_executor = _ThreadPoolExecutor(max_workers=1, thread_name_prefix="metadata")
|
||
|
||
|
||
@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:
|
||
# 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")))
|