import spaces APP_BUILD_ID = "relit-live-zerogpu-te-shim-2026-06-29-1908" import gradio as gr import os import subprocess import sys import uuid import shutil import signal import threading from pathlib import Path from PIL import Image from datetime import datetime import socket # -------------------------- Core Configuration -------------------------- BASE_UPLOAD_DIR = "./datasets/gradio_data/upload_data" BASE_RESULT_DIR = "./datasets/gradio_data/results" DEMO_IMAGES_DIR = "./datasets/gradio_data/assets/images_demo" BASH_SCRIPT_PATH = "./tools/full_inference_modules_gradio.sh" SUPPORTED_IMAGE_FORMATS = [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".hdr"] SUPPORTED_VIDEO_FORMATS = [".mp4", ".mov", ".avi", ".mkv"] SUPPORTED_HDR_FORMATS = [".hdr", ".exr", ".jpg", ".png"] ASSETS_ENV_DIR = "./datasets/gradio_data/assets/envs_demo" BUILTIN_ENV_OPTIONS = [] if os.path.exists(ASSETS_ENV_DIR): for ext in SUPPORTED_HDR_FORMATS: BUILTIN_ENV_OPTIONS.extend([f.stem for f in Path(ASSETS_ENV_DIR).glob(f"*{ext}")]) BUILTIN_ENV_OPTIONS = sorted(set(BUILTIN_ENV_OPTIONS)) DEFAULT_ENV = "Pink_Sunrise" if "Pink_Sunrise" in BUILTIN_ENV_OPTIONS else (BUILTIN_ENV_OPTIONS[0] if BUILTIN_ENV_OPTIONS else None) MODULE1_RESULT_TYPES = { "base_color": {"name": "Base Color", "dir": "Base Color", "glob_pattern": "frame_*"}, "normal": {"name": "Normal Map", "dir": "normal", "glob_pattern": "frame_*"}, "roughness": {"name": "Roughness Map", "dir": "Roughness", "glob_pattern": "frame_*"}, } INPUT_PREVIEW_HEIGHT = 240 MODULE1_VIS_HEIGHT = 180 MODULE2_VIS_HEIGHT = 200 MODULE3_VIS_HEIGHT = 400 DEMO_IMAGE_HEIGHT = 200 RUNNING_PROCESSES = {} RUNNING_PROCESSES_LOCK = threading.RLock() def patch_bash_script_for_hf_spaces(): """Disable hard-coded conda activation lines that do not exist in HF Spaces.""" script = Path(BASH_SCRIPT_PATH) if not script.exists(): return try: text = script.read_text(encoding="utf-8") except Exception as exc: print(f"Warning: could not read {script}: {exc}") return original = text patched_lines = [] changed = False for line in text.splitlines(): stripped = line.strip() if "/home/user/miniconda3/bin/activate" in stripped or stripped.startswith("conda activate"): patched_lines.append("# Disabled by app.py for Hugging Face Spaces: " + line) patched_lines.append("true") changed = True else: patched_lines.append(line) if changed: try: script.write_text("\n".join(patched_lines) + "\n", encoding="utf-8") print(f"Patched {script}: disabled hard-coded conda activation.") except Exception as exc: print(f"Warning: could not patch {script}: {exc}") RUNTIME_SETUP_LOCK = threading.RLock() RUNTIME_SETUP_DONE = False RUNTIME_SETUP_ATTEMPTED = False RUNTIME_SETUP_LOG = "" RELIT_CHECKPOINT_REPO = os.environ.get("RELIT_CHECKPOINT_REPO", "weiqingXiao/Relit-LiVE") WAN_REPO = os.environ.get("RELIT_WAN_REPO", "Wan-AI/Wan2.1-T2V-1.3B") COSMOS_REPO = os.environ.get("RELIT_COSMOS_REPO", "https://github.com/nv-tlabs/cosmos-transfer1-diffusion-renderer.git") DOWNLOAD_RELIT_CHECKPOINTS = os.environ.get("RELIT_DOWNLOAD_CHECKPOINTS", "1") != "0" DOWNLOAD_WAN = os.environ.get("RELIT_DOWNLOAD_WAN", "1") != "0" CLONE_COSMOS = os.environ.get("RELIT_CLONE_COSMOS", "1") != "0" STARTUP_RUNTIME_SETUP = os.environ.get("RELIT_STARTUP_SETUP", "1") != "0" INSTALL_COSMOS_DEPS = os.environ.get("RELIT_INSTALL_COSMOS_DEPS", "1") != "0" INSTALL_TRANSFORMER_ENGINE = os.environ.get("RELIT_INSTALL_TRANSFORMER_ENGINE", "0") == "1" INSTALL_NVDIFFRAST = os.environ.get("RELIT_INSTALL_NVDIFFRAST", "0") == "1" DOWNLOAD_COSMOS_CHECKPOINTS = os.environ.get("RELIT_DOWNLOAD_COSMOS_CHECKPOINTS", "1") != "0" STRICT_COSMOS_DEPS = os.environ.get("RELIT_STRICT_COSMOS_DEPS", "0") == "1" USE_TRANSFORMER_ENGINE_SHIM = os.environ.get("RELIT_USE_TRANSFORMER_ENGINE_SHIM", "1") != "0" # -------------------------- Utility Functions -------------------------- def get_server_ip(): try: s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) s.connect(("8.8.8.8", 80)) local_ip = s.getsockname()[0] s.close() return local_ip except Exception: return "Unknown IP" def generate_test_id(flag): timestamp = datetime.now().strftime("%Y%m%d%H%M%S") random_str = uuid.uuid4().hex[:8] return f"{flag}_{timestamp}_{random_str}" def clear_test_dir(test_id): upload_dir = os.path.join(BASE_UPLOAD_DIR, test_id) result_dir = os.path.join(BASE_RESULT_DIR, test_id) if os.path.exists(upload_dir): shutil.rmtree(upload_dir) if os.path.exists(result_dir): shutil.rmtree(result_dir) os.makedirs(upload_dir, exist_ok=True) os.makedirs(result_dir, exist_ok=True) def convert_to_jpg(file_path, save_path): try: with Image.open(file_path) as img: if img.mode in ("RGBA", "P"): img = img.convert("RGB") img.save(save_path, "JPEG", quality=95) return True except Exception as e: print(f"Image format conversion failed: {e}") return False def save_uploaded_file(file, test_id, is_env=False): if file is None: return None, None, None upload_dir = os.path.join(BASE_UPLOAD_DIR, test_id) os.makedirs(upload_dir, exist_ok=True) original_path = file.name if hasattr(file, "name") else file file_suffix = Path(original_path).suffix.lower() if is_env: env_filename = f"env_{uuid.uuid4().hex[:4]}" save_path = os.path.join(upload_dir, f"{env_filename}{file_suffix}") shutil.copy(original_path, save_path) return file_suffix, save_path, env_filename base_filename = test_id if file_suffix in [".jpg", ".jpeg", ".png"]: save_path = os.path.join(upload_dir, f"{base_filename}.jpg") if not convert_to_jpg(original_path, save_path): shutil.copy(original_path, save_path) return ".jpg", save_path, None save_path = os.path.join(upload_dir, f"{base_filename}{file_suffix}") shutil.copy(original_path, save_path) return file_suffix, save_path, None def clear_module_results(test_id, module_num, test_env=""): try: if module_num == 1: for path in [ os.path.join(BASE_RESULT_DIR, test_id, "frames"), os.path.join(BASE_RESULT_DIR, test_id, "frames_delighting"), os.path.join(BASE_RESULT_DIR, test_id, "rego"), ]: if os.path.exists(path): shutil.rmtree(path) elif module_num == 2: path = os.path.join(BASE_RESULT_DIR, test_id, "envs", test_env) if os.path.exists(path): shutil.rmtree(path) return True except Exception as e: print(f"Failed to clear Module {module_num} results: {e}") return False def find_first_file(root_dir, glob_pattern, formats): if not os.path.exists(root_dir): return None matches = [] for fmt in formats: matches.extend(Path(root_dir).rglob(f"{glob_pattern}{fmt}")) return str(sorted(matches)[0]) if matches else None def find_first_keyword_file(root_dir, keywords): if not os.path.exists(root_dir): return None matches = [] for fmt in SUPPORTED_IMAGE_FORMATS: matches.extend(Path(root_dir).rglob(f"*{fmt}")) for path in sorted(matches): path_lower = str(path).lower() if any(keyword in path_lower for keyword in keywords): return str(path) return None def get_result_files(test_id, params): input_ext = ".jpg" if params["test_type"] == 1 else ".mp4" results = { "input_file": os.path.join(BASE_UPLOAD_DIR, test_id, f"{test_id}{input_ext}"), "input_file_type": "image" if params["test_type"] == 1 else "video", "module1": {"base_color": None, "normal": None, "roughness": None, "status": "Not Executed", "exists": False, "missing_types": []}, "module2": {"ldr_video": None, "env_dir_video": None, "status": "Not Executed", "exists": False}, "module3": {"final": None, "status": "Not Executed", "exists": False, "file_type": None}, "test_id": test_id, "test_env": params["test_env"], } module1_root_candidates = [ os.path.join(BASE_RESULT_DIR, test_id, "rego", f"{test_id}.0"), os.path.join(BASE_RESULT_DIR, test_id, "frames_delighting", "gbuffer_frames", test_id), os.path.join(BASE_RESULT_DIR, test_id, "frames_delighting", "gbuffer_frames"), os.path.join(BASE_RESULT_DIR, test_id, "frames"), ] existing_module1_roots = [path for path in module1_root_candidates if os.path.exists(path)] if params["infer_1"] == 1 or existing_module1_roots: if existing_module1_roots: for result_type, config in MODULE1_RESULT_TYPES.items(): found_file = None for root_dir in existing_module1_roots: exact_dir = os.path.join(root_dir, config["dir"]) found_file = find_first_file(exact_dir, config["glob_pattern"], SUPPORTED_IMAGE_FORMATS) if found_file: break if not found_file: keywords = { "base_color": ["base", "albedo", "color"], "normal": ["normal"], "roughness": ["rough", "roughness"], }[result_type] for root_dir in existing_module1_roots: found_file = find_first_keyword_file(root_dir, keywords) if found_file: break if found_file: results["module1"][result_type] = found_file else: results["module1"]["missing_types"].append(config["name"]) has_result = any(results["module1"][k] for k in ["base_color", "normal", "roughness"]) results["module1"]["exists"] = has_result if not results["module1"]["missing_types"]: results["module1"]["status"] = "Execution Successful (All 3 results generated)" elif has_result: missing = ", ".join(results["module1"]["missing_types"]) results["module1"]["status"] = f"Execution Successful (Partial results missing: {missing})" else: results["module1"]["status"] = "Execution Failed: No results generated" else: checked = "\n".join(f"- {path}" for path in module1_root_candidates) results["module1"]["status"] = f"Execution Failed: Module 1 output directory not found. Checked:\n{checked}" env_dir = os.path.join(BASE_RESULT_DIR, test_id, "envs", params["test_env"]) if params["infer_2"] == 1 or os.path.exists(env_dir): if os.path.exists(env_dir): for fmt in SUPPORTED_VIDEO_FORMATS: candidate = os.path.join(env_dir, f"ldr_video_fix_first_frame{fmt}") if os.path.exists(candidate): results["module2"]["ldr_video"] = candidate break for fmt in SUPPORTED_VIDEO_FORMATS: candidate = os.path.join(env_dir, f"env_dir_video_fix_first_frame{fmt}") if os.path.exists(candidate): results["module2"]["env_dir_video"] = candidate break results["module2"]["exists"] = bool(results["module2"]["ldr_video"] or results["module2"]["env_dir_video"]) if results["module2"]["ldr_video"]: results["module2"]["status"] = "Execution Successful (LDR video generated)" elif results["module2"]["env_dir_video"]: results["module2"]["status"] = "Execution Successful (Only environment direction video generated)" else: results["module2"]["status"] = "Execution Failed: No video files generated" else: results["module2"]["status"] = "Execution Failed: Directory not found" final_files = list(Path(BASE_RESULT_DIR, test_id).glob(f"{test_id}.{params['test_env']}.*")) if params["infer_3"] == 1 or final_files: if final_files: final_file = str(sorted(final_files)[0]) suffix = Path(final_file).suffix.lower() results["module3"]["final"] = final_file results["module3"]["exists"] = True if suffix in SUPPORTED_IMAGE_FORMATS: results["module3"]["file_type"] = "image" elif suffix in SUPPORTED_VIDEO_FORMATS: results["module3"]["file_type"] = "video" results["module3"]["status"] = "Execution Successful" else: results["module3"]["status"] = "Execution Failed: No image/video results generated" return results def estimate_gpu_duration(params, module_num, process_state=None): if not params: return 60 test_type = int(params.get("test_type", 1)) frames = 1 if test_type == 1 else int(params.get("frame", 25)) steps = int(params.get("num_infer_steps", 20)) if module_num == 1: return min(600, max(90, 90 + frames * 6)) if module_num == 2: return min(600, max(90, 90 + frames * 4)) if module_num == 3: return min(1200, max(120, int(120 + frames * steps * 0.7))) return 300 def path_has_files(path, patterns=None): root = Path(path) if not root.exists(): return False if not patterns: return any(item.is_file() for item in root.rglob("*")) for pattern in patterns: if any(root.glob(pattern)): return True return False def run_setup_command(command, cwd=None, env_extra=None): print(f"Runtime setup command: {' '.join(command)}") env = os.environ.copy() if env_extra: env.update(env_extra) completed = subprocess.run( command, cwd=cwd, env=env, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, ) output = completed.stdout or "" if completed.returncode != 0: raise RuntimeError(f"Command failed: {' '.join(command)}\n{output}") return output.strip() def pip_install(args, cwd=None): return run_setup_command([sys.executable, "-m", "pip", "install", *args], cwd=cwd) def snapshot_download_runtime(repo_id, local_dir, allow_patterns=None): try: from huggingface_hub import snapshot_download except Exception as e: raise RuntimeError( "Missing `huggingface_hub`. Add `huggingface_hub` to requirements.txt." ) from e kwargs = { "repo_id": repo_id, "local_dir": local_dir, "local_dir_use_symlinks": False, } if allow_patterns: kwargs["allow_patterns"] = allow_patterns token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") if token: kwargs["token"] = token snapshot_download(**kwargs) def _write_transformer_engine_shim(package_parent): """Write a small PyTorch fallback for the subset of Transformer Engine used by Cosmos.""" package_parent = Path(package_parent) shim_root = package_parent / "transformer_engine" pytorch_dir = shim_root / "pytorch" pytorch_dir.mkdir(parents=True, exist_ok=True) root_init = "from . import pytorch\n__all__ = ['pytorch']\n" pytorch_init = r'''import torch from torch import nn class RMSNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6, **kwargs): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.eps = eps def forward(self, x): dtype = x.dtype variance = x.to(torch.float32).pow(2).mean(dim=-1, keepdim=True) x = x * torch.rsqrt(variance + self.eps).to(dtype) return x * self.weight.to(dtype) LayerNorm = nn.LayerNorm Linear = nn.Linear class LayerNormLinear(nn.Module): def __init__(self, hidden_size, output_size, eps=1e-5, bias=True, return_layernorm_output=False, **kwargs): super().__init__() self.norm = nn.LayerNorm(hidden_size, eps=eps) self.linear = nn.Linear(hidden_size, output_size, bias=bias) self.return_layernorm_output = return_layernorm_output def forward(self, x): y = self.norm(x) out = self.linear(y) if self.return_layernorm_output: return out, y return out try: from .attention import DotProductAttention, apply_rotary_pos_emb except Exception: DotProductAttention = None def apply_rotary_pos_emb(tensor, *args, **kwargs): return tensor ''' attention_code = r'''import torch from torch import nn from torch.nn import functional as F def apply_rotary_pos_emb(tensor, *args, **kwargs): # Fallback for the fused Transformer Engine RoPE helper. # Keeping identity preserves shapes and lets the ZeroGPU demo run without compiling TE. return tensor class DotProductAttention(nn.Module): def __init__( self, num_attention_heads=None, kv_channels=None, *args, qkv_format="bshd", attention_dropout=0.0, **kwargs, ): super().__init__() self.qkv_format = qkv_format or "bshd" self.attention_dropout = float(attention_dropout or 0.0) def _to_bhsd(self, tensor): if tensor.ndim != 4: raise ValueError(f"TE shim expected a 4D qkv tensor, got {tuple(tensor.shape)}") fmt = self.qkv_format if fmt == "bshd": return tensor.permute(0, 2, 1, 3), "bshd" # B,H,S,D if fmt == "sbhd": return tensor.permute(1, 2, 0, 3), "sbhd" # B,H,S,D # Cosmos normally uses bshd. Keep a sane fallback for similar 4D layouts. return tensor.permute(0, 2, 1, 3), "bshd" def _restore_and_flatten(self, tensor, fmt): if fmt == "bshd": tensor = tensor.permute(0, 2, 1, 3).contiguous() # B,S,H,D elif fmt == "sbhd": tensor = tensor.permute(2, 0, 1, 3).contiguous() # S,B,H,D else: raise ValueError(f"Unsupported qkv_format for TE shim: {fmt}") return tensor.reshape(*tensor.shape[:-2], tensor.shape[-2] * tensor.shape[-1]) def forward(self, query_layer, key_layer, value_layer, *args, **kwargs): q, fmt = self._to_bhsd(query_layer) k, _ = self._to_bhsd(key_layer) v, _ = self._to_bhsd(value_layer) out = F.scaled_dot_product_attention( q, k, v, dropout_p=self.attention_dropout if self.training else 0.0, is_causal=False, ) return self._restore_and_flatten(out, fmt) ''' (shim_root / "__init__.py").write_text(root_init, encoding="utf-8") (pytorch_dir / "__init__.py").write_text(pytorch_init, encoding="utf-8") (pytorch_dir / "attention.py").write_text(attention_code, encoding="utf-8") return shim_root def ensure_transformer_engine_shim(cosmos_dir, logs): if not USE_TRANSFORMER_ENGINE_SHIM: logs.append("Transformer Engine shim disabled. Set RELIT_USE_TRANSFORMER_ENGINE_SHIM=1 to enable it.") return repo_root = Path.cwd().resolve() cosmos_dir = Path(cosmos_dir).resolve() written = [] # Root copy wins because run_bash_script puts the app root first in PYTHONPATH. written.append(_write_transformer_engine_shim(repo_root)) # Cosmos-local copy is a fallback when a Cosmos script mutates sys.path/cwd. if cosmos_dir.exists(): written.append(_write_transformer_engine_shim(cosmos_dir)) pythonpath = f"{repo_root}:{cosmos_dir}:{os.environ.get('PYTHONPATH', '')}".rstrip(":") try: run_setup_command( [ sys.executable, "-c", ( "import transformer_engine as te; " "from transformer_engine.pytorch.attention import DotProductAttention; " "print('TE shim import ok', te.__file__, DotProductAttention)" ), ], env_extra={"PYTHONPATH": pythonpath}, ) logs.append("Transformer Engine compatibility shim ready and import-tested.") except Exception as e: logs.append(f"Transformer Engine shim was written but import test failed: {e}") if STRICT_COSMOS_DEPS: raise logs.append("Transformer Engine shim locations: " + ", ".join(str(path) for path in written)) def ensure_cosmos_dependencies(cosmos_dir, logs): cosmos_dir = Path(cosmos_dir) if not cosmos_dir.exists(): return if not INSTALL_COSMOS_DEPS: logs.append("Cosmos dependency installation disabled.") ensure_transformer_engine_shim(cosmos_dir, logs) return requirements_marker = cosmos_dir / ".relit_requirements_installed" requirements_path = cosmos_dir / "requirements.txt" if requirements_path.exists() and not requirements_marker.exists(): logs.append("Installing Cosmos Python requirements...") try: pip_install(["-r", str(requirements_path)]) requirements_marker.write_text("ok\n") logs.append("Cosmos Python requirements installed.") except Exception as e: logs.append(f"Cosmos requirements install failed: {e}") logs.append("Installing minimal fallback dependency: omegaconf...") pip_install(["omegaconf"]) logs.append("Minimal fallback dependency installed.") elif requirements_marker.exists(): logs.append("Cosmos Python requirements already installed.") else: logs.append("No Cosmos requirements.txt found; installing omegaconf fallback...") pip_install(["omegaconf"]) logs.append("OmegaConf installed.") ensure_transformer_engine_shim(cosmos_dir, logs) te_marker = cosmos_dir / ".relit_transformer_engine_installed" if INSTALL_TRANSFORMER_ENGINE and not te_marker.exists(): logs.append("Installing transformer-engine[pytorch]==1.12.0...") try: pip_install(["--no-build-isolation", "transformer-engine[pytorch]==1.12.0"]) te_marker.write_text("ok\n") logs.append("Transformer Engine installed.") except Exception as e: logs.append(f"Transformer Engine install failed: {e}") logs.append("Continuing setup. Set RELIT_STRICT_COSMOS_DEPS=1 to make this fatal.") if STRICT_COSMOS_DEPS: raise elif INSTALL_TRANSFORMER_ENGINE: logs.append("Transformer Engine already installed.") else: logs.append("Transformer Engine install disabled. Using shim unless RELIT_USE_TRANSFORMER_ENGINE_SHIM=0.") nvdiffrast_marker = cosmos_dir / ".relit_nvdiffrast_installed" if INSTALL_NVDIFFRAST and not nvdiffrast_marker.exists(): logs.append("Installing nvdiffrast...") try: pip_install(["--no-build-isolation", "git+https://github.com/NVlabs/nvdiffrast.git"]) nvdiffrast_marker.write_text("ok\n") logs.append("nvdiffrast installed.") except Exception as e: logs.append(f"nvdiffrast install failed: {e}") logs.append("Continuing setup. Set RELIT_STRICT_COSMOS_DEPS=1 to make this fatal.") if STRICT_COSMOS_DEPS: raise elif INSTALL_NVDIFFRAST: logs.append("nvdiffrast already installed.") else: logs.append("nvdiffrast install disabled. Set RELIT_INSTALL_NVDIFFRAST=1 if the renderer asks for it.") def ensure_cosmos_checkpoints(cosmos_dir, logs): if not DOWNLOAD_COSMOS_CHECKPOINTS: logs.append("Cosmos checkpoint download disabled. Set RELIT_DOWNLOAD_COSMOS_CHECKPOINTS=1 to enable it.") return cosmos_dir = Path(cosmos_dir).resolve() checkpoints_dir = cosmos_dir / "checkpoints" inverse_dir = checkpoints_dir / "Diffusion_Renderer_Inverse_Cosmos_7B" forward_dir = checkpoints_dir / "Diffusion_Renderer_Forward_Cosmos_7B" if inverse_dir.exists() and forward_dir.exists(): logs.append("Cosmos DiffusionRenderer checkpoints already present.") return # Recover checkpoints that were downloaded to a nested path by an older runtime setup. nested_inverse_dirs = [ path for path in cosmos_dir.rglob("Diffusion_Renderer_Inverse_Cosmos_7B") if path.is_dir() and path != inverse_dir ] nested_forward_dirs = [ path for path in cosmos_dir.rglob("Diffusion_Renderer_Forward_Cosmos_7B") if path.is_dir() and path != forward_dir ] if nested_inverse_dirs or nested_forward_dirs: checkpoints_dir.mkdir(parents=True, exist_ok=True) for found, expected in [ (nested_inverse_dirs[0] if nested_inverse_dirs else None, inverse_dir), (nested_forward_dirs[0] if nested_forward_dirs else None, forward_dir), ]: if found and not expected.exists(): try: expected.symlink_to(found, target_is_directory=True) logs.append(f"Linked misplaced Cosmos checkpoint: {expected} -> {found}") except Exception: shutil.copytree(found, expected) logs.append(f"Copied misplaced Cosmos checkpoint: {found} -> {expected}") if inverse_dir.exists() and forward_dir.exists(): logs.append("Cosmos DiffusionRenderer checkpoints recovered.") return download_script = cosmos_dir / "scripts" / "download_diffusion_renderer_checkpoints.py" if not download_script.exists(): logs.append("Cosmos checkpoint download script not found; skipping checkpoint download.") return checkpoints_dir.mkdir(parents=True, exist_ok=True) logs.append( "Downloading Cosmos DiffusionRenderer checkpoints to " f"`{checkpoints_dir}`. This is large and may take a while..." ) cosmos_pythonpath = str(cosmos_dir.resolve()) existing_pythonpath = os.environ.get("PYTHONPATH") if existing_pythonpath: cosmos_pythonpath = f"{cosmos_pythonpath}:{existing_pythonpath}" run_setup_command( [ sys.executable, str(download_script.resolve()), "--checkpoint_dir", str(checkpoints_dir.resolve()), ], cwd=str(cosmos_dir), env_extra={"PYTHONPATH": cosmos_pythonpath}, ) logs.append("Cosmos DiffusionRenderer checkpoints ready.") def ensure_runtime_assets(): global RUNTIME_SETUP_DONE, RUNTIME_SETUP_ATTEMPTED, RUNTIME_SETUP_LOG with RUNTIME_SETUP_LOCK: RUNTIME_SETUP_ATTEMPTED = True if RUNTIME_SETUP_DONE: return True, RUNTIME_SETUP_LOG logs = [] try: os.makedirs("checkpoints", exist_ok=True) os.makedirs("models/Wan-AI", exist_ok=True) os.makedirs("third_party", exist_ok=True) checkpoint_patterns = [ "model_frame25_480_832.ckpt", "model_frame57_480_832.ckpt", "model_frame1_1024_1472.ckpt", ] checkpoints_ready = path_has_files("checkpoints", checkpoint_patterns) if DOWNLOAD_RELIT_CHECKPOINTS and not checkpoints_ready: logs.append(f"Downloading Relit-LiVE checkpoints from {RELIT_CHECKPOINT_REPO}...") snapshot_download_runtime( repo_id=RELIT_CHECKPOINT_REPO, local_dir=".", allow_patterns=["checkpoints/*.ckpt"], ) logs.append("Relit-LiVE checkpoints ready.") elif checkpoints_ready: logs.append("Relit-LiVE checkpoints already present.") else: logs.append("Relit-LiVE checkpoint download disabled.") wan_dir = Path("models/Wan-AI/Wan2.1-T2V-1.3B") if DOWNLOAD_WAN and not path_has_files(wan_dir): logs.append(f"Downloading Wan2.1 base model from {WAN_REPO}...") snapshot_download_runtime( repo_id=WAN_REPO, local_dir=str(wan_dir), ) logs.append("Wan2.1 base model ready.") elif path_has_files(wan_dir): logs.append("Wan2.1 base model already present.") else: logs.append("Wan2.1 download disabled.") cosmos_dir = Path("third_party/cosmos-transfer1-diffusion-renderer") if CLONE_COSMOS and not cosmos_dir.exists(): logs.append(f"Cloning Cosmos Transfer renderer from {COSMOS_REPO}...") run_setup_command(["git", "clone", "--depth", "1", COSMOS_REPO, str(cosmos_dir)]) logs.append("Cosmos Transfer renderer cloned.") elif cosmos_dir.exists(): logs.append("Cosmos Transfer renderer already present.") else: logs.append("Cosmos renderer clone disabled.") if cosmos_dir.exists(): ensure_cosmos_dependencies(cosmos_dir, logs) ensure_cosmos_checkpoints(cosmos_dir, logs) RUNTIME_SETUP_DONE = True RUNTIME_SETUP_LOG = "\n".join(logs) print(RUNTIME_SETUP_LOG) return True, RUNTIME_SETUP_LOG except Exception as e: logs.append(f"Runtime setup failed: {e}") RUNTIME_SETUP_LOG = "\n".join(logs) print(RUNTIME_SETUP_LOG) return False, RUNTIME_SETUP_LOG def get_runtime_setup_status_for_execution(): if STARTUP_RUNTIME_SETUP and RUNTIME_SETUP_ATTEMPTED and not RUNTIME_SETUP_DONE: return False, RUNTIME_SETUP_LOG return ensure_runtime_assets() def get_missing_module_prerequisites(module_num): missing = [] if module_num == 1: renderer_dir = Path("third_party/cosmos-transfer1-diffusion-renderer") if not renderer_dir.exists(): missing.append( "Missing `third_party/cosmos-transfer1-diffusion-renderer`. " "Module 1 needs the Cosmos Transfer renderer to generate gbuffer frames." ) inverse_checkpoint_dir = renderer_dir / "checkpoints" / "Diffusion_Renderer_Inverse_Cosmos_7B" has_inverse_checkpoint = inverse_checkpoint_dir.exists() or any( path.is_dir() for path in renderer_dir.rglob("Diffusion_Renderer_Inverse_Cosmos_7B") ) if renderer_dir.exists() and not has_inverse_checkpoint: missing.append( "Missing Cosmos inverse renderer checkpoints under " "`third_party/cosmos-transfer1-diffusion-renderer/checkpoints/Diffusion_Renderer_Inverse_Cosmos_7B`. " "Set `RELIT_DOWNLOAD_COSMOS_CHECKPOINTS=1` or provide the checkpoints manually." ) if module_num in [2, 3]: wan_dir = Path("models/Wan-AI/Wan2.1-T2V-1.3B") if not wan_dir.exists(): missing.append("Missing `models/Wan-AI/Wan2.1-T2V-1.3B`.") if module_num == 3: checkpoints = [ Path("checkpoints/model_frame25_480_832.ckpt"), Path("checkpoints/model_frame57_480_832.ckpt"), Path("checkpoints/model_frame1_1024_1472.ckpt"), ] if not any(path.exists() for path in checkpoints): missing.append("Missing Relit-LiVE checkpoints under `checkpoints/`.") return missing @spaces.GPU(duration=estimate_gpu_duration) def run_bash_script(params, module_num, process_state): patch_bash_script_for_hf_spaces() test_id = params["test_id"] env = os.environ.copy() repo_path = os.getcwd() cosmos_path = str(Path("third_party/cosmos-transfer1-diffusion-renderer").resolve()) # Last-mile guard: make the Transformer Engine shim visible even if the # startup setup was skipped or the Space was restarted from a stale image. if USE_TRANSFORMER_ENGINE_SHIM: try: _write_transformer_engine_shim(Path(repo_path)) if Path(cosmos_path).exists(): _write_transformer_engine_shim(Path(cosmos_path)) print("Transformer Engine shim ensured before module subprocess.") except Exception as shim_error: print(f"Warning: failed to ensure Transformer Engine shim: {shim_error}") pythonpath_parts = [repo_path] if Path(cosmos_path).exists(): pythonpath_parts.append(cosmos_path) if env.get("PYTHONPATH"): pythonpath_parts.append(env["PYTHONPATH"]) env.update({ "TEST_ID": test_id, "TEST_TYPE": str(params["test_type"]), "TEST_ENV": params["test_env"], "USE_OFFICE_ENV": str(params["use_office_env"]), "FRAME": str(params["frame"]), "FRAME_RATE": str(params["frame_rate"]), "INFER_1": str(params["infer_1"]), "INFER_2": str(params["infer_2"]), "INFER_3": str(params["infer_3"]), "ENV_STRENGTH": str(params["env_strength"]), "NUM_INFER_STEPS": str(params["num_infer_steps"]), "WORW": str(params["worw"]), "LIGHT_TYPE": str(params["light_type"]), "REPO_PATH": repo_path, "PYTHONPATH": ":".join(pythonpath_parts), }) try: process = subprocess.Popen( ["bash", BASH_SCRIPT_PATH], env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, preexec_fn=os.setsid, ) with RUNNING_PROCESSES_LOCK: RUNNING_PROCESSES[module_num] = process process_state[module_num] = process.pid print(f"Module {module_num} started with PID: {process.pid}") stdout, stderr = process.communicate() combined_logs = "\n".join(part for part in [stdout, stderr] if part) with RUNNING_PROCESSES_LOCK: RUNNING_PROCESSES.pop(module_num, None) process_state.pop(module_num, None) if process.returncode == 0: print(f"Module {module_num} execution succeeded (PID: {process.pid})") print(f"Script stdout: {stdout}") if stderr: print(f"Script stderr: {stderr}") return True, test_id, "", combined_logs error_msg = f"Module {module_num} execution failed (PID: {process.pid}): {stderr}\n{stdout}" print(error_msg) return False, test_id, error_msg, "" except Exception as e: error_msg = f"Module {module_num} execution error: {e}" print(error_msg) with RUNNING_PROCESSES_LOCK: RUNNING_PROCESSES.pop(module_num, None) process_state.pop(module_num, None) return False, test_id, error_msg, "" def stop_module_execution(module_num, process_state, current_status): with RUNNING_PROCESSES_LOCK: process = RUNNING_PROCESSES.get(module_num) pid = process.pid if process else process_state.get(module_num) if not pid: return process_state, current_status + "\nNo running process found for this module" try: os.killpg(os.getpgid(pid), signal.SIGTERM) with RUNNING_PROCESSES_LOCK: RUNNING_PROCESSES.pop(module_num, None) process_state.pop(module_num, None) return process_state, current_status + f"\nModule {module_num} execution stopped successfully (PID: {pid})" except ProcessLookupError: with RUNNING_PROCESSES_LOCK: RUNNING_PROCESSES.pop(module_num, None) process_state.pop(module_num, None) return process_state, current_status + "\nProcess already completed" except Exception as e: return process_state, current_status + f"\nFailed to stop module: {e}" def get_demo_images(): demo_images = [] if os.path.exists(DEMO_IMAGES_DIR): for ext in SUPPORTED_IMAGE_FORMATS: demo_images.extend(list(Path(DEMO_IMAGES_DIR).glob(f"*{ext}"))) return [str(path) for path in sorted(demo_images)] def update_carousel(index, direction, total_images): if total_images <= 0: return -1 if direction == "next": return (index + 1) % total_images if direction == "prev": return (index - 1) % total_images return index def init_input_file(input_file): if input_file is None: return gr.update(value=None), gr.update(value=None), "Error: Please upload an input image/video file!", gr.update(value={}) test_id = generate_test_id("demo") clear_test_dir(test_id) input_suffix, input_path, _ = save_uploaded_file(input_file, test_id, is_env=False) test_type = 1 if input_suffix in [".jpg", ".jpeg", ".png"] else 0 input_file_type = "image" if test_type == 1 else "video" base_params = { "test_id": test_id, "test_type": test_type, "test_env": DEFAULT_ENV or "", "use_office_env": 1, "frame": 25 if test_type == 0 else 1, "frame_rate": 24, "infer_1": 0, "infer_2": 0, "infer_3": 0, "env_strength": 3.0, "num_infer_steps": 20, "worw": 0.0, "light_type": 0, } return ( input_path if input_file_type == "image" else None, input_path if input_file_type == "video" else None, f"Input file initialized successfully!\nFile type: {input_file_type}\nEnvironment: {DEFAULT_ENV or 'not configured'}", base_params, ) def select_demo_image(image_path): if not image_path or not os.path.exists(image_path): return gr.update(value=None, visible=False), gr.update(value=None, visible=False), "Error: Demo image not found!", gr.update(value={}) class MockFile: def __init__(self, path): self.name = path image_preview, video_preview, status, params = init_input_file(MockFile(image_path)) return ( gr.update(value=image_preview, visible=image_preview is not None, height=INPUT_PREVIEW_HEIGHT), gr.update(value=video_preview, visible=video_preview is not None, height=INPUT_PREVIEW_HEIGHT), f"Selected demo image: {os.path.basename(image_path)}\n{status}", params, ) def update_env_config(base_params, use_builtin_env, builtin_env_choice, env_file): if not base_params: return base_params, "Error: Please initialize input file first!" test_id = base_params["test_id"] if not use_builtin_env and env_file is not None: _, _, test_env = save_uploaded_file(env_file, test_id, is_env=True) use_office_env = 0 env_type = "Custom Environment" else: test_env = builtin_env_choice use_office_env = 1 env_type = f"Built-in Environment ({builtin_env_choice})" if not test_env: return base_params, "Error: Please select or upload an environment first!" base_params.update({"test_env": test_env, "use_office_env": use_office_env}) return base_params, f"Environment configuration updated successfully!\nType: {env_type}" def update_advanced_params(base_params, frame, frame_rate, env_strength, num_infer_steps, worw, light_type): if not base_params: return base_params, "Error: Please initialize input file first!" if light_type not in [0, 1, 2]: light_type = 0 base_params.update({ "frame": frame if base_params["test_type"] == 0 else 1, "frame_rate": frame_rate, "env_strength": env_strength, "num_infer_steps": num_infer_steps, "worw": worw, "light_type": light_type, }) desc = { 0: "Original Scene + Static Light", 1: "Original Scene + Dynamic Light", 2: "Fixed First Frame + Dynamic Light", } return base_params, f"Advanced parameters updated successfully!\nFrame rate: {frame_rate} | Light Type: {desc[light_type]} | Inference Steps: {num_infer_steps} | Env Strength: {env_strength}" def run_single_module(module_num, params, process_state, re_run=True): if not params: if module_num == 1: return None, None, None, process_state, "Error: Please initialize input file first!" if module_num == 2: return None, None, process_state, "Error: Please initialize input file first!" return gr.update(value=None), gr.update(value=None), process_state, "Error: Please initialize input file first!" if module_num in process_state: pid = process_state[module_num] msg = f"Module {module_num} is already running (PID: {pid}). Please stop it first." if module_num == 1: return None, None, None, process_state, msg if module_num == 2: return None, None, process_state, msg return gr.update(value=None), gr.update(value=None), process_state, msg if not re_run: results = get_result_files(params["test_id"], params) if module_num == 1: return results["module1"]["base_color"], results["module1"]["normal"], results["module1"]["roughness"], process_state, f"Show existing results: {results['module1']['status']}" if module_num == 2: return results["module2"]["ldr_video"], results["module2"]["env_dir_video"], process_state, f"Show existing results: {results['module2']['status']}" image_result = results["module3"]["final"] if results["module3"]["file_type"] == "image" else None video_result = results["module3"]["final"] if results["module3"]["file_type"] == "video" else None return gr.update(value=image_result), gr.update(value=video_result), process_state, f"Show existing results: {results['module3']['status']}" setup_ok, setup_log = get_runtime_setup_status_for_execution() if not setup_ok: msg = "Runtime setup failed:\n" + setup_log if module_num == 1: return None, None, None, process_state, msg if module_num == 2: return None, None, process_state, msg return gr.update(value=None), gr.update(value=None), process_state, msg missing_prerequisites = get_missing_module_prerequisites(module_num) if missing_prerequisites: msg = "Missing prerequisites after runtime setup:\n" msg += "\n".join(f"- {item}" for item in missing_prerequisites) if setup_log: msg += "\n\nRuntime setup log:\n" + setup_log if module_num == 1: return None, None, None, process_state, msg if module_num == 2: return None, None, process_state, msg return gr.update(value=None), gr.update(value=None), process_state, msg if params["test_env"] == "" and module_num in [2, 3]: if module_num == 2: return None, None, process_state, "Error: Please configure environment first!" return gr.update(value=None), gr.update(value=None), process_state, "Error: Please configure environment first!" clear_success = clear_module_results(params["test_id"], module_num, params["test_env"]) status_msg = f"Cleared old results for Module {module_num}, starting execution..." if clear_success else f"Could not clear old results for Module {module_num}, attempting execution..." if setup_log: status_msg = f"Runtime setup ready:\n{setup_log}\n\n{status_msg}" params["infer_1"] = 1 if module_num == 1 else 0 params["infer_2"] = 1 if module_num == 2 else 0 params["infer_3"] = 1 if module_num == 3 else 0 success, test_id, error_msg, stdout = run_bash_script(params, module_num, process_state) results = get_result_files(test_id, params) stdout_tail = "" if stdout: stdout_lines = stdout.strip().splitlines() stdout_tail = "\nScript stdout tail:\n" + "\n".join(stdout_lines[-12:]) if module_num == 1: final_status = f"{status_msg}\n{results['module1']['status']}" if error_msg: final_status += f"\n{error_msg}" elif success and not results["module1"]["exists"] and stdout_tail: final_status += stdout_tail return results["module1"]["base_color"], results["module1"]["normal"], results["module1"]["roughness"], process_state, final_status if module_num == 2: final_status = f"{status_msg}\n{results['module2']['status']}" if error_msg: final_status += f"\n{error_msg}" elif success and not results["module2"]["exists"] and stdout_tail: final_status += stdout_tail return results["module2"]["ldr_video"], results["module2"]["env_dir_video"], process_state, final_status final_status = f"{status_msg}\n{results['module3']['status']}" if error_msg: final_status += f"\n{error_msg}" elif success and not results["module3"]["exists"] and stdout_tail: final_status += stdout_tail image_result = results["module3"]["final"] if results["module3"]["file_type"] == "image" else None video_result = results["module3"]["final"] if results["module3"]["file_type"] == "video" else None return gr.update(value=image_result), gr.update(value=video_result), process_state, final_status def one_click_run_all(params, process_state, stop_flag): stop_flag = False if not params: msg = "One-click run failed: Input file not initialized" return None, None, None, None, None, gr.update(), gr.update(), process_state, msg, msg, msg, stop_flag if params["test_env"] == "": msg = "One-click run failed: Environment not configured" return None, None, None, None, None, gr.update(), gr.update(), process_state, msg, msg, msg, stop_flag m1_base, m1_normal, m1_rough, process_state, m1_status = run_single_module(1, params, process_state, re_run=True) if "Failed" in m1_status or "Error" in m1_status: return m1_base, m1_normal, m1_rough, None, None, gr.update(), gr.update(), process_state, m1_status + "\nModule1 failed, one-click run terminated", "Module2 not executed", "Module3 not executed", stop_flag m2_ldr, m2_env, process_state, m2_status = run_single_module(2, params, process_state, re_run=True) if "Failed" in m2_status or "Error" in m2_status: return m1_base, m1_normal, m1_rough, m2_ldr, m2_env, gr.update(), gr.update(), process_state, m1_status, m2_status + "\nModule2 failed, one-click run terminated", "Module3 not executed", stop_flag m3_img, m3_video, process_state, m3_status = run_single_module(3, params, process_state, re_run=True) m3_status += "\nOne-click run completed successfully!" if "Failed" not in m3_status and "Error" not in m3_status else "\nOne-click run completed with Module3 failure" return m1_base, m1_normal, m1_rough, m2_ldr, m2_env, m3_img, m3_video, process_state, m1_status, m2_status, m3_status, stop_flag def stop_one_click_run(process_state, stop_flag, m1_status, m2_status, m3_status): stop_flag = True for module_num in list(process_state.keys()): process_state, _ = stop_module_execution(module_num, process_state, "") return process_state, stop_flag, m1_status + "\nOne-click run manually stopped", m2_status + "\nOne-click run manually stopped", m3_status + "\nOne-click run manually stopped" def init_and_show_preview(input_file): image_preview, video_preview, status, params = init_input_file(input_file) return ( gr.update(value=image_preview, visible=image_preview is not None, height=INPUT_PREVIEW_HEIGHT), gr.update(value=video_preview, visible=video_preview is not None, height=INPUT_PREVIEW_HEIGHT), status, params, ) css = """ .carousel-container { width: 100%; max-width: 600px; margin: 20px auto; overflow: hidden; border-radius: 12px; } .select-image-btn { margin-top: 15px; width: 100%; max-width: 200px; } """ with gr.Blocks(title="Relit-LiVE: Relighting Model Interactive Inference Tool") as demo: gr.Markdown(""" # Relit-LiVE: Relighting Model Interactive Inference Tool > Left Panel (Input + Parameter Configuration) | Right Panel (Module 1 + Module 2 + Module 3) """) base_params = gr.State(value={}) process_state = gr.State(value={}) one_click_stop_flag = gr.State(value=False) demo_image_paths = get_demo_images() total_images = len(demo_image_paths) current_index = gr.State(value=0 if total_images > 0 else -1) with gr.Row(): with gr.Column(scale=1, min_width=400): gr.Markdown("## Input & Parameter Configuration") input_file = gr.File(label="Upload Input File (Supports jpg/png/mp4)", file_types=[".jpg", ".jpeg", ".png", ".mp4"]) init_input_btn = gr.Button("Initialize Input File", variant="primary") with gr.Row(): input_image_preview = gr.Image(label="Image Preview", height=INPUT_PREVIEW_HEIGHT, visible=False) input_video_preview = gr.Video(label="Video Preview", height=INPUT_PREVIEW_HEIGHT, visible=False) input_status = gr.Textbox(label="Initialization Status", lines=2) gr.Markdown("### Demo Images") if demo_image_paths: carousel_image = gr.Image(value=demo_image_paths[0], label="Demo Image", height=DEMO_IMAGE_HEIGHT, interactive=False) with gr.Row(): prev_btn = gr.Button("Previous") next_btn = gr.Button("Next") select_current_btn = gr.Button("Select This Image", variant="primary", elem_classes=["select-image-btn"]) def update_carousel_ui(index): if index < 0 or index >= len(demo_image_paths): return None return demo_image_paths[index] prev_btn.click(lambda idx: update_carousel(idx, "prev", total_images), inputs=[current_index], outputs=[current_index]).then(update_carousel_ui, inputs=[current_index], outputs=[carousel_image]) next_btn.click(lambda idx: update_carousel(idx, "next", total_images), inputs=[current_index], outputs=[current_index]).then(update_carousel_ui, inputs=[current_index], outputs=[carousel_image]) select_current_btn.click(lambda idx: select_demo_image(demo_image_paths[idx] if 0 <= idx < len(demo_image_paths) else None), inputs=[current_index], outputs=[input_image_preview, input_video_preview, input_status, base_params]) else: gr.Markdown("*No demo images found in the specified directory*") gr.Markdown("---\n## Environment Configuration") use_builtin_env = gr.Checkbox(label="Use Built-in Environment (Uncheck to upload custom)", value=True) with gr.Row(): builtin_env_choice = gr.Dropdown(label="Built-in Environment Selection", choices=BUILTIN_ENV_OPTIONS, value=DEFAULT_ENV) env_file = gr.File(label="Custom Environment (hdr/jpg/png)", file_types=[".hdr", ".jpg", ".jpeg", ".png"], visible=False) update_env_btn = gr.Button("Update Environment", variant="primary") env_status = gr.Textbox(label="Environment Config Status", lines=1) gr.Markdown("---\n## Advanced Parameters") frame = gr.Slider(label="Video Frames (Video only, 1-57, 4n+1)", minimum=1, maximum=57, step=4, value=25) frame_rate = gr.Slider(label="Sample Rate of Video Frames (Video only, 10-24)", minimum=10, maximum=24, step=1, value=24) env_strength = gr.Slider(label="Environment Strength (0-5)", minimum=0, maximum=5, step=0.1, value=3.0) num_infer_steps = gr.Slider(label="Inference Steps (1-50)", minimum=1, maximum=50, step=1, value=20) worw = gr.Slider(label="Reference Image Weight (0-5, smaller = more influence. Increase it when the light fails.)", minimum=0, maximum=5, step=0.1, value=0.0) light_type = gr.Radio(label="Light Type", choices=[0, 1, 2], value=0) update_advanced_btn = gr.Button("Update Parameters", variant="primary") advanced_status = gr.Textbox(label="Parameter Update Status", lines=2) gr.Markdown("---\n## One-Click Run All Modules") with gr.Row(): one_click_run_btn = gr.Button("Run Module1 -> Module2 -> Module3", variant="primary", size="lg") one_click_stop_btn = gr.Button("Stop All Running Modules", variant="stop", size="lg") one_click_status = gr.Textbox(label="One-Click Run Status", lines=3) with gr.Column(scale=2, min_width=600): gr.Markdown("## Module 1: Inverse Rendering") with gr.Row(): run_module1_btn = gr.Button("Start Execution", variant="primary") stop_module1_btn = gr.Button("Stop Execution", variant="stop") show_module1_btn = gr.Button("Show Results", variant="secondary") module1_status = gr.Textbox(label="Execution Status", lines=2) with gr.Row(equal_height=True): module1_base_color = gr.Image(label="Base Color", height=MODULE1_VIS_HEIGHT, scale=1) module1_normal = gr.Image(label="Normal Map", height=MODULE1_VIS_HEIGHT, scale=1) module1_roughness = gr.Image(label="Roughness Map", height=MODULE1_VIS_HEIGHT, scale=1) gr.Markdown("---\n## Module 2: Environment Processing") with gr.Row(): run_module2_btn = gr.Button("Start Execution", variant="primary") stop_module2_btn = gr.Button("Stop Execution", variant="stop") show_module2_btn = gr.Button("Show Results", variant="secondary") module2_status = gr.Textbox(label="Execution Status", lines=2) with gr.Row(equal_height=True): module2_ldr_video = gr.Video(label="LDR Video (Core Result)", height=MODULE2_VIS_HEIGHT, scale=1) module2_env_video = gr.Video(label="Environment Direction Video", height=MODULE2_VIS_HEIGHT, scale=1) gr.Markdown("---\n## Module 3: Relighting") with gr.Row(): run_module3_btn = gr.Button("Start Execution", variant="primary") stop_module3_btn = gr.Button("Stop Execution", variant="stop") show_module3_btn = gr.Button("Show Results", variant="secondary") module3_status = gr.Textbox(label="Execution Status", lines=2) with gr.Row(): module3_image_result = gr.Image(label="Relighting Result (Image)", height=MODULE3_VIS_HEIGHT, visible=False, scale=1) module3_video_result = gr.Video(label="Relighting Result (Video)", height=MODULE3_VIS_HEIGHT, visible=False, scale=1) init_input_btn.click(init_and_show_preview, inputs=[input_file], outputs=[input_image_preview, input_video_preview, input_status, base_params]) use_builtin_env.change(lambda x: gr.update(visible=not x), inputs=[use_builtin_env], outputs=[env_file]) update_env_btn.click(update_env_config, inputs=[base_params, use_builtin_env, builtin_env_choice, env_file], outputs=[base_params, env_status]) update_advanced_btn.click(update_advanced_params, inputs=[base_params, frame, frame_rate, env_strength, num_infer_steps, worw, light_type], outputs=[base_params, advanced_status]) run_module1_btn.click(lambda params, ps: run_single_module(1, params, ps, True), inputs=[base_params, process_state], outputs=[module1_base_color, module1_normal, module1_roughness, process_state, module1_status]) stop_module1_btn.click(lambda ps, cs: stop_module_execution(1, ps, cs), inputs=[process_state, module1_status], outputs=[process_state, module1_status], queue=False) show_module1_btn.click(lambda params, ps: run_single_module(1, params, ps, False), inputs=[base_params, process_state], outputs=[module1_base_color, module1_normal, module1_roughness, process_state, module1_status]) run_module2_btn.click(lambda params, ps: run_single_module(2, params, ps, True), inputs=[base_params, process_state], outputs=[module2_ldr_video, module2_env_video, process_state, module2_status]) stop_module2_btn.click(lambda ps, cs: stop_module_execution(2, ps, cs), inputs=[process_state, module2_status], outputs=[process_state, module2_status], queue=False) show_module2_btn.click(lambda params, ps: run_single_module(2, params, ps, False), inputs=[base_params, process_state], outputs=[module2_ldr_video, module2_env_video, process_state, module2_status]) def update_module3_visibility(image_val, video_val, status): return ( gr.update(value=image_val, visible=image_val is not None, height=MODULE3_VIS_HEIGHT), gr.update(value=video_val, visible=video_val is not None, height=MODULE3_VIS_HEIGHT), status, ) run_module3_btn.click(lambda params, ps: run_single_module(3, params, ps, True), inputs=[base_params, process_state], outputs=[module3_image_result, module3_video_result, process_state, module3_status]).then(update_module3_visibility, inputs=[module3_image_result, module3_video_result, module3_status], outputs=[module3_image_result, module3_video_result, module3_status]) stop_module3_btn.click(lambda ps, cs: stop_module_execution(3, ps, cs), inputs=[process_state, module3_status], outputs=[process_state, module3_status], queue=False) show_module3_btn.click(lambda params, ps: run_single_module(3, params, ps, False), inputs=[base_params, process_state], outputs=[module3_image_result, module3_video_result, process_state, module3_status]).then(update_module3_visibility, inputs=[module3_image_result, module3_video_result, module3_status], outputs=[module3_image_result, module3_video_result, module3_status]) one_click_run_btn.click( one_click_run_all, inputs=[base_params, process_state, one_click_stop_flag], outputs=[module1_base_color, module1_normal, module1_roughness, module2_ldr_video, module2_env_video, module3_image_result, module3_video_result, process_state, module1_status, module2_status, module3_status, one_click_stop_flag], ).then( lambda image_val, video_val: (gr.update(visible=image_val is not None, height=MODULE3_VIS_HEIGHT), gr.update(visible=video_val is not None, height=MODULE3_VIS_HEIGHT)), inputs=[module3_image_result, module3_video_result], outputs=[module3_image_result, module3_video_result], ) one_click_stop_btn.click(stop_one_click_run, inputs=[process_state, one_click_stop_flag, module1_status, module2_status, module3_status], outputs=[process_state, one_click_stop_flag, module1_status, module2_status, module3_status], queue=False) if __name__ == "__main__": os.makedirs(BASE_UPLOAD_DIR, exist_ok=True) os.makedirs(BASE_RESULT_DIR, exist_ok=True) print("=" * 60) print("Relit-LiVE: Relighting Model Interactive Inference Tool") print(f"App build: {APP_BUILD_ID}") print("ZeroGPU allocation wraps the module subprocess execution") print("=" * 60) if STARTUP_RUNTIME_SETUP: print("Preparing runtime assets before launching the app...") setup_ok, setup_log = ensure_runtime_assets() print(setup_log) if not setup_ok: print("Runtime setup failed. The app will still launch and show the setup error in module status.") else: print("Startup runtime setup disabled with RELIT_STARTUP_SETUP=0") demo.queue(default_concurrency_limit=1, max_size=10) demo.launch( css=css, server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)), share=False, show_error=True, )