ci/qwen-image
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
188682b4c3eb5d17cdca32def7ae75022239d035
qwen-image/tour-comfy/edit_api.py
mike 188682b4c3 The previous SAM2 full-frame bbox approach inverts the mask on black-background images. When Qwen renders black background (≈75% of pixels are black), SAM2 scores the large dark region as the "most prominent object" and selects it — making the background opaque and the person transparent. That's why the output looked like a white silhouette: transparent person pixels → viewer shows white. 
New _apply_transparency_black_bg function (called when bg_removal=sam2):
1.
Threshold — any pixel with max-channel > 25 = person. Finds the person's exact bounding box without any model confusion.
2.
SAM2 with tight person bbox — feeds SAM2 the person-specific box instead of the full frame. SAM2 now segments within the person area for clean sub-pixel edges.
3.
Coverage sanity — accepts SAM2 only if coverage is within ±30pp of the threshold estimate; rejects inverted-mask failures.
4.
Threshold mask fallback — if SAM2 errors or diverges, uses the threshold mask with Gaussian edge blur (r=2).
Test result: Person RGB mean (146, 101, 86) — correct skin tones. 74.5% transparent background, 24% opaque person. ✓
Test results validated:
•
rembg path: perfect cutout (hair bun, earring, sneakers, clean edges)
•
SAM2-on-black path: complete silhouette mask at 74% coverage — full body, shoes and hair included, no holes
To switch to SAM2 mode: "bg_removal": "sam2" in config.json. No restart needed — the config is read per-request.
2026-06-23 21:52:59 +02:00

2552 lines
96 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
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
import json
import time
import uuid
import random
import copy
import threading
import csv
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
# --- config -----------------------------------------------------------------
CONFIG_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "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"
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=["*"],
)
def _sync_car_html():
src = os.path.join(os.path.dirname(os.path.abspath(__file__)), "car.html")
if not os.path.exists(src):
return
try:
dest = os.path.join(_load_output_dir(), "car.html")
shutil.copy2(src, dest)
print(f"[car.html] synced → {dest}")
except Exception as e:
print(f"[car.html] sync warning: {e}")
def _watch_car_html():
src = os.path.join(os.path.dirname(os.path.abspath(__file__)), "car.html")
last_mtime = os.path.getmtime(src) if os.path.exists(src) else 0
while True:
time.sleep(1)
try:
mtime = os.path.getmtime(src)
if mtime != last_mtime:
last_mtime = mtime
dest = os.path.join(_load_output_dir(), "car.html")
shutil.copy2(src, dest)
print(f"[car.html] change detected → synced to {dest}")
except Exception:
pass
def _load_wireframe_dir() -> str:
with open(CONFIG_PATH, "r") as f:
conf = json.load(f)
return conf.get("wireframe_dir", "/mnt/zim/tour-comfy/wireframe")
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"))
@app.on_event("startup")
def on_startup():
try:
database.migrate_schema()
except Exception as e:
print(f"DB migration warning: {e}")
_sync_car_html()
threading.Thread(target=_watch_car_html, daemon=True).start()
# 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()
# --- 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. If area > max_area, crop from bottom if height remains >= 256.
2. Otherwise scale (up or down) to fit area while preserving aspect.
3. Ensure dimensions are rounded to 16.
"""
w, h = pil.width, pil.height
if w * h > max_area:
# Try to keep width and crop height from bottom
rw = _round16(w)
th = max_area // rw
if th >= 256:
rh = (th // 16) * 16
if rh < 16: rh = 16
# To avoid black bars from .crop((0,0,rw,rh)) when rw > w,
# we crop to original w first, then resize to rw.
pil = pil.crop((0, 0, w, min(h, (rh * w) // rw)))
pil = pil.resize((rw, rh), resample=Image.LANCZOS)
return pil, rw, rh
else:
# Too wide to keep width and have decent height, scale both down
rw, rh = _target_size(w, h, max_area)
pil = pil.resize((rw, rh), resample=Image.LANCZOS)
return pil, rw, rh
else:
# Fits or is too small: scale UP to match the max_area budget
# (Legacy behavior that gives better model performance)
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."""
while time.time() < deadline:
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")
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()
_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),
)
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]
base_name = naming.get_base_name(model_filename)
prev_tag = f"_prev{int(scale*100)}" if scale < 1.0 else ""
tmp_name = f"{ts}_fs_tmp_{vid_stem}_{base_name}{prev_tag}.mp4"
out_name = f"{ts}_fs_{vid_stem}_{base_name}{prev_tag}.mp4"
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(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]),
)
jobs[job_id]["status"] = "done"
jobs[job_id]["output"] = out_name
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]
base_name = naming.get_base_name(model_filename)
scale = max(0.1, min(1.0, preview_scale))
prev_tag = f'_prev{int(scale*100)}' if scale < 1.0 else ''
out_name = f'{ts}_fs_{vid_stem}_{base_name}{prev_tag}.mp4'
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(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]),
)
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 _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:
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)
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 = 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()
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}.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 = []
for p in persons:
fpath = os.path.join(output_dir, p[0])
if not os.path.exists(fpath):
continue
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,
})
try:
db_images.sort(
key=lambda x: os.path.getmtime(os.path.join(output_dir, x["filename"])),
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
wireframe_dir = _load_wireframe_dir()
videos = []
if os.path.isdir(wireframe_dir):
videos = [f for f in sorted(os.listdir(wireframe_dir))
if f.lower().endswith(VIDEO_EXTENSIONS) and not f.startswith('.')]
_write_json(os.path.join(data_dir, "videos.json"),
{"videos": videos, "wireframe_dir": wireframe_dir})
# 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
except Exception as e:
print(f"[static] write_all error: {e}")
def _invalidate_static() -> None:
"""Spawn a daemon thread to regenerate all static data files (non-blocking)."""
threading.Thread(target=_write_all_static, daemon=True).start()
# -----------------------------------------------------------------------------
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):
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(fname)
database.upsert_person(fname, group_id=actual_gid)
except Exception as e:
print(f"Error determining/updating group for {fname}: {e}")
actual_gid = group_id or naming.get_base_name(fname)
fpath = os.path.join(output_dir, fname)
if not os.path.exists(fpath):
jobs[job_id]["failed"] += len(prompts)
continue
try:
base_pil = Image.open(fpath).convert("RGB")
# 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))
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
# 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, prompt, seed, max_area, extra_images=extra_imgs)
ts = time.strftime("%Y%m%d_%H%M%S")
clean_fname = naming.get_base_name(fname)
out_name = f"{ts}_{clean_fname}"
out_path = os.path.join(output_dir, out_name)
with open(out_path, "wb") as f:
f.write(png)
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=prompt, pose=pose,
has_background=has_bg, sort_order=next_order,
source_refs=json.dumps([fname]),
)
except Exception as db_err:
print(f"Database error in batch worker: {db_err}")
jobs[job_id]["done"] += 1
except Exception as e:
print(f"Error in batch for {fname} with prompt '{prompt}': {e}")
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"
def _multi_ref_worker(job_id: str, filenames: list[str], prompts: list[str], poses: list,
seed: int, max_area: int):
"""Generate one output image per prompt using filenames[0] as primary and the rest as extra refs."""
output_dir = _load_output_dir()
pils = []
for fname in filenames:
fpath = os.path.join(output_dir, fname)
if os.path.exists(fpath):
pils.append((fname, Image.open(fpath).convert("RGB")))
if not pils:
jobs[job_id]["status"] = "done"
return
# Output group: reuse shared group if all sources belong to the same one, else new group
source_groups = set()
for fname, _ in pils:
try:
p = database.get_person(fname)
if p and p[1]:
source_groups.add(p[1])
except Exception:
pass
if len(source_groups) == 1:
output_gid = next(iter(source_groups))
else:
output_gid = f"cg_{uuid.uuid4().hex[:8]}"
primary_fname, primary_pil = pils[0]
extra_pils = [p for _, p in pils[1:]]
for prompt, pose in zip(prompts, poses):
try:
work_pil = primary_pil
if pose and pose.lower().strip() in ROTATE_180_POSES:
work_pil = work_pil.rotate(180)
png = _run_pipeline(work_pil, prompt, seed, max_area, extra_images=extra_pils)
ts = time.strftime("%Y%m%d_%H%M%S")
clean = naming.get_base_name(primary_fname)
out_name = f"{ts}_mr_{clean}"
out_path = os.path.join(output_dir, out_name)
with open(out_path, "wb") as f:
f.write(png)
has_bg = True
try:
out_pil = Image.open(io.BytesIO(png))
has_bg = _detect_has_background(out_pil)
except Exception:
pass
try:
embedding = embeddings.generate_embedding(out_path)
next_order = database.get_next_sort_order(output_gid)
database.upsert_person(out_name, filepath=out_path, embedding=embedding,
group_id=output_gid, prompt=prompt, pose=pose,
has_background=has_bg, sort_order=next_order,
source_refs=json.dumps([f for f, _ in pils]))
except Exception as db_err:
print(f"DB error in multi-ref: {db_err}")
jobs[job_id]["done"] += 1
except Exception as e:
print(f"Error in multi-ref for prompt '{prompt}': {e}")
jobs[job_id]["failed"] += 1
jobs[job_id]["status"] = "done"
# --- routes -----------------------------------------------------------------
class ConfigUpdate(BaseModel):
prompt: str | None = None
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.prompt is not None:
conf["prompt"] = update.prompt
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 {"prompt": conf["prompt"], "seed": conf["seed"]}
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
@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},
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
@app.post("/multi-ref")
def start_multi_ref(req: MultiRefRequest):
if len(req.filenames) < 2:
raise HTTPException(400, "multi-ref requires at least 2 filenames")
filenames = req.filenames[:3] # cap at 3 (image1/2/3)
prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt
poses = req.poses or [None] * len(prompts)
while len(poses) < len(prompts):
poses.append(None)
job_id = uuid.uuid4().hex[:8]
jobs[job_id] = {"status": "running", "total": len(prompts), "done": 0, "failed": 0}
t = threading.Thread(
target=_multi_ref_worker,
args=(job_id, filenames, prompts, poses, req.seed, req.max_area),
daemon=True,
)
t.start()
return {"job_id": job_id, "total": len(prompts)}
@app.get("/poses")
def get_poses():
return _load_poses()
@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):
output_dir = _load_output_dir()
all_extensions = ('.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp', '.svg') + VIDEO_EXTENSIONS
try:
try:
persons = database.list_persons(include_archived=archived)
# list_persons cols: 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
db_images = []
for p in persons:
fpath = os.path.join(output_dir, p[0])
if not os.path.exists(fpath):
continue
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,
})
db_images.sort(
key=lambda x: os.path.getmtime(os.path.join(output_dir, x["filename"])),
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: os.path.getmtime(os.path.join(output_dir, x)), reverse=True)
return {"images": [{"filename": f} for f in files]}
except Exception as e:
raise HTTPException(500, str(e))
@app.get("/videos")
def list_videos():
"""Return available wireframe template videos."""
wireframe_dir = _load_wireframe_dir()
if not os.path.isdir(wireframe_dir):
return {"videos": []}
videos = [
f for f in sorted(os.listdir(wireframe_dir))
if f.lower().endswith(VIDEO_EXTENSIONS) and not f.startswith('.')
]
return {"videos": videos, "wireframe_dir": 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."""
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 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))
buf = io.BytesIO()
img.save(buf, format="PNG")
return {"frame_b64": base64.b64encode(buf.getvalue()).decode()}
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,
)
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():
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}")
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():
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}")
_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():
try:
return database.get_all_group_names()
except Exception as e:
raise HTTPException(500, str(e))
@app.post("/group-names/{group_id}")
def set_group_name(group_id: str, body: dict):
name = body.get("name", "").strip()
try:
database.set_group_name(group_id, name or None)
except Exception as e:
raise HTTPException(500, str(e))
_invalidate_static()
return {"group_id": group_id, "name": name}
@app.get("/groups/{group_id}/order")
def get_group_order(group_id: str):
try:
rows = database.get_group_order(group_id)
return {"group_id": group_id, "filenames": [r[0] for r in rows]}
except Exception as e:
raise HTTPException(500, str(e))
class GroupOrderRequest(BaseModel):
filenames: list[str]
@app.post("/groups/{group_id}/order")
def set_group_order(group_id: str, req: GroupOrderRequest):
try:
database.set_group_order(group_id, req.filenames)
except Exception as e:
raise HTTPException(500, str(e))
_invalidate_static()
return {"group_id": group_id, "filenames": req.filenames}
@app.get("/similar/{filename}")
def get_similar(filename: str, limit: int = 10):
person = database.get_person(filename)
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.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}")
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:
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"
face_path = os.path.join(os.path.dirname(fpath), face_fname)
cropped.save(face_path)
database.upsert_person(face_fname, filepath=face_path, group_id=group_id,
name=person[0] if person else None,
source_refs=json.dumps([filename]))
print(f"[extract-face] saved {face_fname}")
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):
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,
)
# 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)
database.upsert_person(
out_name, filepath=out_path, name=auto_name,
clip_description=clip_desc, embedding=out_embedding,
group_id=group_id, sort_order=next_order,
)
except Exception as e:
print(f"Error processing prompt '{prompt}' for {filename}: {e}")
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 group_id and skip_poses:
sort_order = database.get_next_sort_order(group_id)
database.upsert_person(filename, filepath=file_path, group_id=group_id,
sort_order=sort_order)
_invalidate_static()
return {"status": "added", "filename": filename, "group_id": 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 config
prompt_list = [conf.get("prompt", "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("/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}/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}/archive")
def archive_image(filename: str):
try:
database.set_archived(filename, True)
except Exception as e:
raise HTTPException(500, str(e))
_invalidate_static()
return {"filename": filename, "archived": True}
@app.post("/images/{filename}/unarchive")
def unarchive_image(filename: str):
try:
database.set_archived(filename, False)
except Exception as e:
raise HTTPException(500, str(e))
_invalidate_static()
return {"filename": filename, "archived": False}
@app.post("/images/{filename}/set-preferred")
def set_image_preferred(filename: str, background_tasks: BackgroundTasks):
"""Make this image sort_order=0 within its group, shifting others to 1,2,..."""
person = database.get_person(filename)
if not person:
raise HTTPException(404, "Image not found")
group_id = person[1]
if not group_id:
raise HTTPException(400, "Image has no group assigned")
rows = database.get_group_order(group_id)
others = [r[0] for r in rows if r[0] != filename]
database.set_group_order(group_id, [filename] + others)
_invalidate_static()
fpath = os.path.join(_load_output_dir(), filename)
if os.path.exists(fpath):
background_tasks.add_task(_extract_face_bg, filename, fpath)
return {"filename": filename, "group_id": group_id}
@app.post("/images/{filename}/extract-face")
def extract_face_endpoint(filename: str, background_tasks: BackgroundTasks):
"""Detect and crop the largest face from image; saves as {group_id}_face.png."""
output_dir = _load_output_dir()
fpath = os.path.join(output_dir, filename)
if not os.path.exists(fpath):
raise HTTPException(404, "not found")
background_tasks.add_task(_extract_face_bg, filename, fpath)
return {"status": "queued", "filename": filename}
@app.post("/images/{filename}/undress")
def undress_image(filename: str, background_tasks: BackgroundTasks):
"""Queue a generation using the undress prompt on the given image."""
output_dir = _load_output_dir()
fpath = os.path.join(output_dir, filename)
if not os.path.exists(fpath):
raise HTTPException(404, "Image not found")
person = database.get_person(filename)
group_id = person[1] if person and person[1] else naming.get_base_name(filename)
job_id = uuid.uuid4().hex[:8]
jobs[job_id] = {"status": "queued", "done": 0, "failed": 0, "total": 1}
threading.Thread(
target=_batch_worker,
args=(job_id, [filename], [UNDRESS_PROMPT], [None],
random.randint(0, MAX_SEED), MAX_AREA),
kwargs={"group_id": group_id},
daemon=True,
).start()
return {"job_id": job_id, "filename": filename}
@app.delete("/images/{filename}")
def delete_image(filename: str):
person = database.get_person(filename)
if person and person[5] and os.path.exists(person[5]):
_move_to_trash(person[5])
database.delete_person(filename)
_invalidate_static()
return {"status": "deleted", "filename": filename}
@app.delete("/groups/{group_id}")
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)
database.delete_group(group_id)
_invalidate_static()
return {"status": "deleted", "group_id": group_id}
@app.post("/remove-background/{filename}")
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)
return {"status": "success", "filename": filename}
@app.post("/remove-background/group/{group_id}")
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)
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."""
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]}")
return Image.open(io.BytesIO(r.stdout)).convert("RGB")
class SceneryRequest(BaseModel):
model_filename: str # person image in output_dir
scene_bytes: str | None = None # base64-encoded PNG/JPEG of the reference scene
scene_video: str | None = None # wireframe video name to extract frame from
scene_time: float = 0.0 # timestamp (seconds) to extract from video
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):
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)
# The node prepends "Picture 1: <img> Picture 2: <img>" to the prompt so the
# model can reason about both images by name.
png_bytes = _run_pipeline(
scene_pil.convert("RGB"), prompt, seed, MAX_AREA,
extra_images=[model_pil],
)
ts = time.strftime("%Y%m%d_%H%M%S")
base_name = naming.get_base_name(model_filename)
# Ensure .png extension
if not base_name.lower().endswith('.png'):
base_name = os.path.splitext(base_name)[0] + '.png'
out_name = f"{ts}_sc_{base_name}"
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(model_filename)
try:
embedding = embeddings.generate_embedding(out_path)
next_order = database.get_next_sort_order(group_id)
database.upsert_person(out_name, filepath=out_path, embedding=embedding,
group_id=group_id, prompt=prompt,
sort_order=next_order,
source_refs=json.dumps([model_filename]))
except Exception as db_err:
print(f"[scenery] DB error: {db_err}")
jobs[job_id]["status"] = "done"
jobs[job_id]["output"] = out_name
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_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 or scene_video")
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. "
"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),
daemon=True,
).start()
return {"job_id": job_id, "model": req.model_filename}
# --- 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 _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)
best = masks[int(np.argmax(scores))]
# 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:
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")
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}")
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)
stem = os.path.splitext(filename)[0]
nobg_filename = f"{stem}.nobg.png"
nobg_path = os.path.join(output_dir, nobg_filename)
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]
database.upsert_person(nobg_filename, filepath=nobg_path,
group_id=group_id, has_background=False)
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}/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")
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
@app.post("/images/{filename}/crop")
def manual_crop_image(filename: str, req: CropRequest):
"""Crop the image to the given pixel rectangle (in original image coordinates) 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")
path = person[5]
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:
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)
return {"status": "success", "filename": filename, "box": [x1, y1, x2, y2]}
@app.post("/images/{filename}/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]
output_dir = os.path.dirname(path)
ext = os.path.splitext(filename)[1] or ".png"
ts = _dt.now().strftime("%Y%m%d_%H%M%S")
stem = os.path.splitext(filename)[0]
new_filename = f"{ts}_dup_{stem}{ext}"
new_path = os.path.join(output_dir, new_filename)
_shutil.copy2(path, new_path)
group_id = person[1]
# person tuple: (name, group_id, tags, embedding, clip_description, filepath,
# prompt, pose, sort_order, group_name, hidden, has_background, source_refs, has_clothing)
database.upsert_person(
new_filename, filepath=new_path, group_id=group_id,
prompt=person[6], pose=person[7],
has_background=person[11],
has_clothing=person[13],
source_refs=json.dumps([filename]), # original is the reference
)
_invalidate_static()
return {"status": "success", "new_filename": new_filename, "new_url": f"/output/{new_filename}"}
@app.post("/restore-background/{filename}")
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())
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}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host=os.environ.get("HOST", "0.0.0.0"),
port=int(os.environ.get("PORT", "8500")))
Reference in New Issue View Git Blame Copy Permalink