app.py

import argparse
import hashlib
import os
import queue
import re
import shutil
import subprocess
import sys
import tempfile
import threading
import time
import zipfile
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

import gradio as gr
import spaces
from huggingface_hub import snapshot_download


BUILD_ID = "relit-live-readme-gradio-v4"

ROOT = Path(__file__).resolve().parent
DEMO_ROOT = ROOT / "datasets" / "demos"
ENV_ROOT = ROOT / "datasets" / "envs"

WAN_DIR = ROOT / "models" / "Wan-AI" / "Wan2.1-T2V-1.3B"
CHECKPOINTS = {
    "model_frame25_480_832.ckpt": ROOT / "checkpoints" / "model_frame25_480_832.ckpt",
    "model_frame57_480_832.ckpt": ROOT / "checkpoints" / "model_frame57_480_832.ckpt",
    "model_frame1_1024_1472.ckpt": ROOT / "checkpoints" / "model_frame1_1024_1472.ckpt",
}

RUNTIME_ROOT = ROOT / "runtime_readme_gradio"
JOBS_ROOT = RUNTIME_ROOT / "jobs"
OUTPUTS_ROOT = RUNTIME_ROOT / "outputs"
UPLOADS_ROOT = RUNTIME_ROOT / "uploads"

IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".webp", ".exr"}
VIDEO_EXTS = {".mp4", ".mov", ".avi", ".mkv"}

REQUIRED_SAMPLE_DIRS = ["images_4", "Base Color", "depth", "normal"]
OPTIONAL_SAMPLE_DIRS = ["Metallic", "Roughness"]
REQUIRED_ENV_FILES = [
    "ldr_video_fix_first_frame.mp4",
    "hdr_log_video_fix_first_frame.mp4",
    "env_dir_video_fix_first_frame.mp4",
]

PRESETS: Dict[str, Dict[str, Any]] = {
    "Basic 25-frame relighting": {
        "checkpoint": "model_frame25_480_832.ckpt",
        "height": 480,
        "width": 832,
        "frames": 25,
        "flags": [],
        "output_kind": "video",
    },
    "25-frame rotating-light relighting": {
        "checkpoint": "model_frame25_480_832.ckpt",
        "height": 480,
        "width": 832,
        "frames": 25,
        "flags": ["--use_rotate_light"],
        "output_kind": "video",
    },
    "Fixed-frame relighting, width-axis light rotation": {
        "checkpoint": "model_frame25_480_832.ckpt",
        "height": 480,
        "width": 832,
        "frames": 25,
        "flags": ["--use_fixed_frame_and_w_rotate_light"],
        "output_kind": "video",
    },
    "Fixed-frame relighting, height-axis light rotation": {
        "checkpoint": "model_frame25_480_832.ckpt",
        "height": 480,
        "width": 832,
        "frames": 25,
        "flags": ["--use_fixed_frame_and_h_rotate_light"],
        "output_kind": "video",
    },
    "57-frame video relighting": {
        "checkpoint": "model_frame57_480_832.ckpt",
        "height": 480,
        "width": 832,
        "frames": 57,
        "flags": [],
        "output_kind": "video",
    },
    "Single-frame high-resolution relighting": {
        "checkpoint": "model_frame1_1024_1472.ckpt",
        "height": 1024,
        "width": 1472,
        "frames": 1,
        "flags": [],
        "output_kind": "image",
    },
}


for folder in [RUNTIME_ROOT, JOBS_ROOT, OUTPUTS_ROOT, UPLOADS_ROOT]:
    folder.mkdir(parents=True, exist_ok=True)


def tail_text(text: str, max_chars: int = 20000) -> str:
    text = text or ""
    if len(text) <= max_chars:
        return text
    return text[-max_chars:]


def hash_key(*items: Any) -> str:
    raw = "||".join(map(str, items))
    return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16]


def safe_rmtree(path: Path) -> None:
    if path.exists():
        shutil.rmtree(path, ignore_errors=True)


def run_cmd(cmd: list[str], cwd: Path = ROOT) -> Tuple[int, str, str]:
    proc = subprocess.run(cmd, cwd=str(cwd), capture_output=True, text=True)
    return proc.returncode, proc.stdout, proc.stderr


class MonotonicProgress:
    def __init__(self, progress: Optional[Any]):
        self.progress = progress
        self.last_value = 0.0

    def __call__(self, value: float, desc: str = "") -> None:
        if self.progress is None:
            return
        value = max(self.last_value, min(1.0, float(value)))
        self.last_value = value
        try:
            self.progress(value, desc=desc)
        except Exception:
            pass


def update_progress(progress: Optional[Any], value: float, desc: str) -> None:
    if progress is None:
        return
    try:
        progress(value, desc=desc)
    except Exception:
        pass


def strip_ansi(text: str) -> str:
    text = re.sub(r"\x1b\[[0-9;]*[A-Za-z]", "", text)
    return text.replace("\r", " ").strip()


def progress_from_tqdm_line(progress: Optional[Any], line: str) -> None:
    clean = strip_ansi(line)
    if not clean:
        return

    if "Loading models from:" in clean:
        update_progress(progress, 0.12, "Loading models")
        return
    if "processing " in clean:
        update_progress(progress, 0.18, "Preparing sample")
        return
    if "VAE encoding" in clean:
        update_progress(progress, 0.28, "Encoding inputs")
        return
    if "VAE decoding" in clean:
        update_progress(progress, 0.86, "Decoding result")
        return

    match = re.search(r"(\d{1,3})%\|", clean)
    if match:
        pct = min(100, max(0, int(match.group(1))))
        update_progress(progress, 0.30 + 0.50 * (pct / 100.0), clean[:120])


def run_inference_subprocess(
    cmd: list[str],
    env: dict[str, str],
    progress: Optional[Any],
) -> Tuple[int, str, str]:
    proc = subprocess.Popen(
        cmd,
        cwd=str(ROOT),
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        bufsize=1,
    )

    output_queue: queue.Queue[Tuple[str, str]] = queue.Queue()

    def reader(name: str, pipe: Any) -> None:
        try:
            for line in iter(pipe.readline, ""):
                output_queue.put((name, line))
        finally:
            pipe.close()

    threads = [
        threading.Thread(target=reader, args=("stdout", proc.stdout), daemon=True),
        threading.Thread(target=reader, args=("stderr", proc.stderr), daemon=True),
    ]
    for thread in threads:
        thread.start()

    stdout_parts: list[str] = []
    stderr_parts: list[str] = []
    started_at = time.time()
    last_tick = started_at
    update_progress(progress, 0.20, "Starting inference")

    while proc.poll() is None or not output_queue.empty():
        try:
            stream, line = output_queue.get(timeout=0.25)
        except queue.Empty:
            now = time.time()
            if now - last_tick > 5:
                elapsed = int(now - started_at)
                update_progress(progress, 0.35, f"Inference running ({elapsed}s)")
                last_tick = now
            continue

        if stream == "stdout":
            stdout_parts.append(line)
            if "Finish the inference" in line:
                update_progress(progress, 0.92, "Saving output")
        else:
            stderr_parts.append(line)
            progress_from_tqdm_line(progress, line)

    for thread in threads:
        thread.join(timeout=1)

    return proc.returncode or 0, "".join(stdout_parts), "".join(stderr_parts)


def patch_transformers_imports() -> list[str]:
    patched = []
    target = ROOT / "diffsynth" / "models" / "stepvideo_text_encoder.py"
    bad = "from transformers.modeling_utils import PretrainedConfig, PreTrainedModel"
    good = (
        "from transformers.configuration_utils import PretrainedConfig\n"
        "from transformers.modeling_utils import PreTrainedModel"
    )

    if target.exists():
        text = target.read_text(encoding="utf-8")
        if bad in text:
            target.write_text(text.replace(bad, good), encoding="utf-8")
            patched.append(str(target.relative_to(ROOT)))

    sitecustomize = ROOT / "sitecustomize.py"
    sitecustomize.write_text(
        """
try:
    import transformers.modeling_utils as _modeling_utils
    from transformers.configuration_utils import PretrainedConfig as _PretrainedConfig
    if not hasattr(_modeling_utils, "PretrainedConfig"):
        _modeling_utils.PretrainedConfig = _PretrainedConfig
except Exception:
    pass
""".lstrip(),
        encoding="utf-8",
    )
    patched.append(str(sitecustomize.relative_to(ROOT)))
    return patched


def ensure_wan_model(logs: list[str]) -> None:
    if WAN_DIR.exists() and any(WAN_DIR.rglob("*")):
        logs.append("Wan2.1 base model already present.")
        return

    logs.append("Downloading Wan2.1 base model...")
    WAN_DIR.mkdir(parents=True, exist_ok=True)
    snapshot_download(
        repo_id="Wan-AI/Wan2.1-T2V-1.3B",
        local_dir=str(WAN_DIR),
        token=os.getenv("HF_TOKEN"),
    )
    logs.append("Wan2.1 base model ready.")


def ensure_checkpoint(checkpoint_name: str, logs: list[str]) -> Path:
    checkpoint_path = CHECKPOINTS[checkpoint_name]
    if checkpoint_path.exists():
        logs.append(f"Checkpoint already present: {checkpoint_name}")
        return checkpoint_path

    logs.append(f"Downloading checkpoint: {checkpoint_name}")
    snapshot_download(
        repo_id="weiqingXiao/Relit-LiVE",
        local_dir=str(ROOT),
        allow_patterns=[f"checkpoints/{checkpoint_name}"],
        token=os.getenv("HF_TOKEN"),
    )
    if not checkpoint_path.exists():
        raise FileNotFoundError(f"Checkpoint download failed: {checkpoint_path}")

    logs.append(f"Checkpoint ready: {checkpoint_name}")
    return checkpoint_path


def startup_preflight() -> Tuple[bool, str]:
    logs = [
        f"Build: {BUILD_ID}",
        "Mode: README relit_inference.py presets.",
        "Full Cosmos inverse pipeline is not used in this app.",
    ]

    try:
        patched = patch_transformers_imports()
        logs.append("Compatibility patch files:")
        for path in patched:
            logs.append(f"- {path}")

        if not (ROOT / "relit_inference.py").exists():
            return False, "\n".join(logs + ["Missing relit_inference.py at repo root."])

        ensure_wan_model(logs)

        samples = list_demo_samples()
        envs = list_envs()
        logs.append(f"Found {len(samples)} repo demo sample(s).")
        logs.append(f"Found {len(envs)} repo environment(s).")

        code, out, err = run_python_preflight()
        logs.append(f"Python/DiffSynth preflight exit code: {code}")
        logs.append("STDOUT:\n" + tail_text(out, 6000))
        logs.append("STDERR:\n" + tail_text(err, 6000))

        return code == 0, "\n".join(logs)
    except Exception as exc:
        logs.append(f"Startup preflight failed: {repr(exc)}")
        return False, "\n".join(logs)


def run_python_preflight() -> Tuple[int, str, str]:
    code = """
import sys, torch
print("python", sys.version.replace("\\n", " "))
print("torch", torch.__version__)
print("torch_cuda", torch.version.cuda)
print("cuda_available", torch.cuda.is_available())
print("device", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU")
from transformers.modeling_utils import PretrainedConfig, PreTrainedModel
print("legacy_transformers_import_ok", PretrainedConfig.__name__, PreTrainedModel.__name__)
from diffsynth import ModelManager, WanVideoRelitlivePipeline
print("diffsynth_import_ok", ModelManager.__name__, WanVideoRelitlivePipeline.__name__)
"""
    env = os.environ.copy()
    env["PYTHONPATH"] = str(ROOT) + os.pathsep + env.get("PYTHONPATH", "")
    proc = subprocess.run(
        [sys.executable, "-c", code],
        cwd=str(ROOT),
        env=env,
        capture_output=True,
        text=True,
    )
    return proc.returncode, proc.stdout, proc.stderr


def has_images(path: Path) -> bool:
    if not path.exists() or not path.is_dir():
        return False
    for ext in IMAGE_EXTS:
        if any(path.glob(f"*{ext}")):
            return True
    return False


def is_valid_sample_dir(path: Path) -> Tuple[bool, list[str]]:
    missing = []
    for name in REQUIRED_SAMPLE_DIRS:
        if not has_images(path / name):
            missing.append(name)
    return not missing, missing


def is_valid_env_dir(path: Path) -> Tuple[bool, list[str]]:
    missing = [name for name in REQUIRED_ENV_FILES if not (path / name).exists()]
    return not missing, missing


def list_demo_samples() -> list[str]:
    if not DEMO_ROOT.exists():
        return []
    samples = []
    for path in sorted(DEMO_ROOT.iterdir()):
        if not path.is_dir():
            continue
        valid, _ = is_valid_sample_dir(path)
        if valid:
            samples.append(path.name)
    return samples


def list_envs() -> list[str]:
    if not ENV_ROOT.exists():
        return []
    envs = []
    for path in sorted(ENV_ROOT.iterdir()):
        if not path.is_dir():
            continue
        valid, _ = is_valid_env_dir(path)
        if valid:
            envs.append(path.name)
    return envs


def first_image(path: Path) -> Optional[str]:
    if not path.exists():
        return None
    files = []
    for ext in IMAGE_EXTS:
        files.extend(path.glob(f"*{ext}"))
    files = sorted(files)
    return str(files[0]) if files else None


def sample_preview(sample_name: Optional[str]) -> Optional[str]:
    if not sample_name:
        return None
    return first_image(DEMO_ROOT / sample_name / "images_4")


def env_preview(env_name: Optional[str]) -> Optional[str]:
    if not env_name:
        return None
    path = ENV_ROOT / env_name / "ldr_video_fix_first_frame.mp4"
    return str(path) if path.exists() else None


def normalize_zip_dataset(extracted_root: Path, dataset_root: Path) -> Tuple[Path, str]:
    candidates = [p for p in extracted_root.iterdir() if p.is_dir()]

    valid_root, _ = is_valid_sample_dir(extracted_root)
    if valid_root:
        sample_dir = dataset_root / "uploaded_sample"
        shutil.copytree(extracted_root, sample_dir)
        return dataset_root, "ZIP root is one sample. Wrapped as uploaded_sample."

    valid_samples = []
    for candidate in candidates:
        valid, _ = is_valid_sample_dir(candidate)
        if valid:
            valid_samples.append(candidate)

    if not valid_samples:
        required = ", ".join(REQUIRED_SAMPLE_DIRS)
        raise ValueError(
            "No valid Relit-LiVE sample found in ZIP. "
            f"Expected a sample directory containing: {required}."
        )

    for sample in valid_samples:
        shutil.copytree(sample, dataset_root / sample.name)

    return dataset_root, f"Loaded {len(valid_samples)} sample(s) from ZIP."


def normalize_zip_env(extracted_root: Path, env_root: Path) -> Tuple[Path, str]:
    valid_root, _ = is_valid_env_dir(extracted_root)
    if valid_root:
        shutil.copytree(extracted_root, env_root)
        return env_root, "ZIP root is a valid environment directory."

    for candidate in extracted_root.iterdir():
        if not candidate.is_dir():
            continue
        valid, _ = is_valid_env_dir(candidate)
        if valid:
            shutil.copytree(candidate, env_root)
            return env_root, f"Using environment directory from ZIP: {candidate.name}"

    raise ValueError(
        "No valid environment directory found in ZIP. "
        f"Expected: {', '.join(REQUIRED_ENV_FILES)}."
    )


def extract_zip_to(zip_file: Any, target_root: Path) -> Path:
    safe_rmtree(target_root)
    target_root.mkdir(parents=True, exist_ok=True)

    zip_path = Path(zip_file.name if hasattr(zip_file, "name") else zip_file)
    with zipfile.ZipFile(zip_path, "r") as archive:
        for member in archive.infolist():
            member_path = target_root / member.filename
            if not str(member_path.resolve()).startswith(str(target_root.resolve())):
                raise ValueError("Unsafe path found in ZIP archive.")
        archive.extractall(target_root)

    return target_root


def prepare_dataset(
    source_mode: str,
    sample_name: Optional[str],
    dataset_zip: Optional[Any],
    dataset_root: Path,
    extract_root: Path,
) -> Tuple[Path, str]:
    safe_rmtree(dataset_root)
    dataset_root.mkdir(parents=True, exist_ok=True)

    if source_mode == "Repo demo sample":
        if not sample_name:
            raise ValueError("Select a repo demo sample.")
        src = DEMO_ROOT / sample_name
        valid, missing = is_valid_sample_dir(src)
        if not valid:
            raise ValueError(f"Invalid repo sample. Missing: {', '.join(missing)}")
        shutil.copytree(src, dataset_root / sample_name)
        return dataset_root, f"Using repo sample: {sample_name}"

    if source_mode == "Prepared dataset ZIP":
        if dataset_zip is None:
            raise ValueError("Upload a prepared Relit-LiVE dataset ZIP.")
        extracted = extract_zip_to(dataset_zip, extract_root)
        dataset_path, msg = normalize_zip_dataset(extracted, dataset_root)
        return dataset_path, msg

    raise ValueError("Raw image/video upload is not supported by relit_inference.py alone.")


def prepare_env(
    env_mode: str,
    env_name: Optional[str],
    env_zip: Optional[Any],
    env_root: Path,
    extract_root: Path,
) -> Tuple[Path, str]:
    safe_rmtree(env_root)

    if env_mode == "Repo environment":
        if not env_name:
            raise ValueError("Select a repo environment.")
        path = ENV_ROOT / env_name
        valid, missing = is_valid_env_dir(path)
        if not valid:
            raise ValueError(f"Invalid environment. Missing: {', '.join(missing)}")
        return path, f"Using repo environment: {env_name}"

    if env_mode == "Custom environment ZIP":
        if env_zip is None:
            raise ValueError("Upload a prepared environment ZIP.")
        extracted = extract_zip_to(env_zip, extract_root)
        path, msg = normalize_zip_env(extracted, env_root)
        return path, msg

    raise ValueError("Invalid environment mode.")


def build_command(
    preset_name: str,
    dataset_path: Path,
    env_path: Path,
    output_dir: Path,
    output_path: Path,
    steps: int,
    cfg_scale: float,
    quality: int,
    wo_ref_weight: float,
    drop_mr: bool,
    use_multi_ref: bool,
) -> list[str]:
    preset = PRESETS[preset_name]
    checkpoint_path = CHECKPOINTS[preset["checkpoint"]]

    cmd = [
        sys.executable,
        str(ROOT / "relit_inference.py"),
        "--dataset_path",
        str(dataset_path),
        "--ckpt_path",
        str(checkpoint_path),
        "--output_dir",
        str(output_dir),
        "--output_path",
        str(output_path),
        "--cfg_scale",
        str(cfg_scale),
        "--height",
        str(preset["height"]),
        "--width",
        str(preset["width"]),
        "--num_frames",
        str(preset["frames"]),
        "--padding_resolution",
        "--use_ref_image",
        "--env_map_path",
        str(env_path),
        "--frame_interval",
        "1",
        "--num_inference_steps",
        str(steps),
        "--quality",
        str(quality),
        "--wo_ref_weight",
        str(wo_ref_weight),
        "--dataloader_num_workers",
        "0",
    ]

    cmd.extend(preset["flags"])

    if drop_mr:
        cmd.append("--drop_mr")
    if use_multi_ref:
        cmd.append("--use_muti_ref_image")

    return cmd


def find_diagnostic(output_dir: Path, output_path: Path, kind: str) -> Optional[str]:
    if not output_dir.exists():
        return None

    if kind == "image":
        candidates = sorted(p for p in output_dir.rglob("*.png") if p != output_path and "_render" not in p.name)
    else:
        candidates = sorted(p for p in output_dir.rglob("*.mp4") if p != output_path and "_video" not in p.name)

    return str(candidates[0]) if candidates else None


STARTUP_OK, STARTUP_LOG = startup_preflight()


def estimate_gpu_duration_large(
    preset_name: str,
    source_mode: str,
    sample_name: Optional[str],
    dataset_zip: Optional[Any],
    env_mode: str,
    env_name: Optional[str],
    env_zip: Optional[Any],
    steps: int,
    cfg_scale: float,
    quality: int,
    wo_ref_weight: float,
    drop_mr: bool,
    use_multi_ref: bool,
    progress: Optional[Any] = None,
) -> int:
    preset = PRESETS.get(preset_name, PRESETS["Basic 25-frame relighting"])
    frames = int(preset["frames"])
    steps = int(steps)

    if frames == 1:
        return min(210, max(60, 40 + steps * 3))
    if frames <= 25:
        return min(900, max(180, 90 + steps * 10))
    return min(1500, max(300, 180 + steps * 18))


def estimate_gpu_duration_xlarge(
    preset_name: str,
    source_mode: str,
    sample_name: Optional[str],
    dataset_zip: Optional[Any],
    env_mode: str,
    env_name: Optional[str],
    env_zip: Optional[Any],
    steps: int,
    cfg_scale: float,
    quality: int,
    wo_ref_weight: float,
    drop_mr: bool,
    use_multi_ref: bool,
    progress: Optional[Any] = None,
) -> int:
    preset = PRESETS.get(preset_name, PRESETS["Basic 25-frame relighting"])
    frames = int(preset["frames"])
    steps = int(steps)

    if frames == 1:
        return min(150, max(45, 30 + steps * 2))
    if frames <= 25:
        return min(720, max(150, 75 + steps * 8))
    return min(1200, max(240, 150 + steps * 14))


def _run_inference_impl(
    gpu_size_label: str,
    preset_name: str,
    source_mode: str,
    sample_name: Optional[str],
    dataset_zip: Optional[Any],
    env_mode: str,
    env_name: Optional[str],
    env_zip: Optional[Any],
    steps: int,
    cfg_scale: float,
    quality: int,
    wo_ref_weight: float,
    drop_mr: bool,
    use_multi_ref: bool,
    progress: Optional[Any] = None,
) -> Dict[str, Any]:
    progress = MonotonicProgress(progress)
    logs = [STARTUP_LOG]

    try:
        update_progress(progress, 0.01, "Checking setup")
        if not STARTUP_OK:
            raise RuntimeError("Startup preflight failed. See startup log above.")

        if preset_name not in PRESETS:
            raise ValueError("Select a README inference preset.")

        preset = PRESETS[preset_name]
        kind = preset["output_kind"]

        patch_transformers_imports()
        update_progress(progress, 0.04, "Preparing checkpoint")
        ensure_checkpoint(preset["checkpoint"], logs)

        key = hash_key(
            BUILD_ID,
            preset_name,
            source_mode,
            sample_name or "uploaded",
            env_mode,
            env_name or "custom-env",
            steps,
            cfg_scale,
            quality,
            wo_ref_weight,
            drop_mr,
            use_multi_ref,
        )

        job_root = JOBS_ROOT / key
        dataset_root = job_root / "dataset"
        dataset_extract_root = job_root / "dataset_extract"
        env_root = job_root / "custom_env"
        env_extract_root = job_root / "env_extract"
        output_dir = OUTPUTS_ROOT / key
        output_path = output_dir / ("result.png" if kind == "image" else "result.mp4")

        logs.append(f"Cache key: {key}")
        logs.append(f"Preset: {preset_name}")
        logs.append(f"ZeroGPU size: {gpu_size_label}")
        logs.append(f"Expected pure output: {output_path}")

        if output_path.exists() and output_path.stat().st_size > 0:
            update_progress(progress, 1.0, "Returning cached result")
            logs.append("Returning cached result.")
            diagnostic = find_diagnostic(output_dir, output_path, kind)
            return {
                "ok": True,
                "kind": kind,
                "output": str(output_path),
                "diagnostic": diagnostic,
                "log": "\n\n".join(logs),
            }

        safe_rmtree(job_root)
        safe_rmtree(output_dir)
        job_root.mkdir(parents=True, exist_ok=True)
        output_dir.mkdir(parents=True, exist_ok=True)

        update_progress(progress, 0.07, "Preparing inputs")
        dataset_path, dataset_msg = prepare_dataset(
            source_mode=source_mode,
            sample_name=sample_name,
            dataset_zip=dataset_zip,
            dataset_root=dataset_root,
            extract_root=dataset_extract_root,
        )
        env_path, env_msg = prepare_env(
            env_mode=env_mode,
            env_name=env_name,
            env_zip=env_zip,
            env_root=env_root,
            extract_root=env_extract_root,
        )

        logs.append(dataset_msg)
        logs.append(env_msg)

        update_progress(progress, 0.10, "Running GPU preflight")
        code, out, err = run_python_preflight()
        logs.append(f"GPU preflight exit code: {code}")
        logs.append("STDOUT:\n" + tail_text(out, 8000))
        logs.append("STDERR:\n" + tail_text(err, 8000))
        if code != 0:
            raise RuntimeError("GPU preflight failed.")

        cmd = build_command(
            preset_name=preset_name,
            dataset_path=dataset_path,
            env_path=env_path,
            output_dir=output_dir,
            output_path=output_path,
            steps=int(steps),
            cfg_scale=float(cfg_scale),
            quality=int(quality),
            wo_ref_weight=float(wo_ref_weight),
            drop_mr=bool(drop_mr),
            use_multi_ref=bool(use_multi_ref),
        )

        logs.append("Command:")
        logs.append(" ".join(cmd))

        env = os.environ.copy()
        env["PYTHONPATH"] = str(ROOT) + os.pathsep + env.get("PYTHONPATH", "")
        returncode, stdout, stderr = run_inference_subprocess(
            cmd,
            env=env,
            progress=progress,
        )

        logs.append(f"relit_inference.py exit code: {returncode}")
        logs.append("STDOUT tail:\n" + tail_text(stdout))
        logs.append("STDERR tail:\n" + tail_text(stderr))

        if returncode != 0:
            raise RuntimeError("relit_inference.py failed.")

        if not output_path.exists() or output_path.stat().st_size == 0:
            logs.append("Files under output_dir:")
            for path in sorted(output_dir.rglob("*")):
                if path.is_file():
                    logs.append(f"- {path.relative_to(output_dir)} ({path.stat().st_size} bytes)")
            raise RuntimeError("Explicit --output_path was not created.")

        diagnostic = find_diagnostic(output_dir, output_path, kind)
        logs.append(f"Pure output ready: {output_path}")
        if diagnostic:
            logs.append(f"Diagnostic sheet ready: {diagnostic}")
        update_progress(progress, 1.0, "Done")

        return {
            "ok": True,
            "kind": kind,
            "output": str(output_path),
            "diagnostic": diagnostic,
            "log": "\n\n".join(logs),
        }

    except Exception as exc:
        logs.append(f"Error: {repr(exc)}")
        return {
            "ok": False,
            "kind": "video",
            "output": None,
            "diagnostic": None,
            "log": "\n\n".join(logs),
        }


@spaces.GPU(duration=estimate_gpu_duration_large, size="large")
def run_inference_large(
    preset_name: str,
    source_mode: str,
    sample_name: Optional[str],
    dataset_zip: Optional[Any],
    env_mode: str,
    env_name: Optional[str],
    env_zip: Optional[Any],
    steps: int,
    cfg_scale: float,
    quality: int,
    wo_ref_weight: float,
    drop_mr: bool,
    use_multi_ref: bool,
    progress: Optional[Any] = None,
) -> Dict[str, Any]:
    return _run_inference_impl(
        "large",
        preset_name,
        source_mode,
        sample_name,
        dataset_zip,
        env_mode,
        env_name,
        env_zip,
        steps,
        cfg_scale,
        quality,
        wo_ref_weight,
        drop_mr,
        use_multi_ref,
        progress,
    )


@spaces.GPU(duration=estimate_gpu_duration_xlarge, size="xlarge")
def run_inference_xlarge(
    preset_name: str,
    source_mode: str,
    sample_name: Optional[str],
    dataset_zip: Optional[Any],
    env_mode: str,
    env_name: Optional[str],
    env_zip: Optional[Any],
    steps: int,
    cfg_scale: float,
    quality: int,
    wo_ref_weight: float,
    drop_mr: bool,
    use_multi_ref: bool,
    progress: Optional[Any] = None,
) -> Dict[str, Any]:
    return _run_inference_impl(
        "xlarge",
        preset_name,
        source_mode,
        sample_name,
        dataset_zip,
        env_mode,
        env_name,
        env_zip,
        steps,
        cfg_scale,
        quality,
        wo_ref_weight,
        drop_mr,
        use_multi_ref,
        progress,
    )


def run_ui(
    gpu_size,
    preset_name,
    source_mode,
    sample_name,
    dataset_zip,
    env_mode,
    env_name,
    env_zip,
    steps,
    cfg_scale,
    quality,
    wo_ref_weight,
    drop_mr,
    use_multi_ref,
    progress=gr.Progress(track_tqdm=False),
):
    runner = run_inference_xlarge if gpu_size == "Fast - xlarge" else run_inference_large
    update_progress(progress, 0.0, "Queued")
    result = runner(
        preset_name,
        source_mode,
        sample_name,
        dataset_zip,
        env_mode,
        env_name,
        env_zip,
        steps,
        cfg_scale,
        quality,
        wo_ref_weight,
        drop_mr,
        use_multi_ref,
        progress,
    )

    output = result.get("output")
    diagnostic = result.get("diagnostic")
    kind = result.get("kind")
    ok = result.get("ok")
    log = ("OK\n\n" if ok else "FAILED\n\n") + result.get("log", "")

    image_value = output if kind == "image" else None
    video_value = output if kind == "video" else None
    diag_image = diagnostic if diagnostic and diagnostic.endswith(".png") else None
    diag_video = diagnostic if diagnostic and diagnostic.endswith(".mp4") else None

    # Keep the UI aligned with the selected preset.
    # The unused output column is hidden and cleared immediately.
    preset_kind = PRESETS.get(preset_name, {}).get("output_kind", kind or "video")
    is_image = preset_kind == "image"

    return (
        gr.update(visible=is_image),
        gr.update(visible=not is_image),
        gr.update(visible=is_image),
        gr.update(visible=not is_image),
        gr.update(value=image_value if is_image else None, visible=is_image),
        gr.update(value=video_value if not is_image else None, visible=not is_image),
        gr.update(value=diag_image if is_image else None, visible=is_image),
        gr.update(value=diag_video if not is_image else None, visible=not is_image),
        log,
    )


def update_preset_info(preset_name: str):
    preset = PRESETS[preset_name]
    ext = "PNG" if preset["output_kind"] == "image" else "MP4"
    flags = " ".join(preset["flags"]) if preset["flags"] else "none"
    text = (
        f"Checkpoint: {preset['checkpoint']}\n"
        f"Resolution: {preset['height']}x{preset['width']}\n"
        f"Frames: {preset['frames']}\n"
        f"README flags: {flags}\n"
        f"Pure output: {ext}"
    )
    return text


def update_preset_ui(preset_name: str):
    preset_info_text = update_preset_info(preset_name)
    is_image = PRESETS[preset_name]["output_kind"] == "image"

    return (
        preset_info_text,
        gr.update(visible=is_image),
        gr.update(visible=not is_image),
        gr.update(visible=is_image),
        gr.update(visible=not is_image),
        gr.update(value=None, visible=is_image),
        gr.update(value=None, visible=not is_image),
        gr.update(value=None, visible=is_image),
        gr.update(value=None, visible=not is_image),
    )


def update_sample_preview(sample_name):
    value = sample_preview(sample_name)
    return gr.update(value=value, visible=value is not None)


def update_env_preview(env_name):
    value = env_preview(env_name)
    return gr.update(value=value, visible=value is not None)


def update_source_mode(source_mode):
    return (
        gr.update(visible=source_mode == "Repo demo sample"),
        gr.update(visible=source_mode == "Prepared dataset ZIP"),
    )


def update_env_mode(env_mode):
    return (
        gr.update(visible=env_mode == "Repo environment"),
        gr.update(visible=env_mode == "Custom environment ZIP"),
    )


SAMPLES = list_demo_samples()
ENVS = list_envs()
DEFAULT_SAMPLE = SAMPLES[0] if SAMPLES else None
DEFAULT_ENV = "Pink_Sunrise" if "Pink_Sunrise" in ENVS else (ENVS[0] if ENVS else None)
DEFAULT_PRESET = "Basic 25-frame relighting"
DEFAULT_KIND = PRESETS[DEFAULT_PRESET]["output_kind"]


CSS = """
.gradio-container { max-width: 1600px !important; width: 98vw !important; margin: 0 auto !important; }
.app-intro { max-width: 900px; margin: 0 auto 1rem auto; text-align: center; }
#main-layout,
.main-layout {
    display: grid !important;
    grid-template-columns: minmax(340px, 430px) minmax(620px, 1fr) !important;
    width: 100% !important;
    max-width: 1540px !important;
    margin: 0 auto !important;
    column-gap: 24px !important;
    row-gap: 0 !important;
    align-items: flex-start !important;
}
#main-layout > *,
.main-layout > *,
.settings-col,
.results-col {
    min-width: 0 !important;
    width: 100% !important;
    max-width: none !important;
    flex: none !important;
}
#main-layout > *:first-child,
.settings-col {
    grid-column: 1 !important;
}
#main-layout > *:nth-child(2),
.results-col {
    grid-column: 2 !important;
}
.settings-col .block, .results-col .block { margin-bottom: 10px !important; }
.small-note { color: #6b7280; font-size: 0.92rem; }
textarea { font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", monospace !important; }
@media (max-width: 980px) {
    #main-layout,
    .main-layout {
        display: block !important;
    }
    #main-layout > *,
    .main-layout > *,
    .settings-col,
    .results-col { max-width: none !important; width: 100% !important; }
}
"""


with gr.Blocks(title="Relit-LiVE Demo", css=CSS) as demo:
    gr.Markdown(
        f"""
# Relit-LiVE Demo

Try the official inference modes with demo samples, custom prepared datasets, and relighting environments.

This ZeroGPU demo runs the official Relit-LiVE `relit_inference.py` path on prepared demo samples.

It does not run the full Cosmos inverse-rendering pipeline and does not support raw arbitrary video uploads.
""",
        elem_classes=["app-intro"],
    )

    with gr.Row(elem_id="main-layout", elem_classes=["main-layout"]):
        with gr.Column(scale=1, min_width=260, elem_classes=["settings-col"]):
            gr.Markdown("## Inference preset")
            preset = gr.Dropdown(
                label="README case",
                choices=list(PRESETS.keys()),
                value=DEFAULT_PRESET,
            )
            preset_info = gr.Textbox(
                label="Preset details",
                value=update_preset_info(DEFAULT_PRESET),
                lines=6,
                interactive=False,
            )

            gr.Markdown("## Input sample")
            source_mode = gr.Radio(
                label="Source",
                choices=["Repo demo sample", "Prepared dataset ZIP", "Raw image/video upload"],
                value="Repo demo sample",
            )
            with gr.Column(visible=True) as repo_sample_group:
                sample_name = gr.Dropdown(
                    label="Repo demo sample",
                    choices=SAMPLES,
                    value=DEFAULT_SAMPLE,
                )
                sample_img = gr.Image(
                    label="Sample preview",
                    value=sample_preview(DEFAULT_SAMPLE) if DEFAULT_SAMPLE else None,
                    visible=DEFAULT_SAMPLE is not None,
                    height=220,
                )
            with gr.Column(visible=False) as dataset_zip_group:
                dataset_zip = gr.File(
                    label="Prepared Relit-LiVE dataset ZIP",
                    file_types=[".zip"],
                )
                gr.Markdown(
                    """
Expected ZIP layout: either a sample folder at the ZIP root, or one or more sample folders.
Each sample must contain `images_4`, `Base Color`, `depth`, and `normal`.
`Metallic` and `Roughness` are recommended.
""",
                    elem_classes=["small-note"],
                )
            raw_note = gr.Markdown(
                """
Raw video upload needs inverse rendering first. Use the full Cosmos/Module1 pipeline, or upload a
prepared dataset ZIP containing the required maps.
""",
                visible=False,
                elem_classes=["small-note"],
            )

            gr.Markdown("## Environment")
            env_mode = gr.Radio(
                label="Environment source",
                choices=["Repo environment", "Custom environment ZIP"],
                value="Repo environment",
            )
            with gr.Column(visible=True) as repo_env_group:
                env_name = gr.Dropdown(
                    label="Repo environment",
                    choices=ENVS,
                    value=DEFAULT_ENV,
                )
                env_video = gr.Video(
                    label="Environment preview",
                    value=env_preview(DEFAULT_ENV) if DEFAULT_ENV else None,
                    visible=DEFAULT_ENV is not None,
                    height=170,
                )
            with gr.Column(visible=False) as custom_env_group:
                env_zip = gr.File(
                    label="Custom environment ZIP",
                    file_types=[".zip"],
                )
                gr.Markdown(
                    "Expected files: `ldr_video_fix_first_frame.mp4`, "
                    "`hdr_log_video_fix_first_frame.mp4`, `env_dir_video_fix_first_frame.mp4`.",
                    elem_classes=["small-note"],
                )

        with gr.Column(scale=1, min_width=320, elem_classes=["results-col"]):
            gr.Markdown("## Run")
            gpu_size = gr.Radio(
                label="ZeroGPU speed",
                choices=["Eco - large", "Fast - xlarge"],
                value="Eco - large",
                info="xlarge can be faster but uses 2x ZeroGPU quota.",
            )
            steps = gr.Slider(
                label="Quality / steps",
                minimum=8,
                maximum=50,
                step=1,
                value=8,
            )
            with gr.Accordion("Advanced settings", open=False):
                cfg_scale = gr.Slider(
                    label="CFG scale",
                    minimum=0.5,
                    maximum=5.0,
                    step=0.1,
                    value=1.0,
                )
                quality = gr.Slider(
                    label="MP4 quality",
                    minimum=5,
                    maximum=10,
                    step=1,
                    value=10,
                )
                wo_ref_weight = gr.Slider(
                    label="Without-reference branch weight",
                    minimum=0.0,
                    maximum=5.0,
                    step=0.1,
                    value=0.0,
                )
                drop_mr = gr.Checkbox(
                    label="Drop metallic/roughness conditioning",
                    value=False,
                )
                use_multi_ref = gr.Checkbox(
                    label="Use multi-reference image mode",
                    value=False,
                )
            run_btn = gr.Button("Run inference", variant="primary", size="lg")

            gr.Markdown("## Pure result")
            with gr.Column(visible=DEFAULT_KIND == "image") as image_result_col:
                result_image = gr.Image(
                    label="Result image",
                    visible=DEFAULT_KIND == "image",
                    height=430,
                )
            with gr.Column(visible=DEFAULT_KIND == "video") as video_result_col:
                result_video = gr.Video(
                    label="Result video",
                    visible=DEFAULT_KIND == "video",
                    height=430,
                )

            with gr.Accordion("Diagnostic sheet", open=False):
                with gr.Column(visible=DEFAULT_KIND == "image") as diag_image_col:
                    diag_image = gr.Image(
                        label="Diagnostic image",
                        visible=DEFAULT_KIND == "image",
                        height=260,
                    )
                with gr.Column(visible=DEFAULT_KIND == "video") as diag_video_col:
                    diag_video = gr.Video(
                        label="Diagnostic video",
                        visible=DEFAULT_KIND == "video",
                        height=260,
                    )

            with gr.Accordion("Logs", open=False):
                logs = gr.Textbox(
                    label="Runtime logs",
                    value=STARTUP_LOG,
                    lines=18,
                    max_lines=60,
                    autoscroll=True,
                )

    preset.change(
        update_preset_ui,
        inputs=[preset],
        outputs=[
            preset_info,
            image_result_col,
            video_result_col,
            diag_image_col,
            diag_video_col,
            result_image,
            result_video,
            diag_image,
            diag_video,
        ],
    )
    sample_name.change(update_sample_preview, inputs=[sample_name], outputs=[sample_img])
    env_name.change(update_env_preview, inputs=[env_name], outputs=[env_video])
    source_mode.change(
        update_source_mode,
        inputs=[source_mode],
        outputs=[repo_sample_group, dataset_zip_group],
    ).then(
        lambda mode: gr.update(visible=mode == "Raw image/video upload"),
        inputs=[source_mode],
        outputs=[raw_note],
    )
    env_mode.change(
        update_env_mode,
        inputs=[env_mode],
        outputs=[repo_env_group, custom_env_group],
    )

    run_btn.click(
        run_ui,
        inputs=[
            gpu_size,
            preset,
            source_mode,
            sample_name,
            dataset_zip,
            env_mode,
            env_name,
            env_zip,
            steps,
            cfg_scale,
            quality,
            wo_ref_weight,
            drop_mr,
            use_multi_ref,
        ],
        outputs=[
            image_result_col,
            video_result_col,
            diag_image_col,
            diag_video_col,
            result_image,
            result_video,
            diag_image,
            diag_video,
            logs,
        ],
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--server-name", default="0.0.0.0")
    parser.add_argument("--server-port", type=int, default=int(os.getenv("PORT", "7860")))
    args = parser.parse_args()
    demo.queue(default_concurrency_limit=1, max_size=8).launch(
        server_name=args.server_name,
        server_port=args.server_port,
        show_error=True,
    )