clean UI and Advanced Pack for multi refs & multi characters

app.py

@@ -1,4 +1,5 @@
 import gc
+import json
 import logging
 import os
 import shutil
@@ -7,6 +8,7 @@ import sys
 import tempfile
 import traceback
 import uuid
+import zipfile
 from dataclasses import dataclass
 from pathlib import Path
 from types import SimpleNamespace
@@ -41,8 +43,8 @@ SAFETENSORS_REPO_ID = os.getenv("SCAIL_SAFETENSORS_REPO_ID")
 SAFETENSORS_FILENAME = os.getenv("SCAIL_SAFETENSORS_FILENAME", "SCAIL-2.safetensors")
 MODEL_NAME = os.getenv("SCAIL_MODEL_NAME", "SCAIL-14B")
 GPU_SIZE = os.getenv("SCAIL_ZEROGPU_SIZE", "xlarge")
-GPU_DURATION_COLD = int(os.getenv("SCAIL_GPU_DURATION_COLD", "900"))
-GPU_DURATION_WARM = int(os.getenv("SCAIL_GPU_DURATION_WARM", "360"))
+GPU_DURATION_COLD = int(os.getenv("SCAIL_GPU_DURATION_COLD", "600"))
+GPU_DURATION_WARM = int(os.getenv("SCAIL_GPU_DURATION_WARM", "330"))
 DEFAULT_TARGET_H = int(os.getenv("SCAIL_TARGET_H", "512"))
 DEFAULT_TARGET_W = int(os.getenv("SCAIL_TARGET_W", "896"))
 DEFAULT_SEGMENT_LEN = int(os.getenv("SCAIL_SEGMENT_LEN", "81"))
@@ -54,6 +56,8 @@ AUTO_CONVERT = os.getenv("SCAIL_AUTO_CONVERT", "1") == "1"
 PRELOAD_PIPELINE = os.getenv("SCAIL_PRELOAD_PIPELINE", "1") == "1"
 STAGE_SAFETENSORS_FOR_LOAD = os.getenv("SCAIL_STAGE_SAFETENSORS_FOR_LOAD", "1") == "1"
 CONVERT_TO_STAGING_FIRST = os.getenv("SCAIL_CONVERT_TO_STAGING_FIRST", "1") == "1"
+MAX_ADDITIONAL_REFS = int(os.getenv("SCAIL_MAX_ADDITIONAL_REFS", "16"))
+
 CLIP_CKPT_NAME = "models_clip_open-clip-xlm-roberta-large-vit-huge-14-onlyvisual.pth"
 ORIGINAL_DIT_REL_PATH = "model/1/fsdp2_rank_0000_checkpoint.pt"
 BASE_ALLOW_PATTERNS = [
@@ -77,6 +81,13 @@ _SCAIL_CONFIGS = None
 _SCAIL_CONFIG_PATHS = None
 
 
+@dataclass(frozen=True)
+class ReferencePair:
+    label: str
+    image: str
+    mask_image: str
+
+
 @dataclass(frozen=True)
 class PreparedExample:
     label: str
@@ -86,6 +97,7 @@ class PreparedExample:
     mask_video: str
     prompt: str
     replace_flag: bool = False
+    additional_refs: tuple[ReferencePair, ...] = ()
 
 
 PREPARED_EXAMPLES = {
@@ -120,6 +132,23 @@ PREPARED_EXAMPLES = {
         pose="examples/animation_003_multi_ref/rendered_v2.mp4",
         mask_video="examples/animation_003_multi_ref/rendered_mask_v2.mp4",
         prompt="A character performs the motion from the driving video.",
+        additional_refs=(
+            ReferencePair(
+                label="Background",
+                image="examples/animation_003_multi_ref/background.png",
+                mask_image="examples/animation_003_multi_ref/background_mask.png",
+            ),
+            ReferencePair(
+                label="Reference view 1",
+                image="examples/animation_003_multi_ref/character_0.png",
+                mask_image="examples/animation_003_multi_ref/character_0_mask.png",
+            ),
+            ReferencePair(
+                label="Reference view 2",
+                image="examples/animation_003_multi_ref/character_1.png",
+                mask_image="examples/animation_003_multi_ref/character_1_mask.png",
+            ),
+        ),
     ),
     "Replacement 001": PreparedExample(
         label="Replacement 001",
@@ -143,15 +172,307 @@ def _abs(path: str | Path) -> str:
     return str(path)
 
 
+def _example_paths(example: PreparedExample) -> list[str]:
+    paths = [example.image, example.mask_image, example.pose, example.mask_video]
+    for ref in example.additional_refs:
+        paths.extend([ref.image, ref.mask_image])
+    return paths
+
+
 def _existing_examples() -> dict[str, PreparedExample]:
     available = {}
     for name, example in PREPARED_EXAMPLES.items():
-        paths = [example.image, example.mask_image, example.pose, example.mask_video]
-        if all(Path(_abs(p)).exists() for p in paths):
+        if all(Path(_abs(path)).exists() for path in _example_paths(example)):
             available[name] = example
     return available
 
 
+IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp")
+VIDEO_EXTS = (".mp4", ".mov", ".webm", ".mkv")
+PRIMARY_STEMS = ("front", "main", "ref", "reference")
+
+
+def _safe_extract_zip(zip_path: str | Path) -> Path:
+    zip_path = Path(zip_path)
+    if not zip_path.exists():
+        raise RuntimeError(f"Pack does not exist: {zip_path}")
+    if zip_path.suffix.lower() != ".zip":
+        raise RuntimeError("Advanced Pack expects a .zip file.")
+
+    extract_root = Path(tempfile.gettempdir()) / "scail2_input_packs" / 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 _read_metadata(pack_root: Path) -> dict:
+    metadata_path = pack_root / "metadata.json"
+    if not metadata_path.exists():
+        return {}
+    try:
+        return json.loads(metadata_path.read_text(encoding="utf-8"))
+    except Exception as exc:
+        raise RuntimeError(f"Invalid metadata.json: {exc}") from exc
+
+
+def _read_pack_prompt(pack_root: Path, metadata: dict) -> str:
+    if isinstance(metadata.get("prompt"), str):
+        return metadata["prompt"]
+    prompt_path = pack_root / "prompt.txt"
+    if prompt_path.exists():
+        return prompt_path.read_text(encoding="utf-8").strip()
+    return ""
+
+
+def _find_stem_file(directory: Path, stems: tuple[str, ...], exts: tuple[str, ...]) -> Path | None:
+    for stem in stems:
+        stem_path = Path(stem)
+        if stem_path.suffix:
+            candidate = directory / stem
+            if candidate.exists() and candidate.is_file():
+                return candidate
+            continue
+        for ext in exts:
+            candidate = directory / f"{stem}{ext}"
+            if candidate.exists() and candidate.is_file():
+                return candidate
+    return None
+
+
+def _mask_for_image(image_path: Path) -> Path | None:
+    for ext in IMAGE_EXTS:
+        candidate = image_path.with_name(f"{image_path.stem}_mask{ext}")
+        if candidate.exists() and candidate.is_file():
+            return candidate
+    return None
+
+
+def _rel(path: Path, root: Path) -> str:
+    try:
+        return path.relative_to(root).as_posix()
+    except ValueError:
+        return path.name
+
+
+def _collect_pairs_from_dir(directory: Path, root: Path, identity: str) -> list[ReferencePair]:
+    if not directory.exists() or not directory.is_dir():
+        return []
+
+    pairs = []
+    for image_path in sorted(directory.iterdir(), key=lambda p: p.name.lower()):
+        if not image_path.is_file() or image_path.suffix.lower() not in IMAGE_EXTS:
+            continue
+        if image_path.stem.endswith("_mask"):
+            continue
+        mask_path = _mask_for_image(image_path)
+        if mask_path is None:
+            raise RuntimeError(f"Missing mask for `{_rel(image_path, root)}`.")
+        label = f"{identity}: {image_path.stem}"
+        pairs.append(ReferencePair(label=label, image=str(image_path), mask_image=str(mask_path)))
+    return pairs
+
+
+def _collect_character_pairs(pack_root: Path) -> list[ReferencePair]:
+    characters_dir = pack_root / "characters"
+    if not characters_dir.exists():
+        return []
+
+    pairs = []
+    for identity_dir in sorted(characters_dir.iterdir(), key=lambda p: p.name.lower()):
+        if identity_dir.is_dir():
+            pairs.extend(_collect_pairs_from_dir(identity_dir, pack_root, identity_dir.name))
+    return pairs
+
+
+def _collect_environment_pairs(pack_root: Path) -> list[ReferencePair]:
+    pairs = _collect_pairs_from_dir(pack_root / "environment", pack_root, "environment")
+    for stem in ("background", "environment"):
+        image_path = _find_stem_file(pack_root, (stem,), IMAGE_EXTS)
+        if image_path is None:
+            continue
+        mask_path = _mask_for_image(image_path)
+        if mask_path is None:
+            raise RuntimeError(f"Missing mask for `{_rel(image_path, pack_root)}`.")
+        pairs.append(ReferencePair(label=f"environment: {image_path.stem}", image=str(image_path), mask_image=str(mask_path)))
+    return pairs
+
+
+def _collect_legacy_flat_pairs(pack_root: Path) -> list[ReferencePair]:
+    pairs = []
+    for image_path in sorted(pack_root.iterdir(), key=lambda p: p.name.lower()):
+        if not image_path.is_file() or image_path.suffix.lower() not in IMAGE_EXTS:
+            continue
+        if image_path.stem.endswith("_mask"):
+            continue
+        if not image_path.stem.startswith("character_"):
+            continue
+        mask_path = _mask_for_image(image_path)
+        if mask_path is None:
+            raise RuntimeError(f"Missing mask for `{_rel(image_path, pack_root)}`.")
+        pairs.append(ReferencePair(label=f"reference: {image_path.stem}", image=str(image_path), mask_image=str(mask_path)))
+    return pairs
+
+
+def _primary_sort_key(ref: ReferencePair):
+    stem = Path(ref.image).stem.lower()
+    if stem in PRIMARY_STEMS:
+        return (0, PRIMARY_STEMS.index(stem), ref.label.lower())
+    return (1, ref.label.lower())
+
+
+def _resolve_metadata_file(pack_root: Path, rel_path: str | None, label: str) -> Path | None:
+    if not rel_path:
+        return None
+    path = pack_root / rel_path
+    if not path.exists() or not path.is_file():
+        raise RuntimeError(f"metadata.json references missing {label}: `{rel_path}`.")
+    return path
+
+
+def _select_pack_primary(pack_root: Path, metadata: dict, character_pairs: list[ReferencePair]) -> ReferencePair:
+    primary = metadata.get("primary") if isinstance(metadata.get("primary"), dict) else {}
+    primary_image = _resolve_metadata_file(pack_root, primary.get("image"), "primary image")
+    primary_mask = _resolve_metadata_file(pack_root, primary.get("mask"), "primary mask")
+    if primary_image is not None or primary_mask is not None:
+        if primary_image is None or primary_mask is None:
+            raise RuntimeError("metadata.json primary must include both `image` and `mask`.")
+        return ReferencePair(label="metadata primary", image=str(primary_image), mask_image=str(primary_mask))
+
+    ref_image = _find_stem_file(pack_root, ("ref", "main", "reference"), IMAGE_EXTS)
+    if ref_image is not None:
+        ref_mask = _mask_for_image(ref_image)
+        if ref_mask is None:
+            raise RuntimeError(f"Missing mask for primary reference `{_rel(ref_image, pack_root)}`.")
+        return ReferencePair(label="primary reference", image=str(ref_image), mask_image=str(ref_mask))
+
+    if not character_pairs:
+        raise RuntimeError(
+            "No primary reference found. Provide `ref.png` + `ref_mask.png`, "
+            "or at least one pair under `characters/character_0/`."
+        )
+    character_0 = [ref for ref in character_pairs if "/character_0/" in Path(ref.image).as_posix()]
+    candidates = character_0 or character_pairs
+    selected = sorted(candidates, key=_primary_sort_key)[0]
+    return ReferencePair(label=f"{selected.label} (auto primary)", image=selected.image, mask_image=selected.mask_image)
+
+
+def _find_pack_video(pack_root: Path, metadata: dict) -> Path:
+    driving = metadata.get("driving") if isinstance(metadata.get("driving"), dict) else {}
+    metadata_video = _resolve_metadata_file(pack_root, driving.get("video"), "driving video")
+    if metadata_video is not None:
+        return metadata_video
+    video = _find_stem_file(pack_root, ("rendered_v2", "driving", "pose"), VIDEO_EXTS)
+    if video is None:
+        raise RuntimeError("Missing driving video. Expected `rendered_v2.mp4`.")
+    return video
+
+
+def _find_pack_mask_video(pack_root: Path, metadata: dict) -> tuple[Path, bool]:
+    driving = metadata.get("driving") if isinstance(metadata.get("driving"), dict) else {}
+    metadata_mask = _resolve_metadata_file(pack_root, driving.get("mask_video"), "driving mask video")
+    if metadata_mask is not None:
+        return metadata_mask, metadata_mask.name == "replace_mask.mp4" or metadata.get("mode") == "replacement"
+
+    rendered_mask = pack_root / "rendered_mask_v2.mp4"
+    replace_mask = pack_root / "replace_mask.mp4"
+    if rendered_mask.exists() and replace_mask.exists():
+        raise RuntimeError("Found both `rendered_mask_v2.mp4` and `replace_mask.mp4`; keep only one or set metadata.json.")
+    if replace_mask.exists():
+        return replace_mask, True
+    if rendered_mask.exists():
+        return rendered_mask, False
+    raise RuntimeError("Missing mask video. Expected `rendered_mask_v2.mp4` or `replace_mask.mp4`.")
+
+
+def _same_pair(a: ReferencePair, b: ReferencePair) -> bool:
+    return Path(a.image).resolve() == Path(b.image).resolve() and Path(a.mask_image).resolve() == Path(b.mask_image).resolve()
+
+
+def parse_input_pack(pack_zip: str | Path) -> dict:
+    pack_root = _safe_extract_zip(pack_zip)
+    metadata = _read_metadata(pack_root)
+    character_pairs = _collect_character_pairs(pack_root)
+    environment_pairs = _collect_environment_pairs(pack_root)
+    legacy_pairs = _collect_legacy_flat_pairs(pack_root)
+    primary = _select_pack_primary(pack_root, metadata, character_pairs + legacy_pairs)
+    driving_video = _find_pack_video(pack_root, metadata)
+    mask_video, replace_flag = _find_pack_mask_video(pack_root, metadata)
+
+    refs = sorted(character_pairs + legacy_pairs, key=lambda ref: ref.label.lower()) + environment_pairs
+    additional_refs = [ref for ref in refs if not _same_pair(ref, primary)]
+    if len(additional_refs) > MAX_ADDITIONAL_REFS:
+        raise RuntimeError(f"Too many additional references: {len(additional_refs)}. Limit is {MAX_ADDITIONAL_REFS}.")
+
+    return {
+        "root": str(pack_root),
+        "prompt": _read_pack_prompt(pack_root, metadata),
+        "mode": "replacement" if replace_flag else "animation",
+        "replace_flag": replace_flag,
+        "image": primary.image,
+        "mask_image": primary.mask_image,
+        "pose": str(driving_video),
+        "mask_video": str(mask_video),
+        "primary_label": primary.label,
+        "additional_refs": [
+            {"label": ref.label, "image": ref.image, "mask_image": ref.mask_image}
+            for ref in additional_refs
+        ],
+    }
+
+
+def _pack_gallery(pack: dict):
+    items = [
+        (pack["image"], f"Primary: {pack['primary_label']}"),
+        (pack["mask_image"], "Primary mask"),
+    ]
+    for ref in pack["additional_refs"]:
+        items.append((ref["image"], ref["label"]))
+        items.append((ref["mask_image"], f"{ref['label']} mask"))
+    return items
+
+
+def _pack_summary(pack: dict) -> str:
+    lines = [
+        "### Pack validated",
+        f"- Mode: `{pack['mode']}`",
+        f"- Primary: `{pack['primary_label']}`",
+        f"- Driving video: `{Path(pack['pose']).name}`",
+        f"- Mask video: `{Path(pack['mask_video']).name}`",
+        f"- Additional reference pairs: `{len(pack['additional_refs'])}`",
+    ]
+    for ref in pack["additional_refs"]:
+        lines.append(f"  - `{ref['label']}`")
+    return "\n".join(lines)
+
+
+def validate_input_pack(pack_zip):
+    if pack_zip is None:
+        return None, "Upload a `.zip` pack first.", [], None, None, "", "animation"
+    try:
+        pack = parse_input_pack(pack_zip)
+        return (
+            pack,
+            _pack_summary(pack),
+            _pack_gallery(pack),
+            pack["pose"],
+            pack["mask_video"],
+            pack["prompt"],
+            pack["mode"],
+        )
+    except Exception:
+        logging.exception("Advanced pack validation failed")
+        return None, traceback.format_exc(), [], None, None, "", "animation"
+
+
 def _require_repo_layout():
     missing = []
     for rel in ("wan/scail.py", "wan/modules/model_scail2.py", "generate.py", "configs/config-14b.json"):
@@ -174,11 +495,7 @@ def _download_safetensors_if_configured() -> Path | None:
     if local_path.exists():
         return local_path
 
-    logging.info(
-        "Downloading converted SCAIL-2 safetensors from %s/%s",
-        SAFETENSORS_REPO_ID,
-        SAFETENSORS_FILENAME,
-    )
+    logging.info("Downloading converted SCAIL-2 safetensors from %s/%s", SAFETENSORS_REPO_ID, SAFETENSORS_FILENAME)
     downloaded = hf_hub_download(
         repo_id=SAFETENSORS_REPO_ID,
         filename=SAFETENSORS_FILENAME,
@@ -203,10 +520,7 @@ def _find_converted_safetensors(ckpt_dir: Path | None) -> Path | None:
         Path(os.getenv("SCAIL_CONVERTED_DIR", str(STORAGE_ROOT / "scail2_converted"))) / "SCAIL-2.safetensors",
     ]
     if ckpt_dir is not None:
-        candidates += [
-            ckpt_dir / "SCAIL-2.safetensors",
-            ckpt_dir / "model.safetensors",
-        ]
+        candidates += [ckpt_dir / "SCAIL-2.safetensors", ckpt_dir / "model.safetensors"]
     for candidate in candidates:
         if candidate.exists():
             return candidate
@@ -231,11 +545,7 @@ def _copy_file_with_progress(source: Path, dest: Path, description: str) -> Path
     logging.info("%s: %s -> %s", description, source, dest)
     next_log = ((copied // log_every) + 1) * log_every if log_every > 0 else source_size
     if copied:
-        logging.info(
-            "Resuming copy at %.2f/%.2f GB",
-            copied / 1024**3,
-            source_size / 1024**3,
-        )
+        logging.info("Resuming copy at %.2f/%.2f GB", copied / 1024**3, source_size / 1024**3)
 
     with source.open("rb") as src, tmp_dest.open("ab") as dst:
         if copied:
@@ -247,12 +557,7 @@ def _copy_file_with_progress(source: Path, dest: Path, description: str) -> Path
             dst.write(chunk)
             copied += len(chunk)
             if log_every > 0 and copied >= next_log:
-                logging.info(
-                    "%s: %.2f/%.2f GB",
-                    description,
-                    copied / 1024**3,
-                    source_size / 1024**3,
-                )
+                logging.info("%s: %.2f/%.2f GB", description, copied / 1024**3, source_size / 1024**3)
                 next_log += log_every
 
     if tmp_dest.stat().st_size != source_size:
@@ -279,11 +584,9 @@ def _stage_safetensors_for_load(scail_path: Path) -> Path:
         return source
 
     stage_dir = Path(os.getenv("SCAIL_MODEL_LOAD_CACHE", str(STAGING_ROOT / "scail2_model_load")))
-    stage_dir.mkdir(parents=True, exist_ok=True)
     staged = stage_dir / source.name
     if staged.exists() and staged.stat().st_size == source.stat().st_size:
         return staged
-
     return _copy_file_with_progress(source, staged, "Staging SCAIL-2 safetensors for local load")
 
 
@@ -304,13 +607,9 @@ def _download_checkpoint_if_needed(include_original_dit: bool = False) -> Path:
     if has_base_assets:
         return local_dir
 
-    logging.info("Downloading SCAIL-2 base checkpoint assets from %s", MODEL_REPO_ID)
     if include_original_dit:
-        logging.warning(
-            "include_original_dit is deprecated here; original DiT conversion staging "
-            "uses SCAIL_ORIGINAL_DIT_CACHE instead."
-        )
-
+        logging.warning("Original DiT staging uses SCAIL_ORIGINAL_DIT_CACHE; base download will stay narrow.")
+    logging.info("Downloading SCAIL-2 base checkpoint assets from %s", MODEL_REPO_ID)
     snapshot_download(
         repo_id=MODEL_REPO_ID,
         local_dir=str(local_dir),
@@ -335,10 +634,7 @@ def _download_original_dit_for_conversion() -> Path:
     if original_path.exists():
         return local_dir
 
-    logging.info(
-        "Downloading original SCAIL-2 DiT checkpoint for one-time conversion into %s",
-        local_dir,
-    )
+    logging.info("Downloading original SCAIL-2 DiT checkpoint for one-time conversion into %s", local_dir)
     snapshot_download(
         repo_id=MODEL_REPO_ID,
         local_dir=str(local_dir),
@@ -350,7 +646,6 @@ def _download_original_dit_for_conversion() -> Path:
 
 
 def _prepare_assets_for_runtime() -> str:
-    """Download CPU-side assets before any ZeroGPU function starts."""
     global _ASSET_STATUS, _ASSET_ERROR
     try:
         ckpt_dir = _download_checkpoint_if_needed(include_original_dit=False)
@@ -360,16 +655,11 @@ def _prepare_assets_for_runtime() -> str:
             scail_path = _maybe_convert_checkpoint(original_dir, None)
         if scail_path is None:
             _ASSET_STATUS = (
-                "Base checkpoint assets are present, but no converted safetensors file "
-                "was found. Set SCAIL_SAFETENSORS_REPO_ID plus SCAIL_SAFETENSORS_FILENAME, "
-                "or set SCAIL_SAFETENSORS_PATH. Automatic startup conversion is disabled "
-                "only when SCAIL_AUTO_CONVERT=0."
+                "Base checkpoint assets are present, but no converted safetensors file was found. "
+                "Set SCAIL_SAFETENSORS_PATH or SCAIL_SAFETENSORS_REPO_ID."
             )
         else:
-            _ASSET_STATUS = (
-                "Assets ready. Base checkpoint: "
-                f"{ckpt_dir}. Converted DiT safetensors: {scail_path}."
-            )
+            _ASSET_STATUS = f"Assets ready. Base checkpoint: {ckpt_dir}. Converted DiT safetensors: {scail_path}."
         _ASSET_ERROR = None
     except Exception:
         _ASSET_ERROR = traceback.format_exc()
@@ -385,10 +675,8 @@ def _maybe_convert_checkpoint(ckpt_dir: Path, scail_path: Path | None) -> Path:
 
     if not AUTO_CONVERT:
         raise RuntimeError(
-            "Converted SCAIL-2 safetensors file was not found. For the wan-scail2 branch, "
-            "provide SCAIL_SAFETENSORS_PATH=/path/to/SCAIL-2.safetensors, or place "
-            "SCAIL-2.safetensors at the repo root. Automatic startup conversion is disabled "
-            "because SCAIL_AUTO_CONVERT=0."
+            "Converted SCAIL-2 safetensors file was not found. Set SCAIL_SAFETENSORS_PATH, "
+            "or enable SCAIL_AUTO_CONVERT=1 for one-time conversion."
         )
 
     persistent_dir = Path(os.getenv("SCAIL_CONVERTED_DIR", str(STORAGE_ROOT / "scail2_converted")))
@@ -396,9 +684,8 @@ def _maybe_convert_checkpoint(ckpt_dir: Path, scail_path: Path | None) -> Path:
     if persistent_path.exists():
         return persistent_path
 
-    if CONVERT_TO_STAGING_FIRST:
-        save_dir = Path(os.getenv("SCAIL_CONVERSION_WORK_DIR", str(STAGING_ROOT / "scail2_converted_work")))
-    else:
+    save_dir = Path(os.getenv("SCAIL_CONVERSION_WORK_DIR", str(STAGING_ROOT / "scail2_converted_work")))
+    if not CONVERT_TO_STAGING_FIRST:
         save_dir = persistent_dir
     save_dir.mkdir(parents=True, exist_ok=True)
     save_path = save_dir / "SCAIL-2.safetensors"
@@ -410,9 +697,6 @@ def _maybe_convert_checkpoint(ckpt_dir: Path, scail_path: Path | None) -> Path:
 
     logging.info("Converting checkpoint to safetensors: %s", save_path)
     convert_env = os.environ.copy()
-    # PyTorch >= 2.6 changed torch.load's default to weights_only=True.
-    # The official SCAIL-2 FSDP checkpoint needs the legacy trusted-pickle path
-    # during conversion. Only do this for the official checkpoint conversion step.
     convert_env["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1"
     subprocess.run(
         [
@@ -434,11 +718,6 @@ def _maybe_convert_checkpoint(ckpt_dir: Path, scail_path: Path | None) -> Path:
 
 
 def _install_attention_patch():
-    """Use HF Kernels flash-attn2 when available; otherwise fall back to SDPA.
-
-    SCAIL-2 imports `flash_attention` directly in model_scail2.py. This monkey patch
-    avoids requiring a locally built `flash_attn` wheel on Spaces.
-    """
     import wan.modules.attention as attention_mod
 
     hf_flash_attn2 = None
@@ -508,7 +787,6 @@ def _install_attention_patch():
 
             cu_q = torch.cat([q_lens_t.new_zeros([1]), q_lens_t]).cumsum(0, dtype=torch.int32)
             cu_k = torch.cat([k_lens_t.new_zeros([1]), k_lens_t]).cumsum(0, dtype=torch.int32)
-
             try:
                 out = hf_flash_attn2.flash_attn_varlen_func(
                     q=q_var,
@@ -550,9 +828,6 @@ def _install_attention_patch():
         return out.transpose(1, 2).contiguous().type(out_dtype)
 
     attention_mod.flash_attention = patched_flash_attention
-
-    # Several modules import `flash_attention` directly at import time
-    # (`from .attention import flash_attention`), so patch those bound names too.
     for module_name in (
         "wan.modules.clip",
         "wan.modules.model",
@@ -589,7 +864,6 @@ def _import_runtime():
 
 
 def _prepare_runtime_for_startup() -> str:
-    """Import Wan/SCAIL and resolve HF Kernels before any ZeroGPU call."""
     global _RUNTIME_STATUS, _RUNTIME_ERROR
     try:
         _import_runtime()
@@ -613,6 +887,7 @@ def _get_pipeline():
         scail_path = _maybe_convert_checkpoint(original_dir, None)
     else:
         scail_path = _maybe_convert_checkpoint(ckpt_dir, scail_path)
+
     scail_load_path = _stage_safetensors_for_load(scail_path)
     config_path = Path(os.getenv("SCAIL_CONFIG_PATH", _SCAIL_CONFIG_PATHS[MODEL_NAME]))
     if not config_path.is_absolute():
@@ -645,7 +920,6 @@ def _get_pipeline():
 
 
 def _prepare_pipeline_for_startup() -> str:
-    """Preload the CUDA pipeline at module startup for ZeroGPU CUDA emulation."""
     global _PIPELINE_STATUS, _PIPELINE_ERROR
     try:
         _get_pipeline()
@@ -658,6 +932,27 @@ def _prepare_pipeline_for_startup() -> str:
     return _PIPELINE_STATUS if _PIPELINE_ERROR is None else _PIPELINE_STATUS + "\n\n" + _PIPELINE_ERROR
 
 
+def _is_gradio_native_file_path(path: Path) -> bool:
+    path = path.resolve()
+    native_roots = [ROOT.resolve(), Path(tempfile.gettempdir()).resolve()]
+    return any(_is_relative_to(path, root) for root in native_roots)
+
+
+def _prepare_output_for_gradio(path: str | Path) -> str:
+    source = Path(path)
+    if not source.exists():
+        raise RuntimeError(f"Generated video was not found: {source}")
+    if _is_gradio_native_file_path(source):
+        return str(source)
+
+    gradio_dir = Path(os.getenv("SCAIL_GRADIO_OUTPUT_CACHE", str(Path(tempfile.gettempdir()) / "scail2_gradio_outputs")))
+    gradio_dir.mkdir(parents=True, exist_ok=True)
+    dest = gradio_dir / source.name
+    shutil.copy2(source, dest)
+    logging.info("Copied generated video for Gradio display: %s -> %s", source, dest)
+    return str(dest)
+
+
 def _duration_for_job(*args, **kwargs):
     if _PIPELINE is None:
         return int(os.getenv("SCAIL_GPU_DURATION", str(GPU_DURATION_COLD)))
@@ -679,6 +974,7 @@ def _run_scail_job(
     seed,
     segment_len,
     segment_overlap,
+    additional_refs: tuple[ReferencePair, ...] = (),
     progress=None,
 ):
     if progress is not None:
@@ -686,6 +982,7 @@ def _run_scail_job(
     pipeline = _get_pipeline()
     cfg = _SCAIL_CONFIGS[MODEL_NAME]
     save_file = OUTPUT_DIR / f"scail2_{uuid.uuid4().hex}.mp4"
+
     if progress is not None:
         progress(0.12, desc="Preparing inputs")
     args = SimpleNamespace(
@@ -704,6 +1001,13 @@ def _run_scail_job(
         prompt=prompt or "",
     )
 
+    additional_task_input = None
+    if additional_refs:
+        additional_task_input = {
+            "additional_ref_image_paths": [str(ref.image) for ref in additional_refs],
+            "additional_ref_mask_image_paths": [str(ref.mask_image) for ref in additional_refs],
+        }
+
     if progress is not None:
         progress(0.15, desc="Generating video")
     _GENERATE_VIDEO(
@@ -719,16 +1023,20 @@ def _run_scail_job(
         cfg=cfg,
         input_idx=None,
         replace_flag=bool(replace_flag),
-        additional_task_input=None,
+        additional_task_input=additional_task_input,
     )
+
     if progress is not None:
         progress(0.95, desc="Finalizing output")
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
+    if progress is not None:
+        progress(0.98, desc="Preparing video for display")
+    display_file = _prepare_output_for_gradio(save_file)
     if progress is not None:
         progress(1.0, desc="Done")
-    return str(save_file)
+    return display_file
 
 
 @spaces.GPU(duration=_duration_for_job, size=GPU_SIZE)
@@ -751,6 +1059,10 @@ def generate_from_example(
             raise RuntimeError(f"Example is missing from this checkout: {example_name}")
         example = examples[example_name]
         target_w, target_h = [int(v) for v in str(target_size).split("x")]
+        refs = tuple(
+            ReferencePair(ref.label, _abs(ref.image), _abs(ref.mask_image))
+            for ref in example.additional_refs
+        )
         output = _run_scail_job(
             _abs(example.image),
             _abs(example.mask_image),
@@ -766,6 +1078,53 @@ def generate_from_example(
             seed,
             segment_len,
             segment_overlap,
+            additional_refs=refs,
+            progress=progress,
+        )
+        return output, "Done."
+    except Exception:
+        logging.exception("Generation failed")
+        return None, traceback.format_exc()
+
+
+@spaces.GPU(duration=_duration_for_job, size=GPU_SIZE)
+def generate_from_pack(
+    pack,
+    prompt,
+    sample_steps,
+    guide_scale,
+    sample_shift,
+    seed,
+    target_size,
+    segment_len,
+    segment_overlap,
+    progress=gr.Progress(track_tqdm=True),
+):
+    try:
+        progress(0.0, desc="Checking validated pack")
+        if not pack:
+            raise RuntimeError("Validate an Advanced Pack before generating.")
+        target_w, target_h = [int(v) for v in str(target_size).split("x")]
+        refs = tuple(
+            ReferencePair(ref["label"], ref["image"], ref["mask_image"])
+            for ref in pack.get("additional_refs", [])
+        )
+        output = _run_scail_job(
+            pack["image"],
+            pack["mask_image"],
+            pack["pose"],
+            pack["mask_video"],
+            prompt,
+            pack.get("replace_flag", False),
+            target_h,
+            target_w,
+            sample_steps,
+            guide_scale,
+            sample_shift,
+            seed,
+            segment_len,
+            segment_overlap,
+            additional_refs=refs,
             progress=progress,
         )
         return output, "Done."
@@ -827,10 +1186,27 @@ def generate_from_uploads(
         return None, traceback.format_exc()
 
 
+def _reference_gallery(example: PreparedExample):
+    items = [
+        (_abs(example.image), "Primary reference"),
+        (_abs(example.mask_image), "Primary mask"),
+    ]
+    for ref in example.additional_refs:
+        items.append((_abs(ref.image), ref.label))
+        items.append((_abs(ref.mask_image), f"{ref.label} mask"))
+    return items
+
+
+def _reference_note(example: PreparedExample) -> str:
+    if not example.additional_refs:
+        return "Single-reference example."
+    return f"Multi-reference example: {len(example.additional_refs)} additional reference pair(s) are passed to SCAIL-2."
+
+
 def load_example_preview(example_name):
     examples = _existing_examples()
     if example_name not in examples:
-        return None, None, None, None, "", "animation"
+        return None, None, None, None, "", "animation", [], "Example not available."
     example = examples[example_name]
     mode = "replacement" if example.replace_flag else "animation"
     return (
@@ -840,6 +1216,8 @@ def load_example_preview(example_name):
         _abs(example.mask_video),
         example.prompt,
         mode,
+        _reference_gallery(example),
+        _reference_note(example),
     )
 
 
@@ -860,49 +1238,65 @@ def _startup_message():
         return str(exc)
 
 
+def _sampling_controls(prefix: str = ""):
+    with gr.Row():
+        steps = gr.Slider(4, 40, value=8, step=1, label=f"{prefix}Steps".strip())
+        cfg = gr.Slider(1.0, 8.0, value=DEFAULT_GUIDE_SCALE, step=0.1, label=f"{prefix}CFG".strip())
+        shift = gr.Slider(1.0, 6.0, value=DEFAULT_SHIFT, step=0.1, label=f"{prefix}Shift".strip())
+    with gr.Row():
+        seed = gr.Number(value=42, precision=0, label=f"{prefix}Seed".strip())
+        target_size = gr.Dropdown(
+            ["896x512", "512x896", "1280x704", "704x1280"],
+            value=f"{DEFAULT_TARGET_W}x{DEFAULT_TARGET_H}",
+            label=f"{prefix}Target size".strip(),
+        )
+        segment_len = gr.Number(value=DEFAULT_SEGMENT_LEN, precision=0, label=f"{prefix}Segment length".strip())
+        segment_overlap = gr.Number(value=DEFAULT_SEGMENT_OVERLAP, precision=0, label=f"{prefix}Segment overlap".strip())
+    return steps, cfg, shift, seed, target_size, segment_len, segment_overlap
+
+
 def build_ui():
     examples = _existing_examples()
     example_names = list(examples.keys())
     default_example = example_names[0] if example_names else None
-    default_preview = load_example_preview(default_example) if default_example else (None, None, None, None, "", "animation")
+    default_preview = (
+        load_example_preview(default_example)
+        if default_example
+        else (None, None, None, None, "", "animation", [], "No prepared examples found.")
+    )
 
-    with gr.Blocks(title="SCAIL-2 ZeroGPU Demo") as demo:
+    with gr.Blocks(title="SCAIL-2 Character Animation Demo") as demo:
         gr.Markdown(
-            "# SCAIL-2 Demo\n"
-            "V1 inference-only Space. It runs SCAIL-2 from prepared inputs: reference image, "
-            "reference mask, driving/rendered video, and driving mask video. SCAIL-Pose/SAM3 "
-            "preprocessing is intentionally not imported in this version.\n\n"
-            "On ZeroGPU, each generation should be treated as a cold start: the GPU can be "
-            "released between calls, so the 14B stack may need to be loaded again before "
-            "sampling."
+            "# SCAIL-2 Character Animation Demo\n"
+            "Try SCAIL-2 from curated examples or from already-prepared custom inputs. "
+            "The multi-reference example from the repo is handled as a prepared example: "
+            "its additional references are passed to the model automatically."
         )
         startup = gr.Textbox(value=_startup_message(), label="Startup status", interactive=False)
 
         with gr.Tab("Prepared Examples"):
             with gr.Row():
-                example_dropdown = gr.Dropdown(
-                    choices=example_names,
-                    value=default_example,
-                    label="Example",
-                )
+                example_dropdown = gr.Dropdown(choices=example_names, value=default_example, label="Example")
                 mode_view = gr.Textbox(value=default_preview[5], label="Mode", interactive=False)
-            with gr.Row():
-                ref_preview = gr.Image(value=default_preview[0], label="Reference", interactive=False)
-                driving_preview = gr.Video(value=default_preview[1], label="Driving / rendered video")
-            with gr.Row():
-                ref_mask_preview = gr.Image(value=default_preview[2], label="Reference mask", interactive=False)
-                driving_mask_preview = gr.Video(value=default_preview[3], label="Driving mask")
-            prompt = gr.Textbox(value=default_preview[4], label="Prompt", lines=3)
+            reference_note = gr.Markdown(default_preview[7])
+            with gr.Accordion("Input preview", open=False):
+                with gr.Row():
+                    ref_preview = gr.Image(value=default_preview[0], label="Primary reference", interactive=False)
+                    driving_preview = gr.Video(value=default_preview[1], label="Driving / rendered video")
+                with gr.Row():
+                    ref_mask_preview = gr.Image(value=default_preview[2], label="Primary mask", interactive=False)
+                    driving_mask_preview = gr.Video(value=default_preview[3], label="Driving mask")
+                reference_gallery = gr.Gallery(
+                    value=default_preview[6],
+                    label="Reference set",
+                    columns=4,
+                    height=260,
+                    selected_index=0,
+                    preview=True,
+                )
 
-            with gr.Row():
-                sample_steps = gr.Slider(4, 40, value=8, step=1, label="Steps")
-                guide_scale = gr.Slider(1.0, 8.0, value=DEFAULT_GUIDE_SCALE, step=0.1, label="CFG")
-                sample_shift = gr.Slider(1.0, 6.0, value=DEFAULT_SHIFT, step=0.1, label="Shift")
-            with gr.Row():
-                seed = gr.Number(value=42, precision=0, label="Seed")
-                target_size = gr.Dropdown(["896x512", "512x896", "1280x704", "704x1280"], value=f"{DEFAULT_TARGET_W}x{DEFAULT_TARGET_H}", label="Target size")
-                segment_len = gr.Number(value=DEFAULT_SEGMENT_LEN, precision=0, label="Segment length")
-                segment_overlap = gr.Number(value=DEFAULT_SEGMENT_OVERLAP, precision=0, label="Segment overlap")
+            prompt = gr.Textbox(value=default_preview[4], label="Prompt", lines=3)
+            sample_steps, guide_scale, sample_shift, seed, target_size, segment_len, segment_overlap = _sampling_controls()
 
             run_example = gr.Button("Generate", variant="primary")
             output_video = gr.Video(label="Output")
@@ -911,7 +1305,16 @@ def build_ui():
             example_dropdown.change(
                 load_example_preview,
                 inputs=[example_dropdown],
-                outputs=[ref_preview, driving_preview, ref_mask_preview, driving_mask_preview, prompt, mode_view],
+                outputs=[
+                    ref_preview,
+                    driving_preview,
+                    ref_mask_preview,
+                    driving_mask_preview,
+                    prompt,
+                    mode_view,
+                    reference_gallery,
+                    reference_note,
+                ],
             )
             run_example.click(
                 generate_from_example,
@@ -929,9 +1332,11 @@ def build_ui():
                 outputs=[output_video, status],
             )
 
-        with gr.Tab("Advanced Uploads"):
+        with gr.Tab("Custom Uploads"):
             gr.Markdown(
-                "Upload already-preprocessed inputs. For V1, this tab does not run SAM3 or SCAIL-Pose."
+                "Upload a prepared SCAIL-2 input set: reference image, reference mask, "
+                "driving/rendered video, and driving mask video. This tab is intentionally "
+                "single-reference; use prepared examples for the official multi-reference case."
             )
             with gr.Row():
                 up_image = gr.Image(type="filepath", label="Reference image")
@@ -941,15 +1346,8 @@ def build_ui():
                 up_mask_video = gr.Video(label="Driving mask / replace mask")
             up_mode = gr.Radio(["animation", "replacement"], value="animation", label="Mode")
             up_prompt = gr.Textbox(label="Prompt", lines=3)
-            with gr.Row():
-                up_steps = gr.Slider(4, 40, value=8, step=1, label="Steps")
-                up_cfg = gr.Slider(1.0, 8.0, value=DEFAULT_GUIDE_SCALE, step=0.1, label="CFG")
-                up_shift = gr.Slider(1.0, 6.0, value=DEFAULT_SHIFT, step=0.1, label="Shift")
-            with gr.Row():
-                up_seed = gr.Number(value=42, precision=0, label="Seed")
-                up_target_size = gr.Dropdown(["896x512", "512x896", "1280x704", "704x1280"], value=f"{DEFAULT_TARGET_W}x{DEFAULT_TARGET_H}", label="Target size")
-                up_segment_len = gr.Number(value=DEFAULT_SEGMENT_LEN, precision=0, label="Segment length")
-                up_segment_overlap = gr.Number(value=DEFAULT_SEGMENT_OVERLAP, precision=0, label="Segment overlap")
+            up_steps, up_cfg, up_shift, up_seed, up_target_size, up_segment_len, up_segment_overlap = _sampling_controls()
+
             run_upload = gr.Button("Generate from uploads", variant="primary")
             upload_output = gr.Video(label="Output")
             upload_status = gr.Textbox(label="Run status", lines=8)
@@ -973,6 +1371,116 @@ def build_ui():
                 outputs=[upload_output, upload_status],
             )
 
+        with gr.Tab("Advanced Pack"):
+            gr.Markdown(
+                "Upload a `.zip` pack for multi-reference or multi-character inputs. "
+                "The app validates the file structure, selects one primary reference, and "
+                "passes every other image/mask pair as additional references to SCAIL-2."
+            )
+            with gr.Accordion("Pack format", open=False):
+                gr.Markdown(
+                    "### Canonical zip structure\n"
+                    "Use this layout when one or more characters have several reference views. "
+                    "Each image must have a matching mask with the same stem plus `_mask`.\n\n"
+                    "```text\n"
+                    "scail2_input_pack/\n"
+                    "|-- rendered_v2.mp4\n"
+                    "|-- rendered_mask_v2.mp4\n"
+                    "|-- prompt.txt                  # optional\n"
+                    "|-- metadata.json               # optional\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"
+                    "`-- environment/\n"
+                    "    |-- background.png\n"
+                    "    `-- background_mask.png\n"
+                    "```\n\n"
+                    "### Mask convention\n"
+                    "Colors represent identity slots, not individual views. If `character_0` has "
+                    "front, back, and close-up references, all masks for those views should use the "
+                    "same identity color. A different character gets a different color. The driving "
+                    "mask video should use the same color assignments.\n\n"
+                    "### Mapping to SCAIL-2\n"
+                    "SCAIL-2 receives one primary reference plus a list of additional refs. The parser "
+                    "uses `metadata.json` when a primary is declared. Otherwise it uses "
+                    "`characters/character_0/front.*` or the first available view from `character_0`. "
+                    "All remaining image/mask pairs become additional references.\n\n"
+                    "### Legacy repo-style pack\n"
+                    "The official multi-reference example also uses this flat layout, which is supported:\n\n"
+                    "```text\n"
+                    "ref.png\n"
+                    "ref_mask.jpg\n"
+                    "rendered_v2.mp4\n"
+                    "rendered_mask_v2.mp4\n"
+                    "background.png\n"
+                    "background_mask.png\n"
+                    "character_0.png\n"
+                    "character_0_mask.png\n"
+                    "character_1.png\n"
+                    "character_1_mask.png\n"
+                    "```\n"
+                )
+
+            pack_state = gr.State(None)
+            pack_file = gr.File(
+                label="SCAIL-2 input pack (.zip)",
+                file_types=[".zip"],
+                type="filepath",
+            )
+            validate_pack = gr.Button("Validate pack")
+            pack_summary = gr.Markdown("Upload a pack and validate it before generating.")
+            pack_gallery = gr.Gallery(
+                label="Parsed reference set",
+                columns=4,
+                height=260,
+                selected_index=0,
+                preview=True,
+            )
+            with gr.Row():
+                pack_driving_preview = gr.Video(label="Driving / rendered video")
+                pack_mask_preview = gr.Video(label="Driving mask / replace mask")
+            pack_mode = gr.Textbox(value="animation", label="Mode", interactive=False)
+            pack_prompt = gr.Textbox(label="Prompt", lines=3)
+            pack_steps, pack_cfg, pack_shift, pack_seed, pack_target_size, pack_segment_len, pack_segment_overlap = _sampling_controls()
+            run_pack = gr.Button("Generate from pack", variant="primary")
+            pack_output = gr.Video(label="Output")
+            pack_status = gr.Textbox(label="Run status", lines=8)
+
+            validate_pack.click(
+                validate_input_pack,
+                inputs=[pack_file],
+                outputs=[
+                    pack_state,
+                    pack_summary,
+                    pack_gallery,
+                    pack_driving_preview,
+                    pack_mask_preview,
+                    pack_prompt,
+                    pack_mode,
+                ],
+            )
+            run_pack.click(
+                generate_from_pack,
+                inputs=[
+                    pack_state,
+                    pack_prompt,
+                    pack_steps,
+                    pack_cfg,
+                    pack_shift,
+                    pack_seed,
+                    pack_target_size,
+                    pack_segment_len,
+                    pack_segment_overlap,
+                ],
+                outputs=[pack_output, pack_status],
+            )
+
     return demo
 
 
@@ -983,4 +1491,7 @@ if __name__ == "__main__":
         _prepare_runtime_for_startup()
     if PRELOAD_PIPELINE:
         _prepare_pipeline_for_startup()
-    build_ui().queue(max_size=8).launch()
+    build_ui().queue(max_size=8).launch(
+        allowed_paths=[str(OUTPUT_DIR.resolve())],
+        show_error=True,
+    )