app.py

import gc
import copy
import logging
import os
import shutil
import subprocess
import sys
import time
import tempfile
import traceback
import uuid
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", "X-HighVoltage-X/SCAIL-2-14B"
)
SAFETENSORS_FILENAME = os.getenv(
    "SCAIL_SAFETENSORS_FILENAME", "SCAIL-2-14B.safetensors"
)
MODEL_NAME = os.getenv("SCAIL_MODEL_NAME", "SCAIL-14B")
GPU_DURATION_COLD = int(os.getenv("SCAIL_GPU_DURATION_COLD", "600"))

GPU_DURATION_WARM = int(os.getenv("SCAIL_GPU_DURATION_WARM", "600"))

DEFAULT_SOLVER = os.getenv("SCAIL_SAMPLE_SOLVER", "unipc")
PRELOAD_PIPELINE = os.getenv("SCAIL_PRELOAD_PIPELINE", "1") == "1"
STAGE_SAFETENSORS_FOR_LOAD = os.getenv("SCAIL_STAGE_SAFETENSORS_FOR_LOAD", "1") == "1"
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"
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
_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
    negative_prompt: str = ""
    replace_flag: bool = False
    additional_refs: tuple[ReferencePair, ...] = ()


def _prepare_examples():
    return {
        "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,
        ),
    }


PREPARED_EXAMPLES = _prepare_examples()


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


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))

    candidates.append(ROOT / SAFETENSORS_FILENAME)
    candidates.append(
        Path(os.getenv("SCAIL_CONVERTED_DIR", str(STORAGE_ROOT / "scail2_converted")))
        / SAFETENSORS_FILENAME
    )

    if ckpt_dir is not None:
        candidates.append(ckpt_dir / SAFETENSORS_FILENAME)

    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() -> 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

    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 _prepare_assets_for_runtime() -> str:
    global _ASSET_STATUS, _ASSET_ERROR
    try:
        ckpt_dir = _download_checkpoint_if_needed()
        scail_path = _find_converted_safetensors(ckpt_dir)
        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 _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()
    scail_path = _find_converted_safetensors(ckpt_dir)
    if scail_path is None:
        raise RuntimeError(
            "Converted SCAIL-2 safetensors file was not found. Please check the model repository and filename."
        )

    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 with model %s.", scail_load_path)
    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,
        # TODO: True
        init_on_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 _trim_video_to_frames(video_path: Path, max_frames: int) -> Path:
    ffmpeg = _get_ffmpeg_exe()
    if ffmpeg is None:
        return video_path

    output_path = video_path.with_name(
        f"{video_path.stem}_sync{max_frames}{video_path.suffix}"
    )
    if output_path.exists():
        return output_path

    command = [
        ffmpeg,
        "-y",
        "-i",
        str(video_path),
        "-frames:v",
        str(max_frames),
        "-c:v",
        "libx264",
        "-preset",
        "veryfast",
        "-crf",
        "18",
        "-pix_fmt",
        "yuv420p",
        "-c:a",
        "copy",
        str(output_path),
    ]

    try:
        subprocess.run(
            command,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        if output_path.exists() and output_path.stat().st_size > 0:
            logging.info("Trimmed video %s to %s frames", video_path, max_frames)
            return output_path
    except subprocess.CalledProcessError as exc:
        logging.warning(
            "Failed to trim video %s: %s",
            video_path,
            exc.stderr[-2000:] if exc.stderr else "",
        )
    return video_path


def _sync_video_lengths(pose_path: Path, mask_video_path: Path) -> tuple[Path, Path]:
    pose_frames = _video_frame_count(pose_path)
    mask_frames = _video_frame_count(mask_video_path)

    if pose_frames is None or mask_frames is None:
        return pose_path, mask_video_path

    if pose_frames != mask_frames:
        min_frames = min(pose_frames, mask_frames)
        logging.info(
            "Mismatched frame counts detected (pose: %s, mask: %s). Syncing to %s frames.",
            pose_frames,
            mask_frames,
            min_frames,
        )
        if pose_frames > min_frames:
            pose_path = _trim_video_to_frames(pose_path, min_frames)
        if mask_frames > min_frames:
            mask_video_path = _trim_video_to_frames(mask_video_path, min_frames)

    return pose_path, mask_video_path


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_for_job(*args, **kwargs):
    if _PIPELINE is None:
        duration = int(os.getenv("SCAIL_GPU_DURATION", str(GPU_DURATION_COLD)))
    else:
        duration = int(os.getenv("SCAIL_GPU_DURATION", str(GPU_DURATION_WARM)))
    logging.info("ZeroGPU duration estimate: %ss", duration)
    return duration


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,
        example.negative_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(1, 40, value=4, step=1, label=f"{prefix}Steps".strip())
        cfg = gr.Slider(1.0, 8.0, value=5.0, step=0.1, label=f"{prefix}CFG".strip())
        shift = gr.Slider(1.0, 6.0, value=3.0, 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="512x896",
            label=f"{prefix}Target size".strip(),
        )
        segment_len = gr.Number(
            value=81,
            precision=0,
            label=f"{prefix}Segment length".strip(),
        )
        segment_overlap = gr.Number(
            value=5,
            precision=0,
            label=f"{prefix}Segment overlap".strip(),
        )
    return steps, cfg, shift, seed, target_size, segment_len, segment_overlap


#@spaces.GPU(duration=_duration_for_job, size="xlarge")
#@spaces.GPU(duration=200, size="xlarge")
@spaces.GPU(duration=149, size="xlarge")
def run_generation(
    # A dictionary to hold all the inputs from any of the tabs
    inputs: dict,
    progress=gr.Progress(track_tqdm=True),
):
    try:
        task_start_time = time.perf_counter()
        logging.info("=" * 60)
        logging.info("Starting new generation job")
        
        progress(0.0, desc="Checking inputs")

        # Extract common parameters
        prompt = inputs["prompt"]
        negative_prompt = inputs.get("negative_prompt", "")
        sample_steps = inputs["sample_steps"]
        guide_scale = inputs["guide_scale"]
        sample_shift = inputs["sample_shift"]
        seed = inputs["seed"]
        target_size = inputs["target_size"]
        segment_len = inputs["segment_len"]
        segment_overlap = inputs["segment_overlap"]
        
        # Determine the source of the inputs and prepare them
        if "example_name" in inputs:
            # From "Prepared Examples" tab
            example_name = inputs["example_name"]
            examples = _existing_examples()
            if example_name not in examples:
                raise RuntimeError(
                    f"Example is missing from this checkout: {example_name}"
                )
            example = examples[example_name]
            image_path = _abs(example.image)
            mask_image_path = _abs(example.mask_image)
            pose_path = _abs(example.pose)
            mask_video_path = _abs(example.mask_video)
            replace_flag = example.replace_flag
            additional_refs = tuple(
                ReferencePair(ref.label, _abs(ref.image), _abs(ref.mask_image))
                for ref in example.additional_refs
            )
        else:
            # From "Custom Uploads" tab
            image_path = inputs["image"]
            mask_image_path = inputs["mask_image"]
            pose_path = inputs["pose_video"]
            mask_video_path = inputs["mask_video"]
            mode = inputs["mode"]
            replace_flag = mode == "replacement"
            additional_refs = ()
            required = {
                "reference image": image_path,
                "reference mask": mask_image_path,
                "driving/rendered video": pose_path,
                "driving mask video": mask_video_path,
            }
            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")]

        progress(0.01, desc="Syncing video lengths")
        pose_path, mask_video_path = _sync_video_lengths(
            Path(pose_path), Path(mask_video_path)
        )
        pose_path, mask_video_path = str(pose_path), str(mask_video_path)

        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
        )

        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),
            # TODO: True
            offload_model=False,
            save_file=str(save_file),
            save_dir=str(OUTPUT_DIR),
            prompt=prompt or "",
            negative_prompt=negative_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
                ],
            }

        prep_time = time.perf_counter() - task_start_time
        logging.info("Preparation phase completed in %.2f seconds", prep_time)

        progress(0.15, desc="Generating video")
        inference_start_time = time.perf_counter()

        _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,
        )

        inference_time = time.perf_counter() - inference_start_time
        logging.info(
            "Video generation (inference) completed in %.2f seconds", inference_time
        )

        progress(0.95, desc="Finalizing output")
        gc.collect()

        # TODO:
        if torch.cuda.is_available():
            # torch.cuda.empty_cache()
            pass

        if CONFORM_OUTPUT_TO_SOURCE_FPS:
            progress(0.96, desc="Matching source FPS")
            save_file = _conform_output_to_source_fps(save_file, pose_path)
        if COPY_SOURCE_AUDIO:
            progress(0.97, desc="Restoring source audio")
            save_file = _copy_source_audio_to_output(save_file, pose_path)
        progress(0.98, desc="Preparing video for display")
        display_file = _prepare_output_for_gradio(save_file)
        progress(1.0, desc="Done")

        total_time = time.perf_counter() - task_start_time
        logging.info("Total generation job finished in %.2f seconds", total_time)
        logging.info("=" * 60)

        return display_file, "Done."

    except Exception:
        total_time = time.perf_counter() - task_start_time
        logging.info("Generation job failed after %.2f seconds", total_time)
        logging.info("=" * 60)

        logging.exception("Generation failed")
        return None, traceback.format_exc()


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") as demo:
        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_negative_prompt = gr.Textbox(label="Negative prompt", lines=3, value="")
            (
                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)

            def handle_upload_generation(
                image,
                mask_image,
                pose_video,
                mask_video,
                prompt,
                negative_prompt,
                mode,
                sample_steps,
                guide_scale,
                sample_shift,
                seed,
                target_size,
                segment_len,
                segment_overlap,
            ):
                return run_generation(
                    {
                        "image": image,
                        "mask_image": mask_image,
                        "pose_video": pose_video,
                        "mask_video": mask_video,
                        "prompt": prompt,
                        "negative_prompt": negative_prompt,
                        "mode": mode,
                        "sample_steps": sample_steps,
                        "guide_scale": guide_scale,
                        "sample_shift": sample_shift,
                        "seed": seed,
                        "target_size": target_size,
                        "segment_len": segment_len,
                        "segment_overlap": segment_overlap,
                    }
                )

            run_upload.click(
                handle_upload_generation,
                inputs=[
                    up_image,
                    up_mask_image,
                    up_pose_video,
                    up_mask_video,
                    up_prompt,
                    up_negative_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("Prepared Examples"):
            with gr.Row():
                example_dropdown = gr.Dropdown(
                    choices=example_names, value=default_example, label="Example"
                )
                mode_view = gr.Textbox(
                    value=default_preview[6], label="Mode", interactive=False
                )
            reference_note = gr.Markdown(default_preview[8])
            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[7],
                    label="Reference set",
                    columns=4,
                    height=260,
                    selected_index=0,
                    preview=True,
                )

            prompt = gr.Textbox(value=default_preview[4], label="Prompt", lines=3)
            negative_prompt = gr.Textbox(
                value=default_preview[5], label="Negative 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)

            def handle_example_generation(
                example_name,
                prompt,
                negative_prompt,
                sample_steps,
                guide_scale,
                sample_shift,
                seed,
                target_size,
                segment_len,
                segment_overlap,
            ):
                return run_generation(
                    {
                        "example_name": example_name,
                        "prompt": prompt,
                        "negative_prompt": negative_prompt,
                        "sample_steps": sample_steps,
                        "guide_scale": guide_scale,
                        "sample_shift": sample_shift,
                        "seed": seed,
                        "target_size": target_size,
                        "segment_len": segment_len,
                        "segment_overlap": segment_overlap,
                    }
                )

            example_dropdown.change(
                load_example_preview,
                inputs=[example_dropdown],
                outputs=[
                    ref_preview,
                    driving_preview,
                    ref_mask_preview,
                    driving_mask_preview,
                    prompt,
                    negative_prompt,
                    mode_view,
                    reference_gallery,
                    reference_note,
                ],
            )
            run_example.click(
                handle_example_generation,
                inputs=[
                    example_dropdown,
                    prompt,
                    negative_prompt,
                    sample_steps,
                    guide_scale,
                    sample_shift,
                    seed,
                    target_size,
                    segment_len,
                    segment_overlap,
                ],
                outputs=[output_video, 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(
        theme=gr.themes.Soft(primary_hue="blue"),
        allowed_paths=[str(OUTPUT_DIR.resolve())],
        show_error=True,
    )