Update app.py

app.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+import re
 import shutil
 import subprocess
 import sys
@@ -9,6 +10,7 @@ import traceback
 import uuid
 import zipfile
 from pathlib import Path
+from dataclasses import dataclass
 
 import gradio as gr
 import spaces
@@ -34,6 +36,26 @@ GPU_DURATION = int(os.getenv("SCAIL_POSE_GPU_DURATION", "90"))
 IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp"}
 VIDEO_EXTS = {".mp4", ".mov", ".webm", ".mkv"}
 
+SCAIL_COLORS = [
+    (0, 0, 255),
+    (255, 0, 0),
+    (0, 255, 0),
+    (255, 255, 0),
+    (255, 0, 255),
+    (0, 255, 255),
+]
+
+_SAM3_VIDEO_PREDICTOR = None
+_SAM3_IMAGE_PREDICTOR = None
+_SAM3_PREDICTOR_KEY = None
+
+
+@dataclass(frozen=True)
+class CharacterRef:
+    character: str
+    view: str
+    path: Path
+
 
 def _repo_status() -> str:
     required = [
@@ -141,6 +163,116 @@ def _copy_video(upload, dest: Path, label: str) -> Path:
     return dest
 
 
+def _copy_reference_to_png(source: Path, dest: Path) -> Path:
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        image = Image.open(source).convert("RGB")
+        image.save(dest)
+    except Exception as exc:
+        raise RuntimeError(f"Could not read reference image: {source}") from exc
+    return dest
+
+
+def _safe_name(value: str) -> str:
+    value = value.strip().replace(" ", "_")
+    value = re.sub(r"[^A-Za-z0-9_.-]+", "_", value)
+    value = value.strip("._-")
+    return value or "view"
+
+
+def _character_sort_key(character: str) -> tuple[int, str]:
+    match = re.fullmatch(r"character_(\d+)", character)
+    if match:
+        return int(match.group(1)), character
+    return 9999, character
+
+
+def _parse_character_ref_name(path: Path) -> tuple[str, str] | None:
+    if path.suffix.lower() not in IMAGE_EXTS:
+        return None
+    if path.stem.endswith("_mask"):
+        return None
+
+    parent = path.parent.name
+    if re.fullmatch(r"character_\d+", parent):
+        return parent, _safe_name(path.stem)
+
+    stem = path.stem
+    match = re.match(r"^(character_\d+)(?:__|--|_)(.+)$", stem)
+    if not match:
+        return None
+    return match.group(1), _safe_name(match.group(2))
+
+
+def _normalize_upload_list(files) -> list[Path]:
+    if not files:
+        return []
+    if not isinstance(files, list):
+        files = [files]
+    return [_as_path(file, "reference file") for file in files]
+
+
+def _safe_extract_zip(zip_path: str | Path, label: str) -> Path:
+    zip_path = _as_path(zip_path, label)
+    if zip_path.suffix.lower() != ".zip":
+        raise RuntimeError(f"{label} must be a .zip file.")
+
+    extract_root = JOB_ROOT / "uploads" / uuid.uuid4().hex
+    extract_root.mkdir(parents=True, exist_ok=True)
+    with zipfile.ZipFile(zip_path) as zf:
+        for item in zf.infolist():
+            item_path = Path(item.filename)
+            if item_path.is_absolute() or ".." in item_path.parts:
+                raise RuntimeError(f"Unsafe path in zip: {item.filename}")
+        zf.extractall(extract_root)
+
+    visible = [p for p in extract_root.iterdir() if p.name not in {".DS_Store", "__MACOSX"}]
+    if len(visible) == 1 and visible[0].is_dir():
+        return visible[0]
+    return extract_root
+
+
+def _collect_character_refs(reference_files, reference_zip) -> list[CharacterRef]:
+    paths = []
+    if reference_zip is not None:
+        root = _safe_extract_zip(reference_zip, "reference zip")
+        paths.extend(p for p in root.rglob("*") if p.is_file())
+    paths.extend(_normalize_upload_list(reference_files))
+
+    refs = []
+    ignored = []
+    for path in paths:
+        parsed = _parse_character_ref_name(path)
+        if parsed is None:
+            if path.suffix.lower() in IMAGE_EXTS:
+                ignored.append(path.name)
+            continue
+        character, view = parsed
+        refs.append(CharacterRef(character=character, view=view, path=path))
+
+    if not refs:
+        raise RuntimeError(
+            "No character references found. Use names like `character_0__front.png`, "
+            "`character_0__back.png`, `character_1__front.png`, or upload a zip with "
+            "`characters/character_0/front.png`."
+        )
+
+    seen = set()
+    duplicates = []
+    for ref in refs:
+        key = (ref.character, ref.view)
+        if key in seen:
+            duplicates.append(f"{ref.character}/{ref.view}")
+        seen.add(key)
+    if duplicates:
+        raise RuntimeError("Duplicate character view(s): " + ", ".join(sorted(duplicates)))
+
+    if ignored:
+        logging.info("Ignored reference image(s) without character naming: %s", ", ".join(ignored))
+
+    return sorted(refs, key=lambda ref: (_character_sort_key(ref.character), ref.view.lower()))
+
+
 def _text_args(text_prompt: str) -> list[str]:
     words = [part.strip() for part in (text_prompt or "").split() if part.strip()]
     return words or ["human", "character"]
@@ -212,6 +344,84 @@ def _require_outputs(job_dir: Path, names: list[str]) -> None:
         raise RuntimeError("SCAIL-Pose did not produce expected output(s): " + ", ".join(missing))
 
 
+def _load_sam3_predictors(sam3_model: Path):
+    global _SAM3_VIDEO_PREDICTOR, _SAM3_IMAGE_PREDICTOR, _SAM3_PREDICTOR_KEY
+    key = str(sam3_model)
+    if _SAM3_VIDEO_PREDICTOR is not None and _SAM3_PREDICTOR_KEY == key:
+        return _SAM3_VIDEO_PREDICTOR, _SAM3_IMAGE_PREDICTOR
+
+    from ultralytics.models.sam import SAM3SemanticPredictor, SAM3VideoSemanticPredictor
+
+    overrides = dict(
+        conf=0.25,
+        task="segment",
+        mode="predict",
+        imgsz=640,
+        model=str(sam3_model),
+        half=True,
+        save=False,
+        verbose=False,
+    )
+    _SAM3_VIDEO_PREDICTOR = SAM3VideoSemanticPredictor(overrides=overrides, new_det_thresh=1.0)
+    _SAM3_IMAGE_PREDICTOR = SAM3SemanticPredictor(overrides=overrides)
+    _SAM3_PREDICTOR_KEY = key
+    return _SAM3_VIDEO_PREDICTOR, _SAM3_IMAGE_PREDICTOR
+
+
+def _write_canonical_pack_zip(pack_root: Path, mode: str) -> Path:
+    zip_path = pack_root.parent / f"scail2_{mode}_advanced_pack.zip"
+    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
+        for path in sorted(pack_root.rglob("*")):
+            if path.is_file():
+                zf.write(path, arcname=path.relative_to(pack_root).as_posix())
+    return zip_path
+
+
+def _advanced_gallery(pack_root: Path):
+    items = []
+    for image_path in sorted((pack_root / "characters").rglob("*")):
+        if image_path.is_file() and image_path.suffix.lower() in IMAGE_EXTS:
+            items.append((str(image_path), image_path.relative_to(pack_root).as_posix()))
+    return items
+
+
+def _build_advanced_metadata(
+    pack_root: Path,
+    mode: str,
+    prompt: str,
+    primary_ref: CharacterRef,
+    character_names: list[str],
+    text_prompt: str,
+) -> None:
+    primary_image = f"characters/{primary_ref.character}/{primary_ref.view}.png"
+    primary_mask = f"characters/{primary_ref.character}/{primary_ref.view}_mask.png"
+    metadata = {
+        "mode": mode,
+        "source": "scail-pose-gradio-advanced-pack-builder",
+        "primary": {
+            "image": primary_image,
+            "mask": primary_mask,
+        },
+        "driving": {
+            "video": "rendered_v2.mp4",
+            "mask_video": "rendered_mask_v2.mp4" if mode == "animation" else "replace_mask.mp4",
+        },
+        "characters": character_names,
+        "sam3_text": _text_args(text_prompt),
+    }
+    (pack_root / "metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8")
+    (pack_root / "prompt.txt").write_text(prompt or "", encoding="utf-8")
+
+
+def _prepare_advanced_refs(refs: list[CharacterRef], pack_root: Path) -> list[CharacterRef]:
+    prepared = []
+    for ref in refs:
+        dest = pack_root / "characters" / ref.character / f"{ref.view}.png"
+        _copy_reference_to_png(ref.path, dest)
+        prepared.append(CharacterRef(character=ref.character, view=ref.view, path=dest))
+    return prepared
+
+
 def _animation_command(
     job_dir: Path,
     sam3_model: Path,
@@ -396,6 +606,133 @@ def build_replacement_pack(
         return None, None, None, None, traceback.format_exc()
 
 
+@spaces.GPU(duration=GPU_DURATION, size=GPU_SIZE)
+def build_advanced_character_pack(
+    reference_files,
+    reference_zip,
+    driving_video,
+    prompt,
+    sam3_text,
+    max_subjects,
+    progress=gr.Progress(track_tqdm=True),
+):
+    try:
+        progress(0.0, desc="Checking SCAIL-Pose repo")
+        _require_repo_layout()
+        progress(0.04, desc="Preparing SAM3 weights")
+        sam3_model = _ensure_sam3_weights()
+
+        refs = _collect_character_refs(reference_files, reference_zip)
+        character_names = sorted({ref.character for ref in refs}, key=_character_sort_key)
+        if len(character_names) > len(SCAIL_COLORS):
+            raise RuntimeError(f"Too many characters: {len(character_names)}. Limit is {len(SCAIL_COLORS)}.")
+
+        max_targets = max(int(max_subjects or len(character_names)), len(character_names))
+        max_targets = max(1, min(max_targets, len(SCAIL_COLORS)))
+        colors = SCAIL_COLORS[:max_targets]
+
+        job_dir = _new_job_dir("advanced_animation")
+        pack_root = job_dir / "scail2_input_pack"
+        pack_root.mkdir(parents=True, exist_ok=True)
+
+        progress(0.08, desc="Preparing references and driving video")
+        work_driving = job_dir / "driving.mp4"
+        _copy_video(driving_video, work_driving, "driving video")
+        shutil.copy2(work_driving, pack_root / "rendered_v2.mp4")
+
+        prepared_refs = _prepare_advanced_refs(refs, pack_root)
+        primary = sorted(
+            prepared_refs,
+            key=lambda ref: (
+                0 if ref.character == "character_0" else 1,
+                0 if ref.view in {"front", "main", "ref", "reference"} else 1,
+                _character_sort_key(ref.character),
+                ref.view,
+            ),
+        )[0]
+        _build_advanced_metadata(pack_root, "animation", prompt, primary, character_names, sam3_text)
+
+        progress(0.12, desc="Loading SAM3 predictors")
+        video_predictor, image_predictor = _load_sam3_predictors(sam3_model)
+
+        from TrackSam3.track import get_mask_from_image, get_mask_from_image_via_video, get_mask_from_video
+        from NLFPoseExtract.v2_helper import save_colored_mask_image, write_colored_mask_video
+
+        text = _text_args(sam3_text)
+
+        progress(0.20, desc="Tracking subjects in driving video")
+        drv_masks, drv_colors = get_mask_from_video(
+            str(work_driving),
+            video_predictor,
+            max_targets=max_targets,
+            sort_by="x",
+            fixed_colors=colors,
+            text=text,
+        )
+        if len(drv_masks) == 0:
+            raise RuntimeError("SAM3 did not detect any valid subject in the driving video.")
+
+        import cv2
+
+        cap = cv2.VideoCapture(str(work_driving))
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        cap.release()
+        fps_int = max(1, int(round(fps or 24)))
+        write_colored_mask_video(
+            drv_masks,
+            drv_colors,
+            str(pack_root / "rendered_mask_v2.mp4"),
+            fps_int,
+            bg_color=(255, 255, 255),
+        )
+
+        progress(0.52, desc="Generating reference masks")
+        logs = [
+            f"Characters: {', '.join(character_names)}",
+            f"Reference views: {len(prepared_refs)}",
+            f"Driving tracks detected: {len(drv_masks)}",
+            f"Primary: {primary.character}/{primary.view}",
+        ]
+        for idx, ref in enumerate(prepared_refs):
+            character_idx = character_names.index(ref.character)
+            color = drv_colors[character_idx] if character_idx < len(drv_colors) else SCAIL_COLORS[character_idx]
+            mask_path = ref.path.with_name(f"{ref.view}_mask.png")
+            progress(None, desc=f"Masking {ref.character}/{ref.view}")
+            try:
+                ref_masks, _ = get_mask_from_image_via_video(
+                    str(ref.path),
+                    video_predictor,
+                    max_targets=1,
+                    sort_by="x",
+                    fixed_colors=[color],
+                    text=text,
+                )
+            except Exception as exc:
+                logging.warning("Video-based image masking failed for %s, trying image predictor: %s", ref.path, exc)
+                ref_masks, _ = get_mask_from_image(
+                    str(ref.path),
+                    image_predictor,
+                    max_targets=1,
+                    sort_by="x",
+                    fixed_colors=[color],
+                    text=text,
+                )
+            if len(ref_masks) == 0:
+                raise RuntimeError(f"SAM3 did not detect a subject in {ref.character}/{ref.view}.")
+            save_colored_mask_image([ref_masks[0]], [color], str(mask_path), bg_color=(255, 255, 255))
+            logs.append(f"- {ref.character}/{ref.view}: {mask_path.relative_to(pack_root).as_posix()}")
+
+        progress(0.92, desc="Packaging canonical Advanced Pack")
+        zip_path = _write_canonical_pack_zip(pack_root, "animation")
+        gallery = _advanced_gallery(pack_root)
+        progress(1.0, desc="Done")
+        status = f"Done. Advanced pack created at {zip_path}\n\n" + "\n".join(logs)
+        return gallery, str(pack_root / "rendered_v2.mp4"), str(pack_root / "rendered_mask_v2.mp4"), str(zip_path), status
+    except Exception:
+        logging.exception("Advanced character pack generation failed")
+        return [], None, None, None, traceback.format_exc()
+
+
 def build_ui():
     with gr.Blocks(title="SCAIL-Pose Pack Builder") as demo:
         gr.Markdown(
@@ -443,6 +780,74 @@ def build_ui():
                 outputs=[anim_ref_mask, anim_rendered, anim_mask_video, anim_zip, anim_status],
             )
 
+        with gr.Tab("Advanced Character Pack"):
+            gr.Markdown(
+                "Create a canonical SCAIL-2 animation pack for multi-reference or multi-character cases. "
+                "Use this when one character has several views, or when several characters need separate reference slots."
+            )
+            with gr.Accordion("Reference naming", open=True):
+                gr.Markdown(
+                    "Upload reference images with names like:\n\n"
+                    "```text\n"
+                    "character_0__front.png\n"
+                    "character_0__back.png\n"
+                    "character_0__closeup.png\n"
+                    "character_1__front.png\n"
+                    "```\n\n"
+                    "Or upload a zip with this structure:\n\n"
+                    "```text\n"
+                    "characters/\n"
+                    "  character_0/\n"
+                    "    front.png\n"
+                    "    back.png\n"
+                    "  character_1/\n"
+                    "    front.png\n"
+                    "```\n\n"
+                    "Character colors are assigned by slot order: `character_0`, then `character_1`, etc. "
+                    "The driving video is tracked left-to-right with the same color order."
+                )
+            with gr.Row():
+                adv_refs = gr.Files(
+                    label="Reference images",
+                    file_types=[".png", ".jpg", ".jpeg", ".webp"],
+                    type="filepath",
+                )
+                adv_zip_in = gr.File(
+                    label="Reference zip instead of files",
+                    file_types=[".zip"],
+                    type="filepath",
+                )
+            adv_driving = gr.Video(label="Driving video")
+            adv_prompt = gr.Textbox(label="Prompt for SCAIL-2", lines=3)
+            with gr.Row():
+                adv_sam3_text = gr.Textbox(value="human character", label="SAM3 text prompt")
+                adv_max_subjects = gr.Number(value=2, precision=0, label="Max tracked subjects")
+            adv_run = gr.Button("Generate advanced character pack", variant="primary")
+            adv_gallery = gr.Gallery(
+                label="Parsed references and generated masks",
+                columns=4,
+                height=320,
+                selected_index=0,
+                preview=True,
+            )
+            with gr.Row():
+                adv_rendered = gr.Video(label="Rendered / driving video")
+                adv_mask_video = gr.Video(label="Driving mask video")
+            adv_zip = gr.File(label="Download canonical SCAIL-2 Advanced Pack")
+            adv_status = gr.Textbox(label="Run logs", lines=14)
+            adv_run.click(
+                build_advanced_character_pack,
+                inputs=[
+                    adv_refs,
+                    adv_zip_in,
+                    adv_driving,
+                    adv_prompt,
+                    adv_sam3_text,
+                    adv_max_subjects,
+                ],
+                outputs=[adv_gallery, adv_rendered, adv_mask_video, adv_zip, adv_status],
+            )
+
         with gr.Tab("Replacement Pack"):
             gr.Markdown(
                 "Create a replacement pack from a replacement reference image and a source video. "
@@ -478,9 +883,9 @@ def build_ui():
 
         with gr.Tab("Pack Format"):
             gr.Markdown(
-                "The exported zip is intentionally flat for V1, so it is easy to inspect and "
-                "compatible with the SCAIL-2 Advanced Pack parser.\n\n"
-                "Animation pack:\n"
+                "The simple tabs export flat packs. The advanced tab exports a canonical "
+                "character pack with explicit character folders.\n\n"
+                "Simple animation pack:\n"
                 "```text\n"
                 "ref.png\n"
                 "ref_mask.jpg\n"
@@ -498,8 +903,24 @@ def build_ui():
                 "prompt.txt\n"
                 "metadata.json\n"
                 "```\n\n"
-                "For multi-character curation, run several focused jobs first, inspect the masks, "
-                "then combine the selected references into a canonical Advanced Pack."
+                "Advanced character pack:\n"
+                "```text\n"
+                "rendered_v2.mp4\n"
+                "rendered_mask_v2.mp4\n"
+                "prompt.txt\n"
+                "metadata.json\n"
+                "characters/\n"
+                "  character_0/\n"
+                "    front.png\n"
+                "    front_mask.png\n"
+                "    back.png\n"
+                "    back_mask.png\n"
+                "  character_1/\n"
+                "    front.png\n"
+                "    front_mask.png\n"
+                "```\n\n"
+                "For multi-reference, use several views under the same `character_N`. "
+                "For multi-character, use one folder per identity slot."
             )
 
     return demo