"""SCAIL-2 MultiRef Segmented v6 — HuggingFace ZeroGPU Space.

Runs `scail2MultiRefSegmented_v6-2.json` (a ComfyUI API-format export) with
ComfyUI driven in-process via `execution.PromptExecutor`.

The exported workflow contains `WanAniDirector` (node 80), whose SegmentQueueRunner
re-submits per-segment sub-jobs to a live ComfyUI HTTP server (`/prompt`, `/history`)
and depends on `extra_pnginfo.sqr_full_prompt` injected by its frontend JS — neither
exists here.  It also dynamically creates the `LoadImage` nodes feeding
`WanSQRMultiReference.image_1..6`, which is why those inputs are absent from the
export.  `_build_base_workflow()` performs the equivalent rewiring in Python; the
export has 分段数=1, so a single segment is functionally equivalent.
"""

import os
# Disable PyTorch's CUDA memory caching pool entirely.
# ZeroGPU's `large` size is half an RTX Pro 6000 (a partitioned GPU); PyTorch's
# caching allocator calls nvmlDeviceGetMemoryInfo (unsupported on partitions) in
# mallocRetry when cudaMalloc fails, causing an assertion.  cudaMallocAsync avoids
# NVML but retains freed memory in its pool, accumulating tens of GB across model
# loads/offloads within one GPU session.  PYTORCH_NO_CUDA_MEMORY_CACHING bypasses
# all pooling: every alloc is a direct cudaMalloc, every free a direct cudaFree —
# NVML is never queried and offloaded models release VRAM immediately.
os.environ.setdefault("PYTORCH_NO_CUDA_MEMORY_CACHING", "1")

import copy
import json
import math
import random
import shutil
import subprocess
import sys
import time
from pathlib import Path

# ─── Paths ───────────────────────────────────────────────────────────────────

APP_DIR = Path(__file__).parent
WORKFLOW_JSON = APP_DIR / "scail2MultiRefSegmented_v6-2.json"

COMFYUI_DIR = APP_DIR / "ComfyUI"
CUSTOM_NODES_DIR = COMFYUI_DIR / "custom_nodes"
INPUT_DIR = COMFYUI_DIR / "input"
OUTPUT_DIR = COMFYUI_DIR / "output"

# ─── Node IDs (from scail2MultiRefSegmented_v6-2.json) ───────────────────────

N_SAM3_REF        = "2"    # SAM3_VideoTrack  (reference stills)
N_SAM3_DRIVE      = "3"    # SAM3_VideoTrack  (driving video)
N_DIFFUSION       = "4"    # DiffusionModelLoaderKJ
N_CLIP            = "5"    # CLIPLoader       (umt5)
N_VAE             = "6"    # VAELoader
N_CLIP_VISION     = "9"    # CLIPVisionLoader
N_SAM3_CKPT       = "11"   # CheckpointLoaderSimple (sam3.1)
N_TRANSITION      = "12"   # SQRSCAIL2TransitionToVideo
N_SAMPLER         = "13"   # KSampler
N_COMBINE         = "15"   # VHS_VideoCombine  (the output node)
N_SAM3_TEXT_REF   = "16"   # CLIPTextEncode   (reference subject prompt)
N_SAM3_TEXT_DRIVE = "17"   # CLIPTextEncode   (driving subject prompt)
N_POSITIVE        = "21"   # CLIPTextEncode   (positive; was fed by node 80)
N_NEGATIVE        = "22"   # CLIPTextEncode   (negative)
N_SAMPLING_SD3    = "23"   # ModelSamplingSD3
N_POWER_LORA      = "25"   # Power Lora Loader (rgthree) — replaced
N_COLORED_MASK    = "29"   # SQRScail2ColoredMaskAdvanced
N_MULTI_REF       = "50"   # WanSQRMultiReference
N_REF_SPLIT       = "51"   # SQRScail2ReferenceBatchSplit
N_RESOLUTION      = "53"   # LHResolutionSetting
N_LOAD_VIDEO      = "67"   # VHS_LoadVideo
N_CONTEXT_WINDOWS = "75"   # WanContextWindowsManual
N_DIRECTOR        = "80"   # WanAniDirector — removed

# Nodes we synthesise to replace node 25 / the Director's LoadImage injection.
N_LORA_1 = "251"
N_LORA_2 = "252"
REF_IMAGE_NODES = ["101", "102", "103", "104", "105", "106"]
MAX_REFS = len(REF_IMAGE_NODES)

# PreviewImage nodes — pure UI, wasted compute headless.
DROP_NODES = ["19", "20", "34", N_DIRECTOR, N_POWER_LORA]

OUTPUT_PREFIX = "scail2_mrs_v6"

# LHResolutionSetting.RESOLUTIONS keys (ComfyUI-WanAni-SQR/scail_reference_nodes.py)
RESOLUTION_CHOICES = [
    "1:1 480p - 480 x 480",
    "1:1 720p - 720 x 720",
    "1:1 1024 - 1024 x 1024",
    "4:3 480p - 640 x 480",
    "4:3 768p - 1024 x 768",
    "16:9 480p - 854 x 480",
    "16:9 480p safe - 848 x 480",
    "16:9 720p - 1280 x 720",
    "21:9 480p - 1120 x 480",
    "21:9 720p - 1680 x 720",
]
ORIENTATION_CHOICES = ["竖屏 Portrait", "横屏 Landscape"]
IDENTITY_MODES = [
    "multi_person",
    "single_person_multi_reference",
    "multi_person_multi_reference",
]
SORT_BY_CHOICES = ["area", "left_to_right", "none"]
PRECISION_CHOICES = ["nvfp4 (RTX Pro 6000 / Blackwell)", "fp8_scaled"]

DEFAULT_NEGATIVE = (
    "色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,"
    "低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,"
    "毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走"
)

# ─── Helper ───────────────────────────────────────────────────────────────────

def _run(*cmd):
    print("$", " ".join(str(c) for c in cmd), flush=True)
    subprocess.run([str(c) for c in cmd], check=True)

# ─── Phase 1: repos + model download (no CUDA) ───────────────────────────────

CUSTOM_REPOS = {
    "ComfyUI-VideoHelperSuite":
        "https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite",
    "ComfyUI-KJNodes":
        "https://github.com/kijai/ComfyUI-KJNodes",
    # SQRSCAIL2TransitionToVideo, SQRScail2ColoredMaskAdvanced,
    # SQRScail2ReferenceBatchSplit, WanSQRMultiReference, LHResolutionSetting
    "ComfyUI-WanAni-SQR":
        "https://github.com/zere111ai/ComfyUI-WanAni-SQR",
}

# ComfyUI-WanAni-SQR pins `opencv-python>=4.8`; the GUI build is useless in a Space
# and conflicts with opencv-python-headless from requirements.txt.
SKIP_REQUIREMENTS = {"ComfyUI-WanAni-SQR"}


def _setup_repos():
    """Clone ComfyUI and the custom nodes the workflow needs."""
    if not COMFYUI_DIR.exists():
        print("Cloning ComfyUI (Comfy-Org master)…")
        _run("git", "clone", "--depth=1",
             "https://github.com/Comfy-Org/ComfyUI", COMFYUI_DIR)
        _run("pip", "install", "-r", COMFYUI_DIR / "requirements.txt", "-q")

    CUSTOM_NODES_DIR.mkdir(parents=True, exist_ok=True)
    INPUT_DIR.mkdir(parents=True, exist_ok=True)
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    for name, url in CUSTOM_REPOS.items():
        dest = CUSTOM_NODES_DIR / name
        if dest.exists():
            continue
        print(f"Cloning {name}…")
        _run("git", "clone", "--depth=1", url, dest)
        req = dest / "requirements.txt"
        if req.exists() and name not in SKIP_REQUIREMENTS:
            subprocess.run(["pip", "install", "-r", str(req), "-q"])


# repo_id, filename, folder_paths type
MODEL_SPECS = {
    "diffusion": (
        "LHQAQ-Li/wan2.1_14B_SCAIL_2_nvfp4_comfy_V2",
        "wan2.1_14B_SCAIL_2_nvfp4_comfy_V2.safetensors",
        "diffusion_models",
    ),
    # Downloaded on demand — see _model_path("diffusion_fp8").
    "diffusion_fp8": (
        "Comfy-Org/SCAIL-2",
        "diffusion_models/wan2.1_14B_SCAIL_2_fp8_scaled.safetensors",
        "diffusion_models",
    ),
    # The export names `nsfw_wan_umt5-xxl_fp8_scaled.safetensors`, which is not on
    # the Hub; the stock scaled fp8 umt5-xxl is the equivalent text encoder.
    "text_encoder": (
        "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
        "split_files/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors",
        "text_encoders",
    ),
    "vae": (
        "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
        "split_files/vae/wan_2.1_vae.safetensors",
        "vae",
    ),
    "clip_vision": (
        "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
        "split_files/clip_vision/clip_vision_h.safetensors",
        "clip_vision",
    ),
    "sam3": (
        "Comfy-Org/sam3.1",
        "checkpoints/sam3.1_multiplex_fp16.safetensors",
        "checkpoints",
    ),
    "lora_lightx2v": (
        "Kijai/WanVideo_comfy",
        "Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank64_bf16.safetensors",
        "loras",
    ),
    "lora_dpo": (
        "Comfy-Org/SCAIL-2",
        "loras/wan2.1_SCAIL_2_DPO_lora_bf16.safetensors",
        "loras",
    ),
}

# fp8_scaled is a 14 GB alternative to the nvfp4 default; only fetch it if asked for.
LAZY_MODELS = {"diffusion_fp8"}

_MODEL_PATHS: dict[str, Path] = {}


def _model_path(key: str) -> Path:
    """Resolve (downloading and registering on first use) one model file."""
    if key in _MODEL_PATHS:
        return _MODEL_PATHS[key]

    from huggingface_hub import hf_hub_download

    repo_id, filename, model_type = MODEL_SPECS[key]
    print(f"Ensuring {Path(filename).name}…", flush=True)
    path = Path(hf_hub_download(repo_id=repo_id, filename=filename))
    _MODEL_PATHS[key] = path

    # A lazily fetched model lands in a snapshot dir ComfyUI has not seen yet.
    if _comfyui_ready:
        import folder_paths
        folder_paths.add_model_folder_path(model_type, str(path.parent))
    return path


def _download_models():
    for key in MODEL_SPECS:
        if key not in LAZY_MODELS:
            _model_path(key)


def _diffusion_key(model_precision: str) -> str:
    return "diffusion_fp8" if model_precision == "fp8_scaled" else "diffusion"


def _prefetch_precision(model_precision: str) -> str:
    """Fetch a lazily-downloaded diffusion model outside the GPU-billed window.

    Wired to the precision dropdown's change event: a 14 GB download inside
    @spaces.GPU would burn ZeroGPU quota and likely blow the duration budget.
    """
    key = _diffusion_key(model_precision)
    if key in _MODEL_PATHS:
        return model_precision
    print(f"Prefetching {key} outside GPU context…", flush=True)
    _model_path(key)
    return model_precision


# ─── Phase 2: graph rewrite (no CUDA, no ComfyUI import) ─────────────────────

def _build_base_workflow(lora_1_on: bool, lora_1_strength: float,
                         lora_2_on: bool, lora_2_strength: float,
                         ref_count: int) -> dict:
    """Turn the Director-driven export into a self-contained API prompt.

    Drops the PreviewImage nodes, WanAniDirector and the rgthree Power Lora Loader;
    replaces the latter with core LoraLoaderModelOnly nodes and wires `ref_count`
    LoadImage nodes into WanSQRMultiReference.image_1..N.
    """
    with open(WORKFLOW_JSON, encoding="utf-8") as f:
        wf: dict = json.load(f)

    for node_id in DROP_NODES:
        wf.pop(node_id, None)

    # node 25 (rgthree, MODEL+CLIP) fed only node 23's MODEL input; its CLIP output
    # was unused (nodes 21/22 take CLIP straight from node 5).  Two core
    # LoraLoaderModelOnly nodes reproduce the two enabled LoRAs from lora_1/lora_4.
    model_ref = [N_DIFFUSION, 0]
    for node_id, enabled, key, strength in (
        (N_LORA_1, lora_1_on, "lora_lightx2v", lora_1_strength),
        (N_LORA_2, lora_2_on, "lora_dpo", lora_2_strength),
    ):
        if not enabled:
            continue
        wf[node_id] = {
            "inputs": {
                "lora_name": _model_path(key).name,
                "strength_model": float(strength),
                "model": model_ref,
            },
            "class_type": "LoraLoaderModelOnly",
            "_meta": {"title": f"LoRA {key}"},
        }
        model_ref = [node_id, 0]
    wf[N_SAMPLING_SD3]["inputs"]["model"] = model_ref

    # The Director produced node 21's text; it is now a plain widget value.
    wf[N_POSITIVE]["inputs"]["text"] = ""

    # The Director also created these LoadImage nodes and linked them.
    for slot in range(1, MAX_REFS + 1):
        wf[N_MULTI_REF]["inputs"].pop(f"image_{slot}", None)
    for slot in range(1, ref_count + 1):
        node_id = REF_IMAGE_NODES[slot - 1]
        wf[node_id] = {
            "inputs": {"image": ""},
            "class_type": "LoadImage",
            "_meta": {"title": f"Reference Image {slot}"},
        }
        wf[N_MULTI_REF]["inputs"][f"image_{slot}"] = [node_id, 0]

    return wf


# ─── Phase 3: ComfyUI init (deferred until first @spaces.GPU call) ───────────
# model_management.py calls torch.cuda.current_device() at import time → crashes
# without a real GPU.  Defer all ComfyUI imports to the GPU context.

_comfyui_ready = False


def _init_comfyui():
    """Import ComfyUI, register model paths, load nodes. Runs once per worker."""
    global _comfyui_ready

    if _comfyui_ready:
        return

    if str(COMFYUI_DIR) not in sys.path:
        sys.path.insert(0, str(COMFYUI_DIR))

    import folder_paths
    folder_paths.base_path = str(COMFYUI_DIR)

    # Set PromptServer.instance BEFORE loading nodes: VHS, KJNodes and
    # ComfyUI-WanAni-SQR all register aiohttp routes on this singleton at import
    # time (segment_queue_node.py does `@server.PromptServer.instance.routes.get`).
    import server as comfy_server

    class _RouteTableDef:
        """Mimics aiohttp.web.RouteTableDef — supports @routes.get('/path')."""
        def __getattr__(self, method):
            def route(path, **kwargs):
                def decorator(handler):
                    return handler
                return decorator
            return route

    class _MockRouter:
        frozen = True   # stops KJNodes freezing an already-frozen router
        def __getattr__(self, name):
            return lambda *args, **kwargs: None

    class _MockApp:
        router = _MockRouter()
        def __getattr__(self, name):
            return lambda *args, **kwargs: None

    class _MockQueue:
        def __getattr__(self, name):
            return lambda *args, **kwargs: None

    class _MockPromptServer:
        client_id = None
        routes = _RouteTableDef()
        app = _MockApp()
        prompt_queue = _MockQueue()

        def send_sync(self, *args, **kwargs): pass
        def queue_updated(self): pass
        def __getattr__(self, name):
            return lambda *args, **kwargs: None

    comfy_server.PromptServer.instance = _MockPromptServer()

    for key, path in _MODEL_PATHS.items():
        folder_paths.add_model_folder_path(MODEL_SPECS[key][2], str(path.parent))

    import asyncio
    import nodes as comfy_nodes
    asyncio.run(comfy_nodes.init_extra_nodes(init_custom_nodes=True))

    missing = _missing_node_classes(comfy_nodes.NODE_CLASS_MAPPINGS)
    if missing:
        raise RuntimeError(f"Node classes failed to load: {sorted(missing)}")

    # Force LOW_VRAM mode: offload models aggressively between nodes.  ZeroGPU
    # reports the full RTX Pro 6000 to ComfyUI, so it picks HIGH_VRAM and tries to
    # keep the diffusion model, umt5, CLIP Vision, SAM3 and the VAE resident
    # simultaneously — over the per-request limit for the `large` partition.
    import comfy.model_management as _mm
    _mm.vram_state = _mm.VRAMState.LOW_VRAM
    print(f"=== SCAIL2-MRS: VRAM mode set to {_mm.vram_state} ===", flush=True)

    _comfyui_ready = True
    print("=== SCAIL2-MRS: ComfyUI ready ===", flush=True)


def _missing_node_classes(node_class_mappings: dict) -> set[str]:
    """class_types the rewritten workflow needs but ComfyUI did not register."""
    wf = _build_base_workflow(True, 1.0, True, 1.0, MAX_REFS)
    return {
        node["class_type"] for node in wf.values()
        if node["class_type"] not in node_class_mappings
    }


# ─── Inference ────────────────────────────────────────────────────────────────

def _inject(workflow: dict, node_id: str, key: str, value):
    if node_id in workflow:
        workflow[node_id]["inputs"][key] = value
    else:
        print(f"⚠ Node {node_id} not found, skipping {key!r}", flush=True)


def _stage_input(path: str | os.PathLike) -> str:
    """Copy a Gradio upload into ComfyUI/input/ and return its bare filename."""
    INPUT_DIR.mkdir(parents=True, exist_ok=True)
    name = Path(path).name
    dest = INPUT_DIR / name
    if Path(path).resolve() != dest.resolve():
        shutil.copy(path, dest)
    return name


def _ref_paths(ref_files) -> list[str]:
    """Normalise the gr.Files / gr.Gallery value into at most MAX_REFS paths."""
    if not ref_files:
        return []
    paths = []
    for item in ref_files:
        # gr.Gallery yields (path, caption) tuples; gr.Files yields str or FileData.
        if isinstance(item, (tuple, list)):
            item = item[0]
        paths.append(str(getattr(item, "name", item)))
    return paths[:MAX_REFS]


class _NullServer:
    """Minimal PromptServer mock for library-mode execution."""
    client_id = None

    def __getattr__(self, name):
        return lambda *args, **kwargs: None


def _estimate_duration(frame_load_cap, steps, context_length, context_overlap):
    """Seconds of GPU time to request, clamped to ZeroGPU's practical ceiling."""
    frames = max(1, int(frame_load_cap))
    stride = max(1, int(context_length) - int(context_overlap))
    windows = max(1, math.ceil(max(0, frames - int(context_overlap)) / stride))
    est = (
        150                      # model loads / offloads
        + frames * 0.6           # SAM3 tracking, reference + driving passes
        + windows * int(steps) * 28   # sampling
        + frames * 0.4           # VAE decode + mp4 mux
    )
    return int(min(600, max(120, est)))


def _generate_inner(
    video_path, ref_files, positive_prompt, negative_prompt,
    sam3_ref_prompt, sam3_drive_prompt,
    orientation, resolution, force_rate, frame_load_cap,
    seed, steps, cfg,
    identity_mode, sort_by, main_index, background_indices,
    context_length, context_overlap,
    model_precision, lora_1_on, lora_1_strength, lora_2_on, lora_2_strength,
):
    _t0 = time.time()

    def _log(msg):
        print(f"=== SCAIL2-MRS [{time.time() - _t0:6.1f}s]: {msg} ===", flush=True)

    if not video_path:
        raise ValueError("入力動画をアップロードしてください。")

    ref_paths = _ref_paths(ref_files)
    if not ref_paths:
        raise ValueError("参照画像を1枚以上アップロードしてください。")

    # Already resolved outside the GPU window by _prefetch_precision().
    diffusion_path = _model_path(_diffusion_key(model_precision))

    _init_comfyui()
    _log("ComfyUI init done")

    import execution as comfy_execution

    wf = _build_base_workflow(
        bool(lora_1_on), float(lora_1_strength),
        bool(lora_2_on), float(lora_2_strength),
        len(ref_paths),
    )

    # Model files
    _inject(wf, N_DIFFUSION,   "model_name", diffusion_path.name)
    # sage_attention defaults to "auto" in the export; sageattention is not installed.
    _inject(wf, N_DIFFUSION,   "sage_attention", "disabled")
    _inject(wf, N_CLIP,        "clip_name", _model_path("text_encoder").name)
    _inject(wf, N_VAE,         "vae_name",  _model_path("vae").name)
    _inject(wf, N_CLIP_VISION, "clip_name", _model_path("clip_vision").name)
    _inject(wf, N_SAM3_CKPT,   "ckpt_name", _model_path("sam3").name)

    # Inputs
    _inject(wf, N_LOAD_VIDEO, "video",          _stage_input(video_path))
    _inject(wf, N_LOAD_VIDEO, "force_rate",     int(force_rate))
    _inject(wf, N_LOAD_VIDEO, "frame_load_cap", int(frame_load_cap))
    for slot, ref_path in enumerate(ref_paths, start=1):
        _inject(wf, REF_IMAGE_NODES[slot - 1], "image", _stage_input(ref_path))

    # Prompts
    _inject(wf, N_POSITIVE,        "text", positive_prompt or "")
    _inject(wf, N_NEGATIVE,        "text", negative_prompt or "")
    _inject(wf, N_SAM3_TEXT_REF,   "text", sam3_ref_prompt or "human")
    _inject(wf, N_SAM3_TEXT_DRIVE, "text", sam3_drive_prompt or "human")

    # Resolution / sampling
    _inject(wf, N_RESOLUTION, "orientation", orientation)
    _inject(wf, N_RESOLUTION, "resolution",  resolution)
    if int(seed) < 0:
        seed = random.randint(0, 2**53 - 1)
    _inject(wf, N_SAMPLER, "seed",  int(seed))
    _inject(wf, N_SAMPLER, "steps", int(steps))
    _inject(wf, N_SAMPLER, "cfg",   float(cfg))

    # Multi-reference identity handling
    _inject(wf, N_COLORED_MASK, "identity_mode",      identity_mode)
    _inject(wf, N_COLORED_MASK, "sort_by",            sort_by)
    _inject(wf, N_COLORED_MASK, "background_indices", background_indices or "")
    _inject(wf, N_REF_SPLIT,    "main_index",
            max(0, min(int(main_index), len(ref_paths) - 1)))

    _inject(wf, N_CONTEXT_WINDOWS, "context_length",  int(context_length))
    _inject(wf, N_CONTEXT_WINDOWS, "context_overlap", int(context_overlap))

    _inject(wf, N_COMBINE, "filename_prefix", OUTPUT_PREFIX)

    import nest_asyncio
    nest_asyncio.apply()

    import torch
    if torch.cuda.is_available():
        props = torch.cuda.get_device_properties(0)
        _log(f"GPU: {props.name}, total={props.total_memory / 1024**3:.1f}GB, "
             f"capability=sm_{props.major}{props.minor}")

    executor = comfy_execution.PromptExecutor(
        server=_NullServer(),
        cache_args={"ram": 0, "ram_inactive": 0},
    )

    _log(f"executing: {len(ref_paths)} refs, {frame_load_cap} frames, seed={seed}")
    before = set(OUTPUT_DIR.glob(f"{OUTPUT_PREFIX}*"))
    result = executor.execute(
        wf, prompt_id="gradio_run", extra_data={}, execute_outputs=[N_COMBINE],
    )
    _log(f"execute() done, result={result!r}")

    # Newer ComfyUI returns (success, error, node_errors).
    if isinstance(result, tuple) and len(result) >= 2 and result[1]:
        raise RuntimeError(f"ComfyUI execution error: {result[1]}")

    new_videos = sorted(
        (p for p in OUTPUT_DIR.glob(f"{OUTPUT_PREFIX}*")
         if p not in before and p.suffix in (".mp4", ".webm")),
        key=lambda p: p.stat().st_mtime,
    )
    if not new_videos:
        raise RuntimeError("ComfyUI produced no video file.")
    _log(f"output: {new_videos[-1]}")
    return str(new_videos[-1])


print("=== SCAIL2-MRS: cloning repos… ===", flush=True)
_setup_repos()
print("=== SCAIL2-MRS: downloading models… ===", flush=True)
_download_models()
print("=== SCAIL2-MRS: models ready ===", flush=True)

GENERATE_PARAMS = (
    "video_path", "ref_files", "positive_prompt", "negative_prompt",
    "sam3_ref_prompt", "sam3_drive_prompt",
    "orientation", "resolution", "force_rate", "frame_load_cap",
    "seed", "steps", "cfg",
    "identity_mode", "sort_by", "main_index", "background_indices",
    "context_length", "context_overlap",
    "model_precision", "lora_1_on", "lora_1_strength", "lora_2_on", "lora_2_strength",
)


def _get_duration(*args):
    kw = dict(zip(GENERATE_PARAMS, args))
    return _estimate_duration(kw["frame_load_cap"], kw["steps"],
                              kw["context_length"], kw["context_overlap"])


try:
    import spaces as _spaces

    @_spaces.GPU(duration=_get_duration)
    def _generate_gpu(*args):
        import traceback
        try:
            return _generate_inner(*args)
        except Exception:
            traceback.print_exc()
            raise

except ImportError:
    _generate_gpu = _generate_inner


def generate(*args):
    # ZeroGPU pre-warms with empty inputs; don't treat that as a failure.
    if not args or not args[0]:
        print("=== SCAIL2-MRS: generate() with no video (pre-warm) ===", flush=True)
        return None
    return _generate_gpu(*args)

# ─── Gradio UI ────────────────────────────────────────────────────────────────

import gradio as gr

with gr.Blocks(title="SCAIL-2 MultiRef Segmented v6") as demo:
    gr.Markdown(
        "# SCAIL-2 MultiRef Segmented v6\n"
        "WAN2.1 14B SCAIL-2 (NVFP4) — 駆動動画のモーションを最大6枚の参照画像に転写します。\n"
        "SAM3 が参照画像側と駆動動画側の被写体をトラッキングし、色分けマスクとして "
        "SCAIL-2 に渡します。"
    )

    with gr.Row():
        with gr.Column(scale=1):
            inp_video = gr.Video(label="駆動動画 (Driving Video)")
            inp_refs = gr.Files(
                label=f"参照画像 (Reference Images, 最大 {MAX_REFS} 枚)",
                file_types=["image"],
            )
            gal_refs = gr.Gallery(label="参照画像プレビュー", columns=3,
                                  height=180, show_label=False)
            inp_positive = gr.Textbox(
                label="Positive Prompt",
                placeholder="女人在跳舞 / a woman is dancing",
                lines=3,
            )
            inp_negative = gr.Textbox(
                label="Negative Prompt", value=DEFAULT_NEGATIVE, lines=3,
            )
            btn = gr.Button("Generate", variant="primary")

        with gr.Column(scale=1):
            out_video = gr.Video(label="生成結果")

            with gr.Accordion("解像度・フレーム", open=True):
                inp_orientation = gr.Radio(
                    ORIENTATION_CHOICES, value="竖屏 Portrait", label="向き",
                )
                inp_resolution = gr.Dropdown(
                    RESOLUTION_CHOICES, value="16:9 480p safe - 848 x 480",
                    label="解像度 (縦向きでは長辺・短辺が入れ替わります)",
                )
                inp_force_rate = gr.Slider(
                    8, 30, value=24, step=1, label="force_rate (入力動画の再サンプル fps)",
                )
                inp_frame_load_cap = gr.Slider(
                    17, 240, value=144, step=4,
                    label="frame_load_cap (生成フレーム数 — 初回は 144 推奨)",
                )

            with gr.Accordion("多参照・SAM3", open=False):
                inp_identity_mode = gr.Dropdown(
                    IDENTITY_MODES, value="multi_person", label="identity_mode",
                    info="single_person_multi_reference: 1人を複数参照で表現",
                )
                inp_sort_by = gr.Dropdown(
                    SORT_BY_CHOICES, value="area", label="sort_by (被写体の並び順)",
                )
                inp_main_index = gr.Slider(
                    0, MAX_REFS - 1, value=0, step=1,
                    label="main_index (主参照にする画像の 0 始まりの番号)",
                )
                inp_background_indices = gr.Textbox(
                    label="background_indices",
                    placeholder="背景として扱う参照画像の 1 始まり番号 (例: 2 または 1,4)",
                )
                inp_sam3_ref = gr.Textbox(
                    label="SAM3 参照側プロンプト", value="一个女人",
                )
                inp_sam3_drive = gr.Textbox(
                    label="SAM3 駆動側プロンプト", value="human",
                )

            with gr.Accordion("サンプラー・LoRA・モデル", open=False):
                inp_seed = gr.Number(value=-1, precision=0,
                                     label="seed (-1 でランダム)")
                inp_steps = gr.Slider(1, 12, value=4, step=1, label="steps")
                inp_cfg = gr.Slider(1.0, 8.0, value=1.0, step=0.1, label="cfg")
                inp_context_length = gr.Slider(
                    33, 81, value=81, step=4, label="context_length",
                )
                inp_context_overlap = gr.Slider(
                    0, 32, value=16, step=4, label="context_overlap",
                )
                inp_precision = gr.Dropdown(
                    PRECISION_CHOICES, value=PRECISION_CHOICES[0],
                    label="拡散モデル精度",
                    info="fp8_scaled は初回選択時に約14GBを追加ダウンロードします",
                )
                inp_lora_1_on = gr.Checkbox(
                    value=True, label="LoRA: lightx2v I2V 480p cfg-step-distill",
                )
                inp_lora_1_strength = gr.Slider(
                    0.0, 1.5, value=1.0, step=0.05, label="lightx2v strength",
                )
                inp_lora_2_on = gr.Checkbox(
                    value=True, label="LoRA: SCAIL-2 DPO",
                )
                inp_lora_2_strength = gr.Slider(
                    0.0, 1.5, value=1.0, step=0.05, label="SCAIL-2 DPO strength",
                )

    inp_refs.change(
        fn=lambda files: _ref_paths(files),
        inputs=inp_refs,
        outputs=gal_refs,
    )

    # Download a non-default diffusion model as soon as it is picked, so it never
    # happens inside the GPU-billed window.
    inp_precision.change(
        fn=_prefetch_precision, inputs=inp_precision, outputs=inp_precision,
    )

    btn.click(
        fn=generate,
        inputs=[
            inp_video, inp_refs, inp_positive, inp_negative,
            inp_sam3_ref, inp_sam3_drive,
            inp_orientation, inp_resolution, inp_force_rate, inp_frame_load_cap,
            inp_seed, inp_steps, inp_cfg,
            inp_identity_mode, inp_sort_by, inp_main_index, inp_background_indices,
            inp_context_length, inp_context_overlap,
            inp_precision, inp_lora_1_on, inp_lora_1_strength,
            inp_lora_2_on, inp_lora_2_strength,
        ],
        outputs=out_video,
    )

if __name__ == "__main__":
    demo.queue().launch()