This commit is contained in:
mike
2026-06-28 04:51:24 +02:00
parent d522b2a267
commit 256a3b6ac8
4 changed files with 243 additions and 59 deletions

View File

@@ -791,10 +791,18 @@ def _faceswap_worker(job_id: str, model_filename: str, video_name: str, enhance:
# 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)
dir_part = os.path.dirname(model_filename)
basename = os.path.basename(model_filename)
clean_basename = naming.get_base_name(basename)
prev_tag = f"_prev{int(scale*100)}" if scale < 1.0 else ""
tmp_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_basename = f"{ts}_fs_tmp_{vid_stem}_{clean_basename}{prev_tag}.mp4"
out_basename = f"{ts}_fs_{vid_stem}_{clean_basename}{prev_tag}.mp4"
if dir_part:
tmp_name = f"{dir_part}/{tmp_basename}"
out_name = f"{dir_part}/{out_basename}"
else:
tmp_name = tmp_basename
out_name = out_basename
tmp_path = os.path.join(output_dir, tmp_name)
out_path = os.path.join(output_dir, out_name)
@@ -862,7 +870,7 @@ def _faceswap_worker(job_id: str, model_filename: str, video_name: str, enhance:
# 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))
group_id = (person[1] if person and person[1] else naming.get_base_name(os.path.basename(model_filename)))
database.upsert_person(
out_name,
filepath=out_path,
@@ -915,10 +923,16 @@ def _faceswap_worker_ff(job_id: str, model_filename: str, video_name: str,
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)
dir_part = os.path.dirname(model_filename)
basename = os.path.basename(model_filename)
clean_basename = naming.get_base_name(basename)
scale = max(0.1, min(1.0, preview_scale))
prev_tag = f'_prev{int(scale*100)}' if scale < 1.0 else ''
out_name = f'{ts}_fs_{vid_stem}_{base_name}{prev_tag}.mp4'
out_basename = f'{ts}_fs_{vid_stem}_{clean_basename}{prev_tag}.mp4'
if dir_part:
out_name = f'{dir_part}/{out_basename}'
else:
out_name = out_basename
out_path = os.path.join(output_dir, out_name)
processors = ['face_swapper']
@@ -1017,7 +1031,7 @@ def _faceswap_worker_ff(job_id: str, model_filename: str, video_name: str,
_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))
group_id = (person[1] if person and person[1] else naming.get_base_name(os.path.basename(model_filename)))
database.upsert_person(
out_name, filepath=out_path, group_id=group_id,
content_type='video', faceswap_source_video=video_name,
@@ -1531,11 +1545,11 @@ def _batch_worker(job_id: str, filenames: list, prompts: list[str], poses: list,
if person and person[1]:
actual_gid = person[1]
else:
actual_gid = group_id or naming.get_base_name(fname)
actual_gid = group_id or naming.get_base_name(os.path.basename(fname))
database.upsert_person(fname, group_id=actual_gid)
except Exception as e:
print(f"Error determining/updating group for {fname}: {e}")
actual_gid = group_id or naming.get_base_name(fname)
actual_gid = group_id or naming.get_base_name(os.path.basename(fname))
fpath = os.path.join(output_dir, fname)
if not os.path.exists(fpath):
@@ -1586,8 +1600,14 @@ def _batch_worker(job_id: str, filenames: list, prompts: list[str], poses: list,
extra_imgs = [pose_guide_pil] if pose_guide_pil else None
png = _run_pipeline(pil, actual_prompt, seed, max_area, extra_images=extra_imgs)
ts = time.strftime("%Y%m%d_%H%M%S")
clean_fname = naming.get_base_name(fname)
out_name = f"{ts}_{clean_fname}"
dir_part = os.path.dirname(fname)
basename = os.path.basename(fname)
clean_basename = naming.get_base_name(basename)
new_basename = f"{ts}_{clean_basename}"
if dir_part:
out_name = f"{dir_part}/{new_basename}"
else:
out_name = new_basename
out_path = os.path.join(output_dir, out_name)
with open(out_path, "wb") as f:
@@ -1683,8 +1703,14 @@ def _multi_ref_worker(job_id: str, filenames: list[str], prompts: list[str], pos
png = _run_pipeline(work_pil, actual_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}"
dir_part = os.path.dirname(primary_fname)
basename = os.path.basename(primary_fname)
clean_basename = naming.get_base_name(basename)
new_basename = f"{ts}_mr_{clean_basename}"
if dir_part:
out_name = f"{dir_part}/{new_basename}"
else:
out_name = new_basename
out_path = os.path.join(output_dir, out_name)
with open(out_path, "wb") as f:
@@ -2112,11 +2138,16 @@ def generate_video(req: VideoStitchRequest):
paths = paths + list(reversed(paths[1:-1])) if len(paths) > 2 else paths + list(reversed(paths))
ts = time.strftime("%Y%m%d_%H%M%S")
base = naming.get_base_name(req.filenames[0])
base = naming.get_base_name(os.path.basename(req.filenames[0]))
stem = os.path.splitext(base)[0]
out_name = req.output_name or f"{ts}_vid_{stem}.mp4"
if not out_name.lower().endswith(".mp4"):
out_name += ".mp4"
out_basename = req.output_name or f"{ts}_vid_{stem}.mp4"
if not out_basename.lower().endswith(".mp4"):
out_basename += ".mp4"
dir_part = os.path.dirname(req.filenames[0])
if dir_part:
out_name = f"{dir_part}/{out_basename}"
else:
out_name = out_basename
out_path = os.path.join(output_dir, out_name)
# Build ffmpeg concat list in a temp file
@@ -2151,7 +2182,7 @@ def generate_video(req: VideoStitchRequest):
# Register in DB so it shows up in the gallery
person = database.get_person(req.filenames[0])
group_id = person[1] if person and person[1] else naming.get_base_name(req.filenames[0])
group_id = person[1] if person and person[1] else naming.get_base_name(os.path.basename(req.filenames[0]))
try:
next_order = database.get_next_sort_order(group_id)
database.upsert_person(
@@ -2456,13 +2487,13 @@ def db_repair(req: RepairRequest):
if not os.path.exists(fpath):
raise HTTPException(404, "File not found on disk")
# Basic import: just register it in DB
database.upsert_person(req.filename, filepath=fpath, group_id=naming.get_base_name(req.filename), is_source=True)
database.upsert_person(req.filename, filepath=fpath, group_id=naming.get_base_name(os.path.basename(req.filename)), is_source=True)
_invalidate_static()
_run_consistency_check()
return {"status": "imported", "filename": req.filename}
elif req.action == "assign_group":
# Assign a group_id (default to base name if not provided)
gid = req.group_id or naming.get_base_name(req.filename)
gid = req.group_id or naming.get_base_name(os.path.basename(req.filename))
database.upsert_person(req.filename, group_id=gid)
_invalidate_static()
_run_consistency_check()
@@ -2766,7 +2797,7 @@ def set_image_preferred(filename: str):
group_id = person[1]
if not group_id:
# Auto-assign group_id if missing (legacy/orphan)
group_id = naming.get_base_name(filename)
group_id = naming.get_base_name(os.path.basename(filename))
database.upsert_person(filename, group_id=group_id)
print(f"[fix] Auto-assigned group_id={group_id} to orphan {filename} during set-preferred")
@@ -3404,16 +3435,21 @@ def _scenery_worker(job_id: str, model_filename: str, scene_pil: Image.Image,
extra_images=extra_images,
)
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}"
dir_part = os.path.dirname(model_filename)
basename = os.path.basename(model_filename)
clean_basename = naming.get_base_name(basename)
if not clean_basename.lower().endswith('.png'):
clean_basename = os.path.splitext(clean_basename)[0] + '.png'
new_basename = f"{ts}_sc_{clean_basename}"
if dir_part:
out_name = f"{dir_part}/{new_basename}"
else:
out_name = new_basename
out_path = os.path.join(output_dir, out_name)
with open(out_path, "wb") as f:
f.write(png_bytes)
person = database.get_person(model_filename)
group_id = person[1] if person and person[1] else naming.get_base_name(model_filename)
group_id = person[1] if person and person[1] else naming.get_base_name(os.path.basename(model_filename))
try:
embedding = embeddings.generate_embedding(out_path)
next_order = database.get_next_sort_order(group_id)
@@ -3830,9 +3866,16 @@ def remove_background_sam(filename: str):
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)
dir_part = os.path.dirname(filename)
basename = os.path.basename(filename)
stem = os.path.splitext(basename)[0]
nobg_basename = f"{stem}.nobg.png"
if dir_part:
nobg_filename = f"{dir_part}/{nobg_basename}"
else:
nobg_filename = nobg_basename
nobg_path = os.path.join(output_dir, nobg_basename)
with open(path, "rb") as f:
png_bytes = f.read()
@@ -3904,11 +3947,18 @@ def manual_crop_image(filename: str, req: CropRequest):
# Mirror duplicate_image: copy file + register a DB row that points back to the original.
from datetime import datetime as _dt
output_dir = os.path.dirname(src_path)
ext = os.path.splitext(filename)[1] or ".png"
stem = os.path.splitext(filename)[0]
dir_part = os.path.dirname(filename)
basename = os.path.basename(filename)
stem, ext = os.path.splitext(basename)
if not ext:
ext = ".png"
ts = _dt.now().strftime("%Y%m%d_%H%M%S")
new_filename = f"{ts}_crop_{stem}{ext}"
path = os.path.join(output_dir, new_filename)
new_basename = f"{ts}_crop_{stem}{ext}"
if dir_part:
new_filename = f"{dir_part}/{new_basename}"
else:
new_filename = new_basename
path = os.path.join(output_dir, new_basename)
shutil.copy2(src_path, path)
database.upsert_person(
new_filename, filepath=path, group_id=person[1],
@@ -4015,11 +4065,18 @@ def pad_image(filename: str, req: PadRequest):
if req.as_copy:
from datetime import datetime as _dt
output_dir = os.path.dirname(src_path)
ext = os.path.splitext(filename)[1] or ".png"
stem = os.path.splitext(filename)[0]
dir_part = os.path.dirname(filename)
basename = os.path.basename(filename)
stem, ext = os.path.splitext(basename)
if not ext:
ext = ".png"
ts = _dt.now().strftime("%Y%m%d_%H%M%S")
new_filename = f"{ts}_pad_{stem}{ext}"
path = os.path.join(output_dir, new_filename)
new_basename = f"{ts}_pad_{stem}{ext}"
if dir_part:
new_filename = f"{dir_part}/{new_basename}"
else:
new_filename = new_basename
path = os.path.join(output_dir, new_basename)
shutil.copy2(src_path, path)
database.upsert_person(
new_filename, filepath=path, group_id=person[1],
@@ -4104,11 +4161,20 @@ def duplicate_image(filename: str):
raise HTTPException(404, "Image file not found")
path = person[5]
output_dir = os.path.dirname(path)
ext = os.path.splitext(filename)[1] or ".png"
dir_part = os.path.dirname(filename)
basename = os.path.basename(filename)
stem, ext = os.path.splitext(basename)
if not ext:
ext = ".png"
ts = _dt.now().strftime("%Y%m%d_%H%M%S")
stem = os.path.splitext(filename)[0]
new_filename = f"{ts}_dup_{stem}{ext}"
new_path = os.path.join(output_dir, new_filename)
new_basename = f"{ts}_dup_{stem}{ext}"
if dir_part:
new_filename = f"{dir_part}/{new_basename}"
else:
new_filename = new_basename
new_path = os.path.join(output_dir, new_basename)
_shutil.copy2(path, new_path)
group_id = person[1]
# person tuple: (name, group_id, tags, embedding, clip_description, filepath,
@@ -4828,16 +4894,22 @@ def _pose_gen_worker(job_id: str, model_filename: str, prompt: str, seed: int,
png_bytes = _run_pipeline(model_pil, prompt, seed, MAX_AREA, extra_images=extra_images)
ts = time.strftime("%Y%m%d_%H%M%S")
base_name = naming.get_base_name(model_filename)
if not base_name.lower().endswith(".png"):
base_name = os.path.splitext(base_name)[0] + ".png"
out_name = f"{ts}_pose_{base_name}"
dir_part = os.path.dirname(model_filename)
basename = os.path.basename(model_filename)
clean_basename = naming.get_base_name(basename)
if not clean_basename.lower().endswith(".png"):
clean_basename = os.path.splitext(clean_basename)[0] + ".png"
new_basename = f"{ts}_pose_{clean_basename}"
if dir_part:
out_name = f"{dir_part}/{new_basename}"
else:
out_name = new_basename
out_path = os.path.join(output_dir, out_name)
with open(out_path, "wb") as f:
f.write(png_bytes)
person = database.get_person(model_filename)
group_id = person[1] if person and person[1] else naming.get_base_name(model_filename)
group_id = person[1] if person and person[1] else naming.get_base_name(os.path.basename(model_filename))
embedding = embeddings.generate_embedding(out_path)
next_order = database.get_next_sort_order(group_id)