57d98d1ed7
hierarchical family and provides improved segmentation capabilities over the basic models.
1971 lines
71 KiB
Python
1971 lines
71 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
|
||
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}")
|
||
|
||
|
||
# --- 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):
|
||
"""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
|
||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
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)
|
||
tmp_name = f"{ts}_fs_tmp_{vid_stem}_{base_name}.mp4"
|
||
out_name = f"{ts}_fs_{vid_stem}_{base_name}.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:
|
||
result = frame.copy()
|
||
for face in tgt_faces:
|
||
try:
|
||
result = swapper.get(result, 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
|
||
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):
|
||
"""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)
|
||
out_name = f'{ts}_fs_{vid_stem}_{base_name}.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',
|
||
'--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',
|
||
'--face-selector-mode', 'many',
|
||
]
|
||
if enhance:
|
||
cmd += ['--face-enhancer-model', 'gfpgan_1.4']
|
||
|
||
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',
|
||
)
|
||
# Read stdout for progress, stderr for error info
|
||
import threading as _thr
|
||
def _drain_stderr():
|
||
for ln in proc.stderr:
|
||
output_lines.append(ln.rstrip())
|
||
print(f'[facefusion] {ln.rstrip()}')
|
||
_thr.Thread(target=_drain_stderr, daemon=True).start()
|
||
for line in proc.stdout:
|
||
print(f'[facefusion] {line.rstrip()}')
|
||
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
|
||
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]
|
||
|
||
# Transparency detection
|
||
is_transparent = any(kw in prompt.lower() for kw in ["transparent", "no background", "remove background", "alpha channel"])
|
||
if is_transparent:
|
||
graph[NODE_NEGATIVE]["inputs"]["prompt"] = "checkerboard, grid, pattern, texture, background details, watermark, deformed anatomy"
|
||
|
||
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 is_transparent:
|
||
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}")
|
||
|
||
|
||
def _batch_worker(job_id: str, filenames: list, prompts: list[str], poses: list,
|
||
seed: int, max_area: int, group_id: str | None = None):
|
||
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")
|
||
for prompt, pose in zip(prompts, poses):
|
||
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)
|
||
|
||
png = _run_pipeline(pil, prompt, seed, max_area)
|
||
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)
|
||
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
|
||
|
||
|
||
@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}
|
||
t = threading.Thread(
|
||
target=_batch_worker,
|
||
args=(job_id, req.filenames, prompts, poses, req.seed, req.max_area, req.group_id),
|
||
daemon=True,
|
||
)
|
||
t.start()
|
||
return {"job_id": job_id, "total": total_tasks}
|
||
|
||
|
||
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():
|
||
output_dir = _load_output_dir()
|
||
all_extensions = ('.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp', '.svg') + VIDEO_EXTENSIONS
|
||
try:
|
||
try:
|
||
persons = database.list_persons()
|
||
# 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
|
||
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],
|
||
})
|
||
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/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)
|
||
|
||
|
||
@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},
|
||
daemon=True,
|
||
)
|
||
else:
|
||
t = threading.Thread(
|
||
target=_faceswap_worker,
|
||
args=(job_id, req.model_filename, req.video_name),
|
||
kwargs={'enhance': req.enhance},
|
||
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}")
|
||
|
||
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}")
|
||
|
||
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}")
|
||
|
||
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))
|
||
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))
|
||
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])
|
||
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 _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}")
|
||
|
||
|
||
@app.post("/upload")
|
||
def upload_image(
|
||
background_tasks: BackgroundTasks,
|
||
image: UploadFile = File(...),
|
||
prompts: str = Form(""),
|
||
name: str = Form(None),
|
||
):
|
||
# 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)
|
||
# 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)
|
||
|
||
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")]
|
||
|
||
group_id = f"up_{uuid.uuid4().hex[:8]}" # unique per upload; avoids collisions when pasting generic filenames
|
||
background_tasks.add_task(_process_upload, file_path, filename, prompt_list, name, group_id)
|
||
|
||
return {"status": "processing", "filename": filename, "group_id": group_id, "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))
|
||
return {"filename": filename, "hidden": hidden}
|
||
|
||
|
||
@app.post("/images/{filename}/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:
|
||
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)
|
||
return {"filename": filename, "group_id": group_id}
|
||
|
||
|
||
@app.post("/images/{filename}/undress")
|
||
def undress_image(filename: str, background_tasks: BackgroundTasks):
|
||
"""Queue a generation using the undress prompt on the given image."""
|
||
output_dir = _load_output_dir()
|
||
fpath = os.path.join(output_dir, filename)
|
||
if not os.path.exists(fpath):
|
||
raise HTTPException(404, "Image not found")
|
||
person = database.get_person(filename)
|
||
group_id = person[1] if person and person[1] else naming.get_base_name(filename)
|
||
job_id = uuid.uuid4().hex[:8]
|
||
jobs[job_id] = {"status": "queued", "done": 0, "failed": 0, "total": 1}
|
||
threading.Thread(
|
||
target=_batch_worker,
|
||
args=(job_id, [filename], [UNDRESS_PROMPT], [None],
|
||
random.randint(0, MAX_SEED), MAX_AREA),
|
||
kwargs={"group_id": group_id},
|
||
daemon=True,
|
||
).start()
|
||
return {"job_id": job_id, "filename": filename}
|
||
|
||
|
||
@app.delete("/images/{filename}")
|
||
def delete_image(filename: str):
|
||
person = database.get_person(filename)
|
||
if person and person[5] and os.path.exists(person[5]):
|
||
_move_to_trash(person[5])
|
||
|
||
database.delete_person(filename)
|
||
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)
|
||
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 AutomaticMaskGenerator. Returns generator 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.automatic_mask_generator import SAM2AutomaticMaskGenerator
|
||
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 = SAM2AutomaticMaskGenerator(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 using SAM2 (largest-area mask = subject), fallback to rembg."""
|
||
predictor = _load_sam2()
|
||
if predictor is False:
|
||
return _apply_transparency(png_bytes)
|
||
try:
|
||
import numpy as np
|
||
img = Image.open(io.BytesIO(png_bytes)).convert("RGB")
|
||
arr = np.array(img)
|
||
masks = predictor.generate(arr)
|
||
if not masks:
|
||
return _apply_transparency(png_bytes)
|
||
best = max(masks, key=lambda m: m["area"])
|
||
mask_np = (best["segmentation"].astype(np.uint8) * 255)
|
||
rgba = img.convert("RGBA")
|
||
r, g, b, _ = rgba.split()
|
||
alpha = Image.fromarray(mask_np, mode="L")
|
||
out = Image.merge("RGBA", (r, g, b, alpha))
|
||
buf = io.BytesIO()
|
||
out.save(buf, format="PNG")
|
||
return buf.getvalue()
|
||
except Exception as e:
|
||
print(f"[sam2] inference error, falling back to rembg: {e}")
|
||
return _apply_transparency(png_bytes)
|
||
|
||
|
||
@app.post("/remove-background-sam/{filename}")
|
||
def remove_background_sam(filename: str):
|
||
"""SAM2-based background removal (RGBA PNG). Falls back to rembg if SAM2 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]
|
||
with open(path, "rb") as f:
|
||
png_bytes = f.read()
|
||
transparent_png = _apply_transparency_sam2(png_bytes)
|
||
with open(path, "wb") as f:
|
||
f.write(transparent_png)
|
||
used_sam2 = _sam2_predictor is not False and _sam2_predictor is not None
|
||
return {"status": "success", "filename": filename, "used_sam2": used_sam2}
|
||
|
||
|
||
@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")))
|