Update app.py

app.py

@@ -1,375 +1,988 @@
-import gradio as gr
+import argparse
+import hashlib
 import os
-import sys
 import shutil
-import hashlib
 import subprocess
+import sys
+import tempfile
+import zipfile
 from pathlib import Path
-from typing import Tuple, Optional
+from typing import Any, Dict, Optional, Tuple
 
+import gradio as gr
 import spaces
 from huggingface_hub import snapshot_download
 
-BUILD_ID = "relit-live-zerogpu-simple-gradio-v2-transformers-patch"
+
+BUILD_ID = "relit-live-readme-gradio-v4"
 
 ROOT = Path(__file__).resolve().parent
 DEMO_ROOT = ROOT / "datasets" / "demos"
 ENV_ROOT = ROOT / "datasets" / "envs"
-CHECKPOINT_PATH = ROOT / "checkpoints" / "model_frame25_480_832.ckpt"
+
 WAN_DIR = ROOT / "models" / "Wan-AI" / "Wan2.1-T2V-1.3B"
-RUNTIME_ROOT = ROOT / "runtime_simple_demo"
+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"}
-IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}
 
-for p in [RUNTIME_ROOT, JOBS_ROOT, OUTPUTS_ROOT]:
-    p.mkdir(parents=True, exist_ok=True)
-
-
-def list_demo_samples():
-    if not DEMO_ROOT.exists():
-        return []
-    return sorted([p.name for p in DEMO_ROOT.iterdir() if p.is_dir()])
+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 list_envs():
-    if not ENV_ROOT.exists():
-        return []
-    return sorted([p.name for p in ENV_ROOT.iterdir() if p.is_dir()])
+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):
+def safe_rmtree(path: Path) -> None:
     if path.exists():
         shutil.rmtree(path, ignore_errors=True)
 
 
-def hash_key(*items) -> str:
-    raw = "||".join(map(str, items))
-    return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16]
+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
 
 
-def tail_text(text: str, max_chars: int = 16000) -> str:
-    if not text:
-        return ""
-    if len(text) <= max_chars:
-        return text
-    return text[-max_chars:]
+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)))
 
-def patch_transformers_compat() -> list[str]:
-    """Patch the local DiffSynth import that breaks with modern Transformers.
+    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
 
-    The repo imports PretrainedConfig from transformers.modeling_utils. In current
-    Transformers, PretrainedConfig lives in transformers.configuration_utils / public API.
-    We patch the source file and also write sitecustomize.py as a subprocess fallback.
-    """
-    logs = []
 
-    bad = "from transformers.modeling_utils import PretrainedConfig, PreTrainedModel"
-    good = "from transformers.configuration_utils import PretrainedConfig\nfrom transformers.modeling_utils import PreTrainedModel"
-
-    patched_files = []
-    diffsynth_dir = ROOT / "diffsynth"
-    if diffsynth_dir.exists():
-        for py_file in diffsynth_dir.rglob("*.py"):
-            try:
-                text = py_file.read_text(encoding="utf-8")
-            except Exception:
-                continue
-            if bad in text:
-                py_file.write_text(text.replace(bad, good), encoding="utf-8")
-                patched_files.append(str(py_file.relative_to(ROOT)))
-
-    if patched_files:
-        logs.append("Patched DiffSynth Transformers import in:\n- " + "\n- ".join(patched_files))
-    else:
-        logs.append("DiffSynth Transformers import patch: no direct source replacement needed or already 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.",
+    ]
 
-    # Fallback for any remaining legacy import in subprocesses.
-    sitecustomize = ROOT / "sitecustomize.py"
-    marker = "# relit-live transformers compatibility patch"
-    snippet = f'''{marker}\ntry:\n    import transformers.modeling_utils as _relit_mu\n    from transformers.configuration_utils import PretrainedConfig as _RelitPretrainedConfig\n    if not hasattr(_relit_mu, "PretrainedConfig"):\n        _relit_mu.PretrainedConfig = _RelitPretrainedConfig\nexcept Exception:\n    pass\n'''
     try:
-        existing = sitecustomize.read_text(encoding="utf-8") if sitecustomize.exists() else ""
-        if marker not in existing:
-            sitecustomize.write_text(snippet + "\n" + existing, encoding="utf-8")
-            logs.append(f"Wrote subprocess fallback: {sitecustomize.relative_to(ROOT)}")
-        else:
-            logs.append("Subprocess fallback sitecustomize.py already present.")
-    except Exception as e:
-        logs.append(f"Warning: failed to write sitecustomize.py fallback: {repr(e)}")
+        patched = patch_transformers_imports()
+        logs.append("Compatibility patch files:")
+        for path in patched:
+            logs.append(f"- {path}")
 
-    return logs
+        if not (ROOT / "relit_inference.py").exists():
+            return False, "\n".join(logs + ["Missing relit_inference.py at repo root."])
 
+        ensure_wan_model(logs)
 
-def make_env() -> dict:
-    env = os.environ.copy()
-    env["PYTHONPATH"] = str(ROOT) + os.pathsep + env.get("PYTHONPATH", "")
-    env.setdefault("TRANSFORMERS_NO_ADVISORY_WARNINGS", "1")
-    env.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
-    return env
+        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))
 
-def run_preflight_import() -> str:
-    code = r'''
-import sys
-print("python", sys.version)
-import torch
+        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")
-print("capability", torch.cuda.get_device_capability(0) if torch.cuda.is_available() else "N/A")
 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=make_env(),
+        env=env,
         capture_output=True,
         text=True,
     )
-    return (
-        f"Relit/DiffSynth preflight exit code: {proc.returncode}\n"
-        f"STDOUT:\n{tail_text(proc.stdout)}\n"
-        f"STDERR:\n{tail_text(proc.stderr)}"
-    )
+    return proc.returncode, proc.stdout, proc.stderr
 
 
-def ensure_runtime_assets() -> Tuple[bool, str]:
-    logs = [f"Build: {BUILD_ID}", "Mode: README relit_inference.py only. Cosmos is not used."]
-    logs.extend(patch_transformers_compat())
-    token = os.getenv("HF_TOKEN", None)
+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
 
-    try:
-        if CHECKPOINT_PATH.exists():
-            logs.append("Relit-LiVE 25-frame checkpoint already present.")
-        else:
-            logs.append("Downloading Relit-LiVE 25-frame checkpoint...")
-            snapshot_download(
-                repo_id="weiqingXiao/Relit-LiVE",
-                local_dir=str(ROOT),
-                allow_patterns=["checkpoints/model_frame25_480_832.ckpt"],
-                token=token,
-            )
-            if not CHECKPOINT_PATH.exists():
-                return False, "\n".join(logs + ["❌ Failed to fetch model_frame25_480_832.ckpt"])
-            logs.append("Relit-LiVE 25-frame checkpoint ready.")
-
-        if WAN_DIR.exists() and any(WAN_DIR.rglob("*")):
-            logs.append("Wan2.1 base model already present.")
-        else:
-            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=token,
-            )
-            logs.append("Wan2.1 base model ready.")
 
-        if not (ROOT / "relit_inference.py").exists():
-            return False, "\n".join(logs + ["❌ relit_inference.py not found at repo root"])
-        if not DEMO_ROOT.exists():
-            return False, "\n".join(logs + [f"❌ Missing demo root: {DEMO_ROOT}"])
-        if not ENV_ROOT.exists():
-            return False, "\n".join(logs + [f"❌ Missing env root: {ENV_ROOT}"])
+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
 
-        samples = list_demo_samples()
-        envs = list_envs()
-        logs.append(f"Found {len(samples)} demo sample(s).")
-        logs.append(f"Found {len(envs)} environment(s).")
-        if not samples:
-            return False, "\n".join(logs + ["❌ No demo samples found in datasets/demos"])
-        if not envs:
-            return False, "\n".join(logs + ["❌ No environments found in datasets/envs"])
-
-        # Do an import preflight at startup so dependency errors appear before inference.
-        logs.append(run_preflight_import())
-
-        return True, "\n".join(logs)
-    except Exception as e:
-        logs.append(f"❌ Runtime asset preparation failed: {repr(e)}")
-        return False, "\n".join(logs)
 
+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
 
-STARTUP_OK, STARTUP_LOG = ensure_runtime_assets()
 
+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 detect_output(output_dir: Path) -> Tuple[Optional[Path], Optional[str]]:
-    if not output_dir.exists():
-        return None, None
-    files = [p for p in output_dir.rglob("*") if p.is_file()]
-    if not files:
-        return None, None
-    video_files = sorted([p for p in files if p.suffix.lower() in VIDEO_EXTS])
-    if video_files:
-        return video_files[0], "video"
-    image_files = sorted([p for p in files if p.suffix.lower() in IMAGE_EXTS])
-    if image_files:
-        return image_files[0], "image"
-    return None, None
-
-
-def build_command(dataset_path: Path, env_path: Path, output_dir: Path, num_frames: int, steps: int, cfg_scale: float):
-    return [
+
+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),
-        "--cfg_scale", str(cfg_scale),
-        "--height", "480",
-        "--width", "832",
-        "--num_frames", str(num_frames),
+        "--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", "10",
+        "--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"])
 
-def make_job_dataset(sample_name: str, job_dataset_root: Path) -> Path:
-    src = DEMO_ROOT / sample_name
-    if not src.exists():
-        raise FileNotFoundError(f"Sample not found: {src}")
-    safe_rmtree(job_dataset_root)
-    job_dataset_root.mkdir(parents=True, exist_ok=True)
-    shutil.copytree(src, job_dataset_root / sample_name)
-    return job_dataset_root
+    if drop_mr:
+        cmd.append("--drop_mr")
+    if use_multi_ref:
+        cmd.append("--use_muti_ref_image")
 
+    return cmd
 
-@spaces.GPU(duration=600, size="xlarge")
-def run_relit_cli(sample_name: str, env_name: str, num_frames: int, steps: int, cfg_scale: float):
-    # Re-apply the patch inside the worker process. Some ZeroGPU workers are separate processes.
-    local_logs = []
-    local_logs.extend(patch_transformers_compat())
-    logs = [STARTUP_LOG, "\n".join(local_logs)]
 
-    if not STARTUP_OK:
-        logs.append("❌ Startup preparation failed. Aborting inference.")
-        return None, None, "\n\n".join(logs)
+def find_diagnostic(output_dir: Path, output_path: Path, kind: str) -> Optional[str]:
+    if not output_dir.exists():
+        return None
 
-    if not sample_name:
-        return None, None, "\n\n".join(logs + ["❌ No sample selected."])
-    if not env_name:
-        return None, None, "\n\n".join(logs + ["❌ No environment selected."])
+    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)
 
-    env_path = ENV_ROOT / env_name
-    if not env_path.exists():
-        return None, None, "\n\n".join(logs + [f"❌ Environment not found: {env_path}"])
+    return str(candidates[0]) if candidates else None
 
-    key = hash_key(sample_name, env_name, num_frames, steps, cfg_scale, BUILD_ID)
-    logs.append(f"Cache key: {key}")
 
-    output_dir = OUTPUTS_ROOT / key
-    existing_output, existing_kind = detect_output(output_dir)
-    if existing_output is not None:
-        logs.append("✅ Returning cached result.")
-        return (None, str(existing_output), "\n\n".join(logs)) if existing_kind == "video" else (str(existing_output), None, "\n\n".join(logs))
+STARTUP_OK, STARTUP_LOG = startup_preflight()
 
-    job_root = JOBS_ROOT / key
-    dataset_root = job_root / "dataset"
-    safe_rmtree(job_root)
-    safe_rmtree(output_dir)
-    job_root.mkdir(parents=True, exist_ok=True)
-    output_dir.mkdir(parents=True, exist_ok=True)
+
+@spaces.GPU(duration=3600, size="xlarge")
+def run_inference(
+    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,
+) -> Dict[str, Any]:
+    logs = [STARTUP_LOG]
 
     try:
-        make_job_dataset(sample_name, dataset_root)
-    except Exception as e:
-        return None, None, "\n\n".join(logs + [f"❌ Failed to prepare job dataset: {repr(e)}"])
+        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()
+        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,
+        )
 
-    cmd = build_command(dataset_root, env_path, output_dir, int(num_frames), int(steps), float(cfg_scale))
-    logs.append("Command:\n" + " ".join(cmd))
+        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"Expected pure output: {output_path}")
+
+        if output_path.exists() and output_path.stat().st_size > 0:
+            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)
+
+        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,
+        )
 
-    preflight = run_preflight_import()
-    logs.append(preflight)
-    if "Relit/DiffSynth preflight exit code: 0" not in preflight:
-        return None, None, "❌ Dependency preflight failed.\n\n" + "\n\n".join(logs)
+        logs.append(dataset_msg)
+        logs.append(env_msg)
+
+        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),
+        )
 
-    proc = subprocess.run(
-        cmd,
-        cwd=str(ROOT),
-        env=make_env(),
-        capture_output=True,
-        text=True,
+        logs.append("Command:")
+        logs.append(" ".join(cmd))
+
+        env = os.environ.copy()
+        env["PYTHONPATH"] = str(ROOT) + os.pathsep + env.get("PYTHONPATH", "")
+        proc = subprocess.run(
+            cmd,
+            cwd=str(ROOT),
+            env=env,
+            capture_output=True,
+            text=True,
+        )
+
+        logs.append(f"relit_inference.py exit code: {proc.returncode}")
+        logs.append("STDOUT tail:\n" + tail_text(proc.stdout))
+        logs.append("STDERR tail:\n" + tail_text(proc.stderr))
+
+        if proc.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}")
+
+        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),
+        }
+
+
+def run_ui(
+    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,
+):
+    result = run_inference(
+        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,
     )
 
-    logs.append(f"relit_inference.py exit code: {proc.returncode}")
-    logs.append("STDOUT tail:\n" + tail_text(proc.stdout))
-    logs.append("STDERR tail:\n" + tail_text(proc.stderr))
+    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", "")
 
-    if proc.returncode != 0:
-        return None, None, "❌ Inference failed.\n\n" + "\n\n".join(logs)
+    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
 
-    output_path, output_kind = detect_output(output_dir)
-    if output_path is None:
-        logs.append(f"❌ Inference finished but no output file found under {output_dir}")
-        return None, None, "❌ Inference failed.\n\n" + "\n\n".join(logs)
+    return (
+        gr.update(value=image_value, visible=image_value is not None),
+        gr.update(value=video_value, visible=video_value is not None),
+        gr.update(value=diag_image, visible=diag_image is not None),
+        gr.update(value=diag_video, visible=diag_video is not None),
+        log,
+    )
 
-    logs.append(f"✅ Output detected: {output_path}")
-    return (None, str(output_path), "\n\n".join(logs)) if output_kind == "video" else (str(output_path), None, "\n\n".join(logs))
 
+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_sample_preview(sample_name):
+    value = sample_preview(sample_name)
+    return gr.update(value=value, visible=value is not None)
 
-def run_demo(sample_name, env_name, num_frames, steps, cfg_scale):
-    try:
-        image_path, video_path, log = run_relit_cli(sample_name, env_name, int(num_frames), int(steps), float(cfg_scale))
-        return (
-            gr.update(value=image_path, visible=image_path is not None),
-            gr.update(value=video_path, visible=video_path is not None),
-            log,
-        )
-    except Exception as e:
-        return (
-            gr.update(value=None, visible=False),
-            gr.update(value=None, visible=False),
-            f"❌ Gradio wrapper failed.\n\n{STARTUP_LOG}\n\n{repr(e)}",
-        )
+
+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"
 
-DESCRIPTION = f"""
-# Relit-LiVE ZeroGPU Demo
-**Build:** `{BUILD_ID}`
 
-This Space runs the simple README path with `relit_inference.py` on curated samples from the cloned repo.
+CSS = """
+.gradio-container { max-width: 1280px !important; }
+.small-note { color: #6b7280; font-size: 0.92rem; }
+textarea { font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", monospace !important; }
+"""
+
+
+with gr.Blocks(title="Relit-LiVE README Inference", css=CSS) as demo:
+    gr.Markdown(
+        f"""
+# Relit-LiVE README Inference
 
-- Uses `datasets/demos/`
-- Uses `datasets/envs/`
-- Uses Relit-LiVE + Wan2.1
-- Does not use Cosmos
-- Does not support arbitrary uploads
+Build `{BUILD_ID}`. This app exposes the inference cases documented in the repository README:
+basic 25-frame, rotating light, fixed-frame width/height light rotation, 57-frame video,
+and single-frame high-resolution inference.
+
+Raw image/video upload is intentionally not treated as a valid direct input. `relit_inference.py`
+requires a prepared Relit-LiVE sample with RGB frames plus base color, depth, normal, and optional
+metallic/roughness maps. Upload a prepared dataset ZIP for custom content.
 """
+    )
 
-with gr.Blocks(title="Relit-LiVE ZeroGPU Demo") as demo:
-    gr.Markdown(DESCRIPTION)
     with gr.Row():
         with gr.Column(scale=1):
-            sample_name = gr.Dropdown(label="Demo sample", choices=SAMPLES, value=SAMPLES[0] if SAMPLES else None)
-            env_name = gr.Dropdown(label="Environment", choices=ENVS, value=ENVS[0] if ENVS else None)
-            num_frames = gr.Radio(label="Number of frames", choices=[1, 9, 17, 25], value=1)
-            steps = gr.Slider(label="Inference steps", minimum=4, maximum=30, step=1, value=8)
-            cfg_scale = gr.Slider(label="CFG scale", minimum=0.5, maximum=3.0, step=0.1, value=1.0)
-            run_btn = gr.Button("Run relighting", variant="primary")
+            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.Group(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.Group(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.Group(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.Group(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):
-            result_image = gr.Image(label="Result image", visible=False)
-            result_video = gr.Video(label="Result video", visible=False)
-    logs = gr.Textbox(label="Logs", value=STARTUP_LOG, lines=30, max_lines=50, autoscroll=True)
+            gr.Markdown("## Controls")
+            steps = gr.Slider(
+                label="Inference steps",
+                minimum=8,
+                maximum=50,
+                step=1,
+                value=50,
+            )
+            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")
+            result_image = gr.Image(label="Result image", visible=False, height=430)
+            result_video = gr.Video(label="Result video", visible=False, height=430)
+
+            gr.Markdown("## Diagnostic sheet")
+            diag_image = gr.Image(label="Diagnostic image", visible=False, height=260)
+            diag_video = gr.Video(label="Diagnostic video", visible=False, height=260)
+
+    logs = gr.Textbox(
+        label="Logs",
+        value=STARTUP_LOG,
+        lines=26,
+        max_lines=60,
+        autoscroll=True,
+    )
+
+    preset.change(update_preset_info, inputs=[preset], outputs=[preset_info])
+    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(
-        fn=run_demo,
-        inputs=[sample_name, env_name, num_frames, steps, cfg_scale],
-        outputs=[result_image, result_video, logs],
+        run_ui,
+        inputs=[
+            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=[result_image, result_video, diag_image, diag_video, logs],
     )
 
+
 if __name__ == "__main__":
-    demo.queue(default_concurrency_limit=1).launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)), show_error=True)
+    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,
+    )