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