import gc import copy import json import logging import os import shutil import subprocess import sys import tempfile import traceback import uuid import zipfile from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace import gradio as gr import spaces import torch from huggingface_hub import hf_hub_download, snapshot_download logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(levelname)s: %(message)s") ROOT = Path(__file__).resolve().parent def _default_storage_root() -> Path: if os.getenv("SCAIL_STORAGE_ROOT"): return Path(os.environ["SCAIL_STORAGE_ROOT"]) data_mount = Path("/data") if data_mount.exists() and os.access(data_mount, os.W_OK): return data_mount return Path("/tmp") STORAGE_ROOT = _default_storage_root() STAGING_ROOT = Path(os.getenv("SCAIL_STAGING_ROOT", "/tmp")) OUTPUT_DIR = Path(os.getenv("SCAIL_OUTPUT_DIR", str(STORAGE_ROOT / "scail2_outputs"))) OUTPUT_DIR.mkdir(parents=True, exist_ok=True) MODEL_REPO_ID = os.getenv("SCAIL_MODEL_REPO_ID", "zai-org/SCAIL-2") 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", "600")) GPU_DURATION_WARM = int(os.getenv("SCAIL_GPU_DURATION_WARM", "330")) GPU_DURATION_MAX = int(os.getenv("SCAIL_GPU_DURATION_MAX", "1200")) GPU_DURATION_MULTI_CHARACTER_MULTIPLIER = float(os.getenv("SCAIL_GPU_DURATION_MULTI_CHARACTER_MULTIPLIER", "1.0")) 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")) DEFAULT_SEGMENT_OVERLAP = int(os.getenv("SCAIL_SEGMENT_OVERLAP", "5")) DEFAULT_SHIFT = float(os.getenv("SCAIL_SAMPLE_SHIFT", "3.0")) DEFAULT_GUIDE_SCALE = float(os.getenv("SCAIL_GUIDE_SCALE", "5.0")) DEFAULT_SOLVER = os.getenv("SCAIL_SAMPLE_SOLVER", "unipc") 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")) COPY_SOURCE_AUDIO = os.getenv("SCAIL_COPY_SOURCE_AUDIO", "1") == "1" MATCH_SOURCE_FPS_AT_EXPORT = os.getenv("SCAIL_MATCH_SOURCE_FPS_AT_EXPORT", "1") == "1" CONFORM_OUTPUT_TO_SOURCE_FPS = os.getenv("SCAIL_CONFORM_OUTPUT_TO_SOURCE_FPS", "1") == "1" FPS_MATCH_EPSILON = float(os.getenv("SCAIL_FPS_MATCH_EPSILON", "0.05")) AUTO_AVOID_SEGMENT_TRUNCATION = os.getenv("SCAIL_AUTO_AVOID_SEGMENT_TRUNCATION", "1") == "1" MAX_AUTO_SEGMENT_LEN = int(os.getenv("SCAIL_MAX_AUTO_SEGMENT_LEN", "161")) 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 = [ "Wan2.1_VAE.pth", "umt5-xxl/**", CLIP_CKPT_NAME, ] _PIPELINE = None _PIPELINE_KEY = None _ASSET_STATUS = "Assets were not prepared yet." _ASSET_ERROR = None _RUNTIME_STATUS = "Runtime was not prepared yet." _RUNTIME_ERROR = None _PIPELINE_STATUS = "Pipeline was not preloaded." _PIPELINE_ERROR = None _LAST_CONVERTED_SAFETENSORS = None _WAN = None _GENERATE_VIDEO = None _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 image: str mask_image: str pose: str mask_video: str prompt: str replace_flag: bool = False additional_refs: tuple[ReferencePair, ...] = () PREPARED_EXAMPLES = { "Animation 001 - end-to-end": PreparedExample( label="Animation 001 - end-to-end", image="examples/animation_001/ref.jpg", mask_image="examples/animation_001/ref_mask.jpg", pose="examples/animation_001/rendered_v2.mp4", mask_video="examples/animation_001/rendered_mask_v2.mp4", prompt="A young woman is dancing with energetic body movement.", ), "Animation 001 - pose-driven": PreparedExample( label="Animation 001 - pose-driven", image="examples/animation_001_posedriven/ref.jpg", mask_image="examples/animation_001_posedriven/ref_mask.jpg", pose="examples/animation_001_posedriven/rendered_v2.mp4", mask_video="examples/animation_001_posedriven/rendered_mask_v2.mp4", prompt="A young woman is dancing with energetic body movement.", ), "Animation 002 - end-to-end": PreparedExample( label="Animation 002 - end-to-end", image="examples/animation_002/ref.jpg", mask_image="examples/animation_002/ref_mask.jpg", pose="examples/animation_002/rendered_v2.mp4", mask_video="examples/animation_002/rendered_mask_v2.mp4", prompt="A character performs the motion from the driving video.", ), "Animation 003 - multi-reference": PreparedExample( label="Animation 003 - multi-reference", image="examples/animation_003_multi_ref/ref.png", mask_image="examples/animation_003_multi_ref/ref_mask.jpg", 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", image="examples/replace_001/ref.png", mask_image="examples/replace_001/ref_mask.png", pose="examples/replace_001/rendered_v2.mp4", mask_video="examples/replace_001/replace_mask.mp4", prompt=( "A blond white male wearing a black suit, trousers, and leather shoes " "is playing the violin on the street while pedestrians walk past him." ), replace_flag=True, ), } def _abs(path: str | Path) -> str: path = Path(path) if not path.is_absolute(): path = ROOT / path 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(): 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 _character_ids_from_pairs(character_pairs: list[ReferencePair]) -> list[str]: ids = set() for ref in character_pairs: identity = ref.label.split(":", 1)[0].strip() if identity.startswith("character_"): ids.add(identity) return sorted(ids) 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) character_ids = _character_ids_from_pairs(character_pairs) 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) character_count = len(character_ids) estimated_passes = max(1, character_count) if not replace_flag else 1 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, "character_ids": character_ids, "character_count": character_count, "estimated_passes": estimated_passes, "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'])}`", f"- Character slots: `{pack.get('character_count', 0)}`", f"- Estimated generation passes: `{pack.get('estimated_passes', 1)}`", ] 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"): if not (ROOT / rel).exists(): missing.append(rel) if missing: raise RuntimeError( "This app.py is meant to live at the root of the SCAIL-2 repository. " f"Missing: {', '.join(missing)}" ) def _download_safetensors_if_configured() -> Path | None: if not SAFETENSORS_REPO_ID: return None local_dir = Path(os.getenv("SCAIL_SAFETENSORS_CACHE", str(STORAGE_ROOT / "scail2_safetensors"))) local_dir.mkdir(parents=True, exist_ok=True) local_path = local_dir / SAFETENSORS_FILENAME if local_path.exists(): return local_path 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, local_dir=str(local_dir), local_dir_use_symlinks=False, resume_download=True, ) return Path(downloaded) def _find_converted_safetensors(ckpt_dir: Path | None) -> Path | None: candidates = [] env_path = os.getenv("SCAIL_SAFETENSORS_PATH") if env_path: candidates.append(Path(env_path)) if _LAST_CONVERTED_SAFETENSORS is not None: candidates.append(Path(_LAST_CONVERTED_SAFETENSORS)) candidates += [ ROOT / "SCAIL-2.safetensors", ROOT / "models" / "SCAIL-2.safetensors", ROOT / "model.safetensors", 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"] for candidate in candidates: if candidate.exists(): return candidate return _download_safetensors_if_configured() def _copy_file_with_progress(source: Path, dest: Path, description: str) -> Path: source_size = source.stat().st_size chunk_size = int(os.getenv("SCAIL_STAGE_COPY_CHUNK_MB", "64")) * 1024 * 1024 log_every = int(os.getenv("SCAIL_STAGE_COPY_LOG_GB", "1")) * 1024 * 1024 * 1024 dest.parent.mkdir(parents=True, exist_ok=True) if dest.exists() and dest.stat().st_size == source_size: return dest tmp_dest = dest.with_suffix(dest.suffix + ".tmp") copied = tmp_dest.stat().st_size if tmp_dest.exists() else 0 if copied > source_size: tmp_dest.unlink() copied = 0 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) with source.open("rb") as src, tmp_dest.open("ab") as dst: if copied: src.seek(copied) while copied < source_size: chunk = src.read(min(chunk_size, source_size - copied)) if not chunk: raise RuntimeError(f"Unexpected EOF while copying {source}: {copied} of {source_size} bytes") 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) next_log += log_every if tmp_dest.stat().st_size != source_size: raise RuntimeError(f"Copied file size mismatch: {tmp_dest.stat().st_size} != {source_size}") tmp_dest.replace(dest) logging.info("Finished %s: %s", description, dest) return dest def _is_relative_to(path: Path, parent: Path) -> bool: try: path.resolve().relative_to(parent.resolve()) return True except ValueError: return False def _stage_safetensors_for_load(scail_path: Path) -> Path: if not STAGE_SAFETENSORS_FOR_LOAD: return scail_path source = Path(scail_path) if _is_relative_to(source, STAGING_ROOT): return source stage_dir = Path(os.getenv("SCAIL_MODEL_LOAD_CACHE", str(STAGING_ROOT / "scail2_model_load"))) 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") def _download_checkpoint_if_needed(include_original_dit: bool = False) -> Path: env_dir = os.getenv("SCAIL_CKPT_DIR") if env_dir: ckpt_dir = Path(env_dir) if not ckpt_dir.exists(): raise RuntimeError(f"SCAIL_CKPT_DIR does not exist: {ckpt_dir}") return ckpt_dir local_dir = Path(os.getenv("SCAIL_CKPT_CACHE", str(STORAGE_ROOT / "scail2_ckpt"))) has_base_assets = ( (local_dir / "Wan2.1_VAE.pth").exists() and (local_dir / "umt5-xxl").exists() and (local_dir / CLIP_CKPT_NAME).exists() ) if has_base_assets: return local_dir if include_original_dit: 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), local_dir_use_symlinks=False, resume_download=True, allow_patterns=BASE_ALLOW_PATTERNS, ) return local_dir def _download_original_dit_for_conversion() -> Path: env_dir = os.getenv("SCAIL_ORIGINAL_DIT_DIR") if env_dir: original_dir = Path(env_dir) original_path = original_dir / ORIGINAL_DIT_REL_PATH if not original_path.exists(): raise RuntimeError(f"SCAIL_ORIGINAL_DIT_DIR is missing {ORIGINAL_DIT_REL_PATH}: {original_dir}") return original_dir local_dir = Path(os.getenv("SCAIL_ORIGINAL_DIT_CACHE", str(STAGING_ROOT / "scail2_original_dit"))) original_path = local_dir / ORIGINAL_DIT_REL_PATH if original_path.exists(): return 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), local_dir_use_symlinks=False, resume_download=True, allow_patterns=[ORIGINAL_DIT_REL_PATH], ) return local_dir def _prepare_assets_for_runtime() -> str: global _ASSET_STATUS, _ASSET_ERROR try: ckpt_dir = _download_checkpoint_if_needed(include_original_dit=False) scail_path = _find_converted_safetensors(ckpt_dir) if scail_path is None and AUTO_CONVERT: original_dir = _download_original_dit_for_conversion() 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_PATH or SCAIL_SAFETENSORS_REPO_ID." ) else: _ASSET_STATUS = f"Assets ready. Base checkpoint: {ckpt_dir}. Converted DiT safetensors: {scail_path}." _ASSET_ERROR = None except Exception: _ASSET_ERROR = traceback.format_exc() _ASSET_STATUS = "Asset preparation failed. See the traceback below." logging.exception("Asset preparation failed") return _ASSET_STATUS if _ASSET_ERROR is None else _ASSET_STATUS + "\n\n" + _ASSET_ERROR def _maybe_convert_checkpoint(ckpt_dir: Path, scail_path: Path | None) -> Path: global _LAST_CONVERTED_SAFETENSORS if scail_path is not None: return scail_path if not AUTO_CONVERT: raise RuntimeError( "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"))) persistent_path = persistent_dir / "SCAIL-2.safetensors" if persistent_path.exists(): return persistent_path 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" if save_path.exists(): _LAST_CONVERTED_SAFETENSORS = save_path if save_path != persistent_path: _copy_file_with_progress(save_path, persistent_path, "Persisting converted SCAIL-2 safetensors to storage") return save_path logging.info("Converting checkpoint to safetensors: %s", save_path) convert_env = os.environ.copy() convert_env["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1" subprocess.run( [ sys.executable, str(ROOT / "convert.py"), "--scail-dir", str(ckpt_dir), "--save-path", str(save_path), ], check=True, cwd=str(ROOT), env=convert_env, ) _LAST_CONVERTED_SAFETENSORS = save_path if save_path != persistent_path: _copy_file_with_progress(save_path, persistent_path, "Persisting converted SCAIL-2 safetensors to storage") return save_path def _install_attention_patch(): import wan.modules.attention as attention_mod hf_flash_attn2 = None try: from kernels import get_kernel hf_flash_attn2 = get_kernel("kernels-community/flash-attn2", version=2) logging.info("Using kernels-community/flash-attn2 through HF Kernels.") except Exception as exc: if torch.cuda.is_available(): device_name = torch.cuda.get_device_name(0) capability = torch.cuda.get_device_capability(0) else: device_name = "no cuda" capability = None logging.warning("Could not initialize HF Kernels flash-attn2: %r", exc) logging.warning( "Attention fallback environment: torch=%s cuda=%s device=%s capability=%s", torch.__version__, torch.version.cuda, device_name, capability, ) def patched_flash_attention( q, k, v, q_lens=None, k_lens=None, dropout_p=0.0, softmax_scale=None, q_scale=None, causal=False, window_size=(-1, -1), deterministic=False, dtype=torch.bfloat16, version=None, ): half_dtypes = (torch.float16, torch.bfloat16) b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype def half(x): return x if x.dtype in half_dtypes else x.to(dtype) if hf_flash_attn2 is not None and q.device.type == "cuda": if q_lens is None: q_var = half(q.flatten(0, 1)) q_lens_t = torch.full((b,), lq, dtype=torch.int32, device=q.device) else: q_lens_t = q_lens.to(device=q.device, dtype=torch.int32) q_var = half(torch.cat([u[: int(n)] for u, n in zip(q, q_lens_t)])) if k_lens is None: k_var = half(k.flatten(0, 1)) v_var = half(v.flatten(0, 1)) k_lens_t = torch.full((b,), lk, dtype=torch.int32, device=k.device) else: k_lens_t = k_lens.to(device=k.device, dtype=torch.int32) k_var = half(torch.cat([u[: int(n)] for u, n in zip(k, k_lens_t)])) v_var = half(torch.cat([u[: int(n)] for u, n in zip(v, k_lens_t)])) q_var = q_var.to(v_var.dtype) k_var = k_var.to(v_var.dtype) if q_scale is not None: q_var = q_var * q_scale 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, k=k_var, v=v_var, cu_seqlens_q=cu_q, cu_seqlens_k=cu_k, max_seqlen_q=lq, max_seqlen_k=lk, dropout_p=dropout_p, softmax_scale=softmax_scale, causal=causal, window_size=window_size, deterministic=deterministic, ) if isinstance(out, tuple): out = out[0] return out.unflatten(0, (b, lq)).type(out_dtype) except Exception as exc: logging.warning("HF Kernels flash-attn2 failed, falling back to SDPA: %s", exc) if q_lens is not None and not torch.all(q_lens == lq): logging.warning("SDPA fallback ignores variable q_lens; demo batch size should stay at 1.") if k_lens is not None and not torch.all(k_lens == lk): logging.warning("SDPA fallback ignores variable k_lens; demo batch size should stay at 1.") q_sdpa = q.transpose(1, 2).to(dtype) k_sdpa = k.transpose(1, 2).to(dtype) v_sdpa = v.transpose(1, 2).to(dtype) out = torch.nn.functional.scaled_dot_product_attention( q_sdpa, k_sdpa, v_sdpa, attn_mask=None, dropout_p=dropout_p, is_causal=causal, scale=softmax_scale, ) return out.transpose(1, 2).contiguous().type(out_dtype) attention_mod.flash_attention = patched_flash_attention for module_name in ( "wan.modules.clip", "wan.modules.model", "wan.modules.model_scail", "wan.modules.model_scail2", ): try: module = __import__(module_name, fromlist=["flash_attention"]) if hasattr(module, "flash_attention"): module.flash_attention = patched_flash_attention logging.info("Patched %s.flash_attention", module_name) except Exception as exc: logging.warning("Could not patch %s.flash_attention: %s", module_name, exc) def _import_runtime(): global _WAN, _GENERATE_VIDEO, _SCAIL_CONFIGS, _SCAIL_CONFIG_PATHS if _WAN is not None: return _require_repo_layout() if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) import wan from generate import generate_video from wan.configs import SCAIL_CONFIGS, SCAIL_CONFIG_PATHS _install_attention_patch() _WAN = wan _GENERATE_VIDEO = generate_video _SCAIL_CONFIGS = SCAIL_CONFIGS _SCAIL_CONFIG_PATHS = SCAIL_CONFIG_PATHS def _prepare_runtime_for_startup() -> str: global _RUNTIME_STATUS, _RUNTIME_ERROR try: _import_runtime() _RUNTIME_STATUS = "Runtime ready. Attention backend has been initialized at startup." _RUNTIME_ERROR = None except Exception: _RUNTIME_ERROR = traceback.format_exc() _RUNTIME_STATUS = "Runtime preparation failed. See the traceback below." logging.exception("Runtime preparation failed") return _RUNTIME_STATUS if _RUNTIME_ERROR is None else _RUNTIME_STATUS + "\n\n" + _RUNTIME_ERROR def _get_pipeline(): global _PIPELINE, _PIPELINE_KEY _import_runtime() ckpt_dir = _download_checkpoint_if_needed(include_original_dit=False) scail_path = _find_converted_safetensors(ckpt_dir) if scail_path is None and AUTO_CONVERT: original_dir = _download_original_dit_for_conversion() 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(): config_path = ROOT / config_path lora_path = os.getenv("SCAIL_LORA_PATH") or None lora_alpha = float(os.getenv("SCAIL_LORA_ALPHA", "1.0")) key = (str(ckpt_dir), str(scail_load_path), str(config_path), lora_path, lora_alpha) if _PIPELINE is not None and _PIPELINE_KEY == key: return _PIPELINE logging.info("Loading SCAIL-2 pipeline.") cfg = _SCAIL_CONFIGS[MODEL_NAME] _PIPELINE = _WAN.SCAIL2Pipeline( config=cfg, checkpoint_dir=str(ckpt_dir), scail_safetensors_path=str(scail_load_path), scail_config_path=str(config_path), device_id=0, rank=0, t5_fsdp=False, dit_fsdp=False, use_usp=False, t5_cpu=False, lora_path=lora_path, lora_alpha=lora_alpha, ) _PIPELINE_KEY = key return _PIPELINE def _prepare_pipeline_for_startup() -> str: global _PIPELINE_STATUS, _PIPELINE_ERROR try: _get_pipeline() _PIPELINE_STATUS = "Pipeline preloaded at startup." _PIPELINE_ERROR = None except Exception: _PIPELINE_ERROR = traceback.format_exc() _PIPELINE_STATUS = "Pipeline preload failed. See the traceback below." logging.exception("Pipeline preload failed") 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 _get_ffmpeg_exe() -> str | None: try: import imageio_ffmpeg return imageio_ffmpeg.get_ffmpeg_exe() except Exception as exc: logging.warning("Could not locate ffmpeg: %s", exc) return None def _video_fps(video_path: str | Path) -> float | None: try: import cv2 capture = cv2.VideoCapture(str(video_path)) fps = float(capture.get(cv2.CAP_PROP_FPS) or 0.0) capture.release() if fps > 0: return fps except Exception as exc: logging.warning("Could not read FPS for %s: %s", video_path, exc) return None def _video_frame_count(video_path: str | Path) -> int | None: try: import cv2 capture = cv2.VideoCapture(str(video_path)) frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT) or 0) capture.release() if frame_count > 0: return frame_count except Exception as exc: logging.warning("Could not read frame count for %s: %s", video_path, exc) return None def _effective_segment_len(requested_segment_len: int, segment_overlap: int, source_video_path: str | Path) -> int: requested_segment_len = int(requested_segment_len) segment_overlap = int(segment_overlap) if not AUTO_AVOID_SEGMENT_TRUNCATION: return requested_segment_len frame_count = _video_frame_count(source_video_path) if frame_count is None or frame_count <= requested_segment_len: return requested_segment_len if frame_count <= MAX_AUTO_SEGMENT_LEN: effective = max(frame_count, segment_overlap + 1) logging.info( "Increasing segment_len to avoid SCAIL-2 tail truncation: requested=%s frame_count=%s effective=%s", requested_segment_len, frame_count, effective, ) return effective logging.warning( "Source has %s frames, but segment_len=%s and max_auto_segment_len=%s. " "SCAIL-2 may drop the final partial segment unless you increase segment_len or trim/resample the input.", frame_count, requested_segment_len, MAX_AUTO_SEGMENT_LEN, ) return requested_segment_len def _generation_config_for_source_video(source_video_path: str | Path): cfg = copy.deepcopy(_SCAIL_CONFIGS[MODEL_NAME]) if not MATCH_SOURCE_FPS_AT_EXPORT: return cfg source_fps = _video_fps(source_video_path) if source_fps is None: logging.warning("Could not detect source FPS; keeping default SCAIL export FPS: %s", cfg.sample_fps) return cfg cfg.sample_fps = source_fps logging.info("Using source FPS for generated video export: %.3f", source_fps) return cfg def _conform_output_to_source_fps(video_path: str | Path, source_video_path: str | Path) -> Path: video_path = Path(video_path) source_video_path = Path(source_video_path) if not CONFORM_OUTPUT_TO_SOURCE_FPS: return video_path if not video_path.exists() or not source_video_path.exists(): return video_path source_fps = _video_fps(source_video_path) output_fps = _video_fps(video_path) if not source_fps or not output_fps: return video_path if abs(source_fps - output_fps) <= FPS_MATCH_EPSILON: return video_path ffmpeg = _get_ffmpeg_exe() if ffmpeg is None: return video_path speed_multiplier = output_fps / source_fps output_path = video_path.with_name(f"{video_path.stem}_fps{source_fps:.3f}{video_path.suffix}") command = [ ffmpeg, "-y", "-i", str(video_path), "-an", "-vf", f"setpts={speed_multiplier:.10f}*PTS", "-r", f"{source_fps:.6f}", "-c:v", "libx264", "-preset", "veryfast", "-crf", "18", "-pix_fmt", "yuv420p", str(output_path), ] try: subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) except subprocess.CalledProcessError as exc: stderr = exc.stderr[-2000:] if exc.stderr else "" logging.warning("FPS conform failed; returning generated video at original output FPS. stderr=%s", stderr) return video_path if output_path.exists() and output_path.stat().st_size > 0: logging.info( "Conformed generated video FPS to source: output_fps=%.3f source_fps=%.3f file=%s", output_fps, source_fps, output_path, ) return output_path return video_path def _copy_source_audio_to_output(video_path: str | Path, source_video_path: str | Path) -> Path: video_path = Path(video_path) source_video_path = Path(source_video_path) if not COPY_SOURCE_AUDIO: return video_path if not video_path.exists() or not source_video_path.exists(): return video_path ffmpeg = _get_ffmpeg_exe() if ffmpeg is None: return video_path probe = subprocess.run( [ffmpeg, "-hide_banner", "-i", str(source_video_path)], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, ) if "Audio:" not in probe.stderr: logging.info("Source video has no audio stream; keeping generated video silent.") return video_path output_path = video_path.with_name(f"{video_path.stem}_audio{video_path.suffix}") command = [ ffmpeg, "-y", "-i", str(video_path), "-i", str(source_video_path), "-map", "0:v:0", "-map", "1:a:0?", "-c:v", "copy", "-c:a", "aac", "-shortest", str(output_path), ] try: subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) except subprocess.CalledProcessError as exc: stderr = exc.stderr[-2000:] if exc.stderr else "" logging.warning("Audio remux failed; returning generated video without source audio. stderr=%s", stderr) return video_path if output_path.exists() and output_path.stat().st_size > 0: logging.info("Copied source audio onto generated video: %s", output_path) return output_path return video_path def _duration_base_seconds() -> int: if _PIPELINE is None: return int(os.getenv("SCAIL_GPU_DURATION", str(GPU_DURATION_COLD))) return int(os.getenv("SCAIL_GPU_DURATION", str(GPU_DURATION_WARM))) def _estimated_passes_from_duration_args(args, kwargs) -> int: candidates = list(args) + list(kwargs.values()) for value in candidates: if isinstance(value, dict) and "estimated_passes" in value: try: return max(1, int(value.get("estimated_passes") or 1)) except Exception: return 1 return 1 def _duration_for_job(*args, **kwargs): base = _duration_base_seconds() passes = _estimated_passes_from_duration_args(args, kwargs) if passes <= 1: return base duration = int(base * passes * GPU_DURATION_MULTI_CHARACTER_MULTIPLIER) duration = max(base, duration) duration = min(duration, GPU_DURATION_MAX) logging.info( "ZeroGPU duration estimate: base=%ss passes=%s multiplier=%s duration=%ss", base, passes, GPU_DURATION_MULTI_CHARACTER_MULTIPLIER, duration, ) return duration def _run_scail_job( image_path, mask_image_path, pose_path, mask_video_path, prompt, replace_flag, target_h, target_w, sample_steps, guide_scale, sample_shift, seed, segment_len, segment_overlap, additional_refs: tuple[ReferencePair, ...] = (), progress=None, ): if progress is not None: progress(0.02, desc="Loading SCAIL-2 pipeline") pipeline = _get_pipeline() cfg = _generation_config_for_source_video(pose_path) save_file = OUTPUT_DIR / f"scail2_{uuid.uuid4().hex}.mp4" effective_segment_len = _effective_segment_len(segment_len, segment_overlap, pose_path) if progress is not None: progress(0.12, desc="Preparing inputs") args = SimpleNamespace( target_h=int(target_h), target_w=int(target_w), sample_shift=float(sample_shift), sample_solver=DEFAULT_SOLVER, segment_len=effective_segment_len, segment_overlap=int(segment_overlap), sample_steps=int(sample_steps), sample_guide_scale=float(guide_scale), base_seed=int(seed), offload_model=True, save_file=str(save_file), save_dir=str(OUTPUT_DIR), 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( pipeline, prompt or "", str(image_path), str(mask_image_path), str(pose_path), str(mask_video_path), args, device=0, rank=0, cfg=cfg, input_idx=None, replace_flag=bool(replace_flag), 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 CONFORM_OUTPUT_TO_SOURCE_FPS: if progress is not None: progress(0.96, desc="Matching source FPS") save_file = _conform_output_to_source_fps(save_file, pose_path) if COPY_SOURCE_AUDIO: if progress is not None: progress(0.97, desc="Restoring source audio") save_file = _copy_source_audio_to_output(save_file, pose_path) 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 display_file @spaces.GPU(duration=_duration_for_job, size=GPU_SIZE) def generate_from_example( example_name, 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 example inputs") examples = _existing_examples() if example_name not in examples: 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), _abs(example.pose), _abs(example.mask_video), prompt, example.replace_flag, target_h, target_w, sample_steps, guide_scale, sample_shift, 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." except Exception: logging.exception("Generation failed") return None, traceback.format_exc() @spaces.GPU(duration=_duration_for_job, size=GPU_SIZE) def generate_from_uploads( image, mask_image, pose_video, mask_video, prompt, mode, 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 uploaded inputs") required = { "reference image": image, "reference mask": mask_image, "driving/rendered video": pose_video, "driving mask video": mask_video, } missing = [name for name, value in required.items() if value is None] if missing: raise RuntimeError("Missing required input(s): " + ", ".join(missing)) target_w, target_h = [int(v) for v in str(target_size).split("x")] output = _run_scail_job( image, mask_image, pose_video, mask_video, prompt, mode == "replacement", target_h, target_w, sample_steps, guide_scale, sample_shift, seed, segment_len, segment_overlap, progress=progress, ) return output, "Done." except Exception: logging.exception("Generation failed") 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", [], "Example not available." example = examples[example_name] mode = "replacement" if example.replace_flag else "animation" return ( _abs(example.image), _abs(example.pose), _abs(example.mask_image), _abs(example.mask_video), example.prompt, mode, _reference_gallery(example), _reference_note(example), ) def _startup_message(): try: _require_repo_layout() examples = _existing_examples() if not examples: return "Repo layout detected, but no prepared examples were found." return ( f"Ready. Found {len(examples)} prepared example(s). " f"Storage root: {STORAGE_ROOT}. Staging root: {STAGING_ROOT}. " f"Output dir: {OUTPUT_DIR}.\n\n" f"{_ASSET_STATUS}\n\n{_RUNTIME_STATUS}\n\n{_PIPELINE_STATUS}\n\n" "Attention backend: HF Kernels flash-attn2 when available, otherwise SDPA." ) except Exception as exc: 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", [], "No prepared examples found.") ) with gr.Blocks(title="SCAIL-2 Character Animation Demo") as demo: gr.Markdown( "# 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") mode_view = gr.Textbox(value=default_preview[5], label="Mode", interactive=False) 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, ) 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") status = gr.Textbox(label="Run status", lines=8) example_dropdown.change( load_example_preview, inputs=[example_dropdown], outputs=[ ref_preview, driving_preview, ref_mask_preview, driving_mask_preview, prompt, mode_view, reference_gallery, reference_note, ], ) run_example.click( generate_from_example, inputs=[ example_dropdown, prompt, sample_steps, guide_scale, sample_shift, seed, target_size, segment_len, segment_overlap, ], outputs=[output_video, status], ) with gr.Tab("Custom Uploads"): gr.Markdown( "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") up_mask_image = gr.Image(type="filepath", label="Reference mask") with gr.Row(): up_pose_video = gr.Video(label="Driving / rendered video") 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) 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) run_upload.click( generate_from_uploads, inputs=[ up_image, up_mask_image, up_pose_video, up_mask_video, up_prompt, up_mode, up_steps, up_cfg, up_shift, up_seed, up_target_size, up_segment_len, up_segment_overlap, ], 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 if __name__ == "__main__": if os.getenv("SCAIL_PRELOAD_ASSETS", "1") == "1": _prepare_assets_for_runtime() if os.getenv("SCAIL_PRELOAD_RUNTIME", "1") == "1": _prepare_runtime_for_startup() if PRELOAD_PIPELINE: _prepare_pipeline_for_startup() build_ui().queue(max_size=8).launch( allowed_paths=[str(OUTPUT_DIR.resolve())], show_error=True, )