import gradio as gr import os import subprocess import uuid import shutil import signal 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 Directory Configuration DEMO_IMAGES_DIR = "./datasets/gradio_data/assets/images_demo" BASH_SCRIPT_PATH = "./tools/full_inference_modules_gradio.sh" # Supported Image and Video Formats SUPPORTED_IMAGE_FORMATS = [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".hdr"] SUPPORTED_VIDEO_FORMATS = [".mp4", ".mov", ".avi", ".mkv"] SUPPORTED_HDR_FORMATS = [".hdr", ".exr", ".jpg", ".png"] # Built-in Environment Options 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}")]) # Module 1 Result Types (Fixed Directory Structure) MODULE1_RESULT_TYPES = { "base_color": { "name": "Base Color", "dir": "Base Color", # Corresponding directory name "glob_pattern": "frame_*" # File name matching pattern }, "normal": { "name": "Normal Map", "dir": "normal", "glob_pattern": "frame_*" }, "roughness": { "name": "Roughness Map", "dir": "Roughness", "glob_pattern": "frame_*" } } # -------------------------- Visualization Window Configuration (Adaptive Width) -------------------------- # Base Height Configuration (Ensure Vertical Proportion) INPUT_PREVIEW_HEIGHT = 240 # Input preview height MODULE1_VIS_HEIGHT = 180 # Module 1 visualization height (3 results split equally horizontally) MODULE2_VIS_HEIGHT = 200 # Module 2 visualization height (2 videos split equally horizontally) MODULE3_VIS_HEIGHT = 400 # Module 3 visualization height (1 result full width) - Adjusted to 400px DEMO_IMAGE_HEIGHT = 200 # Demo image display height (Increased for carousel mode) # -------------------------- Global Process State (For Stop Functionality) -------------------------- # Store current running process IDs for each module (module_num: pid) process_state = gr.State(value={}) # Flag to indicate if one-click run should be stopped one_click_stop_flag = gr.State(value=False) # -------------------------- Utility Functions (Preserve Original Logic) -------------------------- def get_server_ip(): """Get server's public IP for easy access""" 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: return "Unknown IP (please check server public IP manually)" def generate_test_id(flag): """Generate TEST_ID with timestamp, format: flag_YYYYMMDDHHMMSS_random8chars""" # Get current timestamp (accurate to seconds, format: YYYYMMDDHHMMSS) timestamp = datetime.now().strftime("%Y%m%d%H%M%S") # Keep the original 8-character random string to ensure uniqueness random_str = uuid.uuid4().hex[:8] # Combine to form the new TEST_ID return f"{flag}_{timestamp}_{random_str}" def init_test_dir(test_id): upload_dir = os.path.join(BASE_UPLOAD_DIR, test_id) result_dir = os.path.join(BASE_RESULT_DIR, test_id) os.makedirs(upload_dir, exist_ok=True) os.makedirs(result_dir, exist_ok=True) 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) 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) file_suffix = Path(file.name).suffix.lower() if hasattr(file, 'name') else Path(file).suffix.lower() original_path = file.name if hasattr(file, 'name') else file 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) final_suffix = file_suffix test_env = env_filename else: 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) final_suffix = ".jpg" else: save_path = os.path.join(upload_dir, f"{base_filename}{file_suffix}") shutil.copy(original_path, save_path) final_suffix = file_suffix test_env = None return final_suffix, save_path, test_env def clear_module_results(test_id, module_num, test_env=""): try: if module_num == 1: module1_paths = [ 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") ] for path in module1_paths: if os.path.exists(path): shutil.rmtree(path) elif module_num == 2: module2_path = os.path.join(BASE_RESULT_DIR, test_id, "envs", test_env) if os.path.exists(module2_path): shutil.rmtree(module2_path) # elif module_num == 3: # module3_files = list(Path(BASE_RESULT_DIR, test_id).glob(f"{test_id}.{test_env}.*")) # module3_dir = os.path.join(BASE_RESULT_DIR, test_id, f"relighting.{test_env}") # for file in module3_files: # if file.exists(): # os.remove(file) # if os.path.exists(module3_dir): # shutil.rmtree(module3_dir) return True except Exception as e: print(f"Failed to clear Module {module_num} results: {str(e)}") return False def get_result_files(test_id, params): results = { "input_file": os.path.join(BASE_UPLOAD_DIR, test_id, f"{test_id}{'.jpg' if params['test_type']==1 else '.mp4'}"), "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, "main_type": "video", "secondary_type": "video" }, "module3": { "final": None, "status": "Not Executed", "exists": False, "file_type": None }, "test_id": test_id, "test_env": params["test_env"] } # Module 1 Results if params["infer_1"] == 1 or os.path.exists(os.path.join(BASE_RESULT_DIR, test_id, "rego")): module1_root_dir = os.path.join(BASE_RESULT_DIR, test_id, "rego", f"{test_id}.0") if os.path.exists(module1_root_dir): for result_type, config in MODULE1_RESULT_TYPES.items(): result_dir = os.path.join(module1_root_dir, config["dir"]) if os.path.exists(result_dir): frame_files = [] for fmt in SUPPORTED_IMAGE_FORMATS: pattern = f"{config['glob_pattern']}{fmt}" frame_files.extend(list(Path(result_dir).rglob(pattern))) if frame_files: frame_files.sort() results["module1"][result_type] = str(frame_files[0]) else: results["module1"]["missing_types"].append(config["name"]) else: results["module1"]["missing_types"].append(config["name"]) has_result = any([ results["module1"]["base_color"], results["module1"]["normal"], results["module1"]["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_str = ", ".join(results["module1"]["missing_types"]) results["module1"]["status"] = f"Execution Successful (Partial results missing: {missing_str})" else: results["module1"]["status"] = "Execution Failed: No results generated" else: results["module1"]["status"] = "Execution Failed: Core directory not found" # Module 2 Results if params["infer_2"] == 1 or os.path.exists(os.path.join(BASE_RESULT_DIR, test_id, "envs", params["test_env"])): env_dir = os.path.join(BASE_RESULT_DIR, test_id, "envs", params["test_env"]) if os.path.exists(env_dir): ldr_video = None for fmt in SUPPORTED_VIDEO_FORMATS: video_path = os.path.join(env_dir, f"ldr_video_fix_first_frame{fmt}") if os.path.exists(video_path): ldr_video = video_path break env_dir_video = None for fmt in SUPPORTED_VIDEO_FORMATS: video_path = os.path.join(env_dir, f"env_dir_video_fix_first_frame{fmt}") if os.path.exists(video_path): env_dir_video = video_path break results["module2"]["ldr_video"] = ldr_video results["module2"]["env_dir_video"] = env_dir_video results["module2"]["exists"] = True if (ldr_video or env_dir_video) else False if ldr_video: results["module2"]["status"] = "Execution Successful (LDR video generated)" elif env_dir_video: results["module2"]["status"] = "Execution Successful (Only environment direction video generated, no LDR video)" else: results["module2"]["status"] = "Execution Failed: No video files generated" else: results["module2"]["status"] = "Execution Failed: Directory not found" # Module 3 Results if params["infer_3"] == 1 or any( os.path.exists(os.path.join(BASE_RESULT_DIR, test_id, f"{test_id}.{params['test_env']}{fmt}")) for fmt in SUPPORTED_IMAGE_FORMATS + SUPPORTED_VIDEO_FORMATS ): final_files = list(Path(BASE_RESULT_DIR, test_id).glob(f"{test_id}.{params['test_env']}.*")) if final_files: final_file = str(final_files[0]) file_suffix = Path(final_file).suffix.lower() results["module3"]["final"] = final_file results["module3"]["exists"] = True if file_suffix in SUPPORTED_IMAGE_FORMATS: results["module3"]["file_type"] = "image" elif file_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 run_bash_script(params, module_num, process_state): """Run bash script in background and store process ID (Compatible with Python <3.7)""" test_id = params["test_id"] env = os.environ.copy() 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"]), "CUDA_VISIBLE_DEVICES": "0", "REPO_PATH": os.getcwd() }) try: # Replace capture_output=True with stdout/stderr pipe redirection (Compatible with all Python 3 versions) process = subprocess.Popen( ["bash", BASH_SCRIPT_PATH], env=env, stdout=subprocess.PIPE, # Capture standard output stderr=subprocess.PIPE, # Capture standard error text=True, preexec_fn=os.setsid # Create new process group for easy termination ) # Update process state: store PID for current module process_state[module_num] = process.pid print(f"Module {module_num} started with PID: {process.pid}") # Wait for process completion and get output stdout, stderr = process.communicate() # Clear PID from state after completion if module_num in process_state: del process_state[module_num] if process.returncode == 0: print(f"Module {module_num} execution succeeded (PID: {process.pid})") print(f"Script stdout: {stdout}") return True, test_id, "", stdout else: 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: {str(e)}" print(error_msg) # Clear PID if error occurs if module_num in process_state: del process_state[module_num] return False, test_id, error_msg, "" def stop_module_execution(module_num, process_state, current_status): """Stop running module process (Remove temporary result cleanup logic)""" # Check if module is running if module_num not in process_state: return process_state, current_status + "\n⚠️ No running process found for this module" pid = process_state[module_num] try: # Terminate entire process group (including child processes) os.killpg(os.getpgid(pid), signal.SIGTERM) print(f"Module {module_num} process terminated (PID: {pid})") # Clear PID from state del process_state[module_num] # Remove temporary result cleanup logic when stopping, preserve already generated files return process_state, current_status + f"\n✅ Module {module_num} execution stopped successfully (PID: {pid})\nℹ️ Temporary results preserved, no cleanup performed" except ProcessLookupError: # Process already completed del process_state[module_num] return process_state, current_status + "\n⚠️ Process already completed" except Exception as e: return process_state, current_status + f"\n❌ Failed to stop module: {str(e)}" # -------------------------- Carousel Related Functions EN: -------------------------- def get_demo_images(): """Get all image files in the demo images directory""" 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}"))) # Convert to string paths and sort demo_image_paths = [str(path) for path in sorted(demo_images)] return demo_image_paths def update_carousel(index, direction, total_images): """Update carousel index""" if direction == "next": new_index = (index + 1) % total_images elif direction == "prev": new_index = (index - 1) % total_images else: new_index = index return new_index def get_current_image_path(index, image_paths): """Get the image path corresponding to the current index""" if not image_paths or index < 0 or index >= len(image_paths): return None return image_paths[index] # -------------------------- Core Business Functions (Simplified Status Output) -------------------------- 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={}) debug = False flag = "demo" # choice from "test", "debug" and "demo" # Generate unique TEST_ID and clear directory (Only reset during input file initialization) if debug: test_id = 'test_9c7f265f' else: test_id = generate_test_id(flag) 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": "", "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 } image_preview = input_path if input_file_type == "image" else None video_preview = input_path if input_file_type == "video" else None # Simplified status output: Only core status information return ( image_preview, video_preview, f"✅ Input file initialized successfully!\nFile type: {input_file_type}", base_params ) # -------------------------- Select Demo Image Initialization -------------------------- def select_demo_image(image_path): """Select a demo image and initialize""" 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={}) ) # Simulate file object for the original initialization function class MockFile: def __init__(self, path): self.name = path mock_file = MockFile(image_path) image_preview, video_preview, status, params = init_input_file(mock_file) image_visible = image_preview is not None video_visible = video_preview is not None return ( gr.update(value=image_preview, visible=image_visible, height=INPUT_PREVIEW_HEIGHT), gr.update(value=video_preview, visible=video_visible, 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"] test_env = "" use_office_env = 1 if not use_builtin_env and env_file is not None: _, env_path, 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})" base_params.update({ "test_env": test_env, "use_office_env": use_office_env }) # Simplified status output: Only configuration type 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!" print("Current light_type:", light_type, "Type:", type(light_type)) if light_type not in [0, 1, 2]: light_type = 0 print(f"Invalid light type, automatically set to default value 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 }) light_type_desc = { 0: "Original Scene + Static Light", 1: "Original Scene + Dynamic Light", 2: "Fixed First Frame + Dynamic Light" } # Simplified status output: Only key parameter summary return base_params, f"✅ Advanced parameters updated successfully!\nFrame rate: {frame_rate} | Light Type: {light_type_desc[light_type]} | Inference Steps: {num_infer_steps} | Env Strength: {env_strength}" # -------------------------- Module Execution Functions (With Stop Support) -------------------------- 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!" elif module_num == 2: return None, None, process_state, "Error: Please initialize input file first!" elif module_num == 3: return gr.update(value=None), gr.update(value=None), process_state, "Error: Please initialize input file first!" else: return None, None, None, process_state, "Error: Invalid module number!" # Check if module is already running if module_num in process_state: pid = process_state[module_num] if module_num == 1: return None, None, None, process_state, f"⚠️ Module {module_num} is already running (PID: {pid})! Please stop it first." elif module_num == 2: return None, None, process_state, f"⚠️ Module {module_num} is already running (PID: {pid})! Please stop it first." elif module_num == 3: return gr.update(value=None), gr.update(value=None), process_state, f"⚠️ Module {module_num} is already running (PID: {pid})! Please stop it first." if not re_run: results = get_result_files(params["test_id"], params) if module_num == 1: status = results["module1"]["status"] return ( results["module1"]["base_color"] if results["module1"]["exists"] else None, results["module1"]["normal"] if results["module1"]["exists"] else None, results["module1"]["roughness"] if results["module1"]["exists"] else None, process_state, f"📋 Show existing results: {status}" ) elif module_num == 2: status = results["module2"]["status"] return ( results["module2"]["ldr_video"] if results["module2"]["exists"] else None, results["module2"]["env_dir_video"] if results["module2"]["exists"] else None, process_state, f"📋 Show existing results: {status}" ) elif module_num == 3: status = results["module3"]["status"] image_result = results["module3"]["final"] if (results["module3"]["exists"] and results["module3"]["file_type"] == "image") else None video_result = results["module3"]["final"] if (results["module3"]["exists"] and results["module3"]["file_type"] == "video") else None return ( gr.update(value=image_result), gr.update(value=video_result), process_state, f"📋 Show existing results: {status}" ) else: return None, None, None, process_state, "Error: Invalid module number!" # Keep the logic of clearing old results when re-running (Avoid confusion between new and old results) clear_success = clear_module_results(params["test_id"], module_num, params["test_env"]) if not clear_success: status_msg = f"⚠️ Failed to clear old results for Module {module_num}, still attempting execution..." else: status_msg = f"✅ Cleared old results for Module {module_num}, starting re-execution..." if params["test_env"] == "" and (module_num == 2 or module_num == 3): if module_num == 2: return None, None, process_state, "Error: Please configure environment first!" elif module_num == 3: return gr.update(value=None), gr.update(value=None), process_state, "Error: Please configure environment first!" 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 # Run script and get result (process_state is updated internally) success, test_id, error_msg, stdout = run_bash_script(params, module_num, process_state) results = get_result_files(test_id, params) if module_num == 1: final_status = f"{status_msg}\n{results['module1']['status']}" if error_msg: final_status += f"\n{error_msg}" return ( results["module1"]["base_color"], results["module1"]["normal"], results["module1"]["roughness"], process_state, final_status ) elif module_num == 2: final_status = f"{status_msg}\n{results['module2']['status']}" if error_msg: final_status += f"\n{error_msg}" return ( results["module2"]["ldr_video"], results["module2"]["env_dir_video"], process_state, final_status ) elif module_num == 3: final_status = f"{status_msg}\n{results['module3']['status']}" if error_msg: final_status += f"\n{error_msg}" image_result = results["module3"]["final"] if (results["module3"]["exists"] and results["module3"]["file_type"] == "image") else None video_result = results["module3"]["final"] if (results["module3"]["exists"] and results["module3"]["file_type"] == "video") else None return ( gr.update(value=image_result), gr.update(value=video_result), process_state, final_status ) else: return None, None, None, process_state, "Invalid module number" # -------------------------- One-Click Run All Modules -------------------------- def one_click_run_all(params, process_state, stop_flag): """One-click run all modules: Module1 → Module2 → Module3 in sequence""" # Reset stop flag stop_flag = False # Pre-check if not params: return ( None, None, None, None, None, gr.update(), gr.update(), process_state, "❌ One-click run failed: Input file not initialized", "❌ One-click run failed: Input file not initialized", "❌ One-click run failed: Input file not initialized", stop_flag ) if params["test_env"] == "": return ( None, None, None, None, None, gr.update(), gr.update(), process_state, "❌ One-click run failed: Environment not configured", "❌ One-click run failed: Environment not configured", "❌ One-click run failed: Environment not configured", stop_flag ) # Run Module 1 m1_base, m1_normal, m1_rough, process_state, m1_status = run_single_module(1, params, process_state, re_run=True) if stop_flag or "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 + "\n❌ Module1 execution failed, one-click run terminated", "⏹️ Module2 not executed (previous module failed)", "⏹️ Module3 not executed (previous module failed)", stop_flag ) # Run Module 2 m2_ldr, m2_env, process_state, m2_status = run_single_module(2, params, process_state, re_run=True) if stop_flag or "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 + "\n❌ Module2 execution failed, one-click run terminated", "⏹️ Module3 not executed (previous module failed)", stop_flag ) # Run Module 3 m3_img, m3_video, process_state, m3_status = run_single_module(3, params, process_state, re_run=True) # Final status if "Failed" in m3_status or "Error" in m3_status: m3_status += "\n⚠️ One-click run completed (Module3 execution failed)" else: m3_status += "\n🎉 One-click run completed successfully!" 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 one-click run""" stop_flag = True # Stop all running modules for module_num in list(process_state.keys()): process_state, _ = stop_module_execution(module_num, process_state, "") # Update status m1_status += "\n⚠️ One-click run manually stopped" m2_status += "\n⚠️ One-click run manually stopped" m3_status += "\n⚠️ One-click run manually stopped" return ( process_state, stop_flag, m1_status, m2_status, m3_status ) # -------------------------- Single Page Layout Construction -------------------------- with gr.Blocks(title="Relit-LiVE: Relighting Model Interactive Inference Tool") as demo: gr.HTML(""" """) # Add public access information (minimal change) server_ip = get_server_ip() gr.Markdown(f""" # Relit-LiVE: Relighting Model Interactive Inference Tool > Left Panel (Input + Parameter Configuration) | Right Panel (Module 1 + Module 2 + Module 3) """) # Global State Variables base_params = gr.State(value={}) process_state = gr.State(value={}) # Track running processes one_click_stop_flag = gr.State(value=False) # One-click run stop flag # Get demo image list demo_image_paths = get_demo_images() total_images = len(demo_image_paths) # Carousel state current_index = gr.State(value=0 if total_images > 0 else -1) # Main Layout: Two Columns (Width Ratio 1:1.5, Right column wider for 3 modules) with gr.Row(): # -------------------------- Left Column: Input + Parameter Configuration -------------------------- with gr.Column(scale=1, min_width=400): gr.Markdown("## 📥 Input & Parameter Configuration") # 1. Input File Upload 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") # -------------------------- Predefine Preview Components (Solve Undefined Variable Issue) -------------------------- # Input Preview (Adaptive width, fixed height) 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) # -------------------------- Optimize: Demo Images as Carousel -------------------------- gr.Markdown("### 📸 Demo Images") if demo_image_paths and total_images > 0: # Carousel container gr.HTML('
') # Main image display carousel_image = gr.Image( value=demo_image_paths[0] if total_images > 0 else None, label="Demo Image", height=DEMO_IMAGE_HEIGHT, interactive=False ) # Navigation buttons with gr.Row(): prev_btn = gr.Button("◀️ Previous", elem_classes=["carousel-btn", "carousel-prev"]) next_btn = gr.Button("Next ▶️", elem_classes=["carousel-btn", "carousel-next"]) # Indicators gr.HTML('
') indicator_btns = [] for i in range(total_images): btn = gr.Button( "", elem_classes=["indicator-dot"] + (["active"] if i == 0 else []), size="sm", visible=False # Hide default button style, use CSS style ) indicator_btns.append(btn) # Click indicator to switch image btn.click( fn=lambda idx=i: ( idx, # Update current index demo_image_paths[idx] if idx < len(demo_image_paths) else None # Update image ), inputs=[], outputs=[current_index, carousel_image] ) gr.HTML('
') # Select current image button select_current_btn = gr.Button( "✅ Select This Image", variant="primary", elem_classes=["select-image-btn"] ) gr.HTML('
') # Carousel navigation logic def update_carousel_ui(index): """Update carousel UI""" if index < 0 or index >= len(demo_image_paths): return None return demo_image_paths[index] # Previous button prev_btn.click( fn=lambda idx: update_carousel(idx, "prev", total_images), inputs=[current_index], outputs=[current_index] ).then( fn=update_carousel_ui, inputs=[current_index], outputs=[carousel_image] ) # Next button next_btn.click( fn=lambda idx: update_carousel(idx, "next", total_images), inputs=[current_index], outputs=[current_index] ).then( fn=update_carousel_ui, inputs=[current_index], outputs=[carousel_image] ) # Select current image select_current_btn.click( fn=lambda idx: select_demo_image(demo_image_paths[idx] if 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*") # 2. Environment Configuration 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="Pink_Sunrise" ) 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) # 3. Advanced Parameters 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 (0=Original+Static, 1=Original+Dynamic, 2=Fixed First Frame+Dynamic)", 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) # One-Click Run All Modules 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, placeholder="Click button above to start one-click run...") # -------------------------- Right Column: Module 1 + Module 2 + Module 3 (Vertical Arrangement) -------------------------- with gr.Column(scale=1.5, min_width=600): # Module 1: Inverse Rendering (3 results split horizontally) 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) # Module 1 Visualization Results gr.Markdown("### Partial Visualization Results (Base Color | Normal Map | Roughness Map)") 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) # Module 2: Environment Processing (2 videos split horizontally) 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) # Module 2 Visualization Results gr.Markdown("### Visualization Results (LDR Video | Environment Direction Video)") 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) # Module 3: Relighting (1 result full width) 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) # Module 3 Visualization Results gr.Markdown("### Visualization Results (Auto-adapt Image/Video)") 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 ) # -------------------------- Event Bindings -------------------------- # Input File Initialization def init_and_show_preview(input_file): image_preview, video_preview, status, params = init_input_file(input_file) image_visible = image_preview is not None video_visible = video_preview is not None return ( gr.update(value=image_preview, visible=image_visible, height=INPUT_PREVIEW_HEIGHT), gr.update(value=video_preview, visible=video_visible, height=INPUT_PREVIEW_HEIGHT), status, params ) init_input_btn.click( fn=init_and_show_preview, inputs=[input_file], outputs=[input_image_preview, input_video_preview, input_status, base_params] ) # Environment Configuration Visibility use_builtin_env.change( fn=lambda x: gr.update(visible=not x), inputs=[use_builtin_env], outputs=[env_file] ) update_env_btn.click( fn=update_env_config, inputs=[base_params, use_builtin_env, builtin_env_choice, env_file], outputs=[base_params, env_status] ) # Advanced Parameters Update update_advanced_btn.click( fn=update_advanced_params, inputs=[base_params, frame, frame_rate, env_strength, num_infer_steps, worw, light_type], outputs=[base_params, advanced_status] ) # -------------------------- Module 1: Execution/Stop/Display -------------------------- run_module1_btn.click( fn=lambda params, process_state: run_single_module(1, params, process_state, re_run=True), inputs=[base_params, process_state], outputs=[module1_base_color, module1_normal, module1_roughness, process_state, module1_status] ) stop_module1_btn.click( fn=lambda ps, cs: stop_module_execution(1, ps, cs), inputs=[process_state, module1_status], outputs=[process_state, module1_status] ) show_module1_btn.click( fn=lambda params, process_state: run_single_module(1, params, process_state, re_run=False), inputs=[base_params, process_state], outputs=[module1_base_color, module1_normal, module1_roughness, process_state, module1_status] ) # -------------------------- Module 2: Execution/Stop/Display -------------------------- run_module2_btn.click( fn=lambda params, process_state: run_single_module(2, params, process_state, re_run=True), inputs=[base_params, process_state], outputs=[module2_ldr_video, module2_env_video, process_state, module2_status] ) stop_module2_btn.click( fn=lambda ps, cs: stop_module_execution(2, ps, cs), inputs=[process_state, module2_status], outputs=[process_state, module2_status] ) show_module2_btn.click( fn=lambda params, process_state: run_single_module(2, params, process_state, re_run=False), inputs=[base_params, process_state], outputs=[module2_ldr_video, module2_env_video, process_state, module2_status] ) # -------------------------- Module 3: Execution/Stop/Display -------------------------- run_module3_btn.click( fn=lambda params, process_state: run_single_module(3, params, process_state, re_run=True), inputs=[base_params, process_state], outputs=[module3_image_result, module3_video_result, process_state, module3_status] ).then( fn=lambda image_val, video_val, status: ( 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 ), inputs=[module3_image_result, module3_video_result, module3_status], outputs=[module3_image_result, module3_video_result, module3_status] ) stop_module3_btn.click( fn=lambda ps, cs: stop_module_execution(3, ps, cs), inputs=[process_state, module3_status], outputs=[process_state, module3_status] ) show_module3_btn.click( fn=lambda params, process_state: run_single_module(3, params, process_state, re_run=False), inputs=[base_params, process_state], outputs=[module3_image_result, module3_video_result, process_state, module3_status] ).then( fn=lambda image_val, video_val, status: ( 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 ), inputs=[module3_image_result, module3_video_result, module3_status], outputs=[module3_image_result, module3_video_result, module3_status] ) # -------------------------- One-Click Run Event Bindings -------------------------- one_click_run_btn.click( fn=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( fn=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( fn=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] ) # -------------------------- Run Gradio (Public Access Enabled) -------------------------- if __name__ == "__main__": os.makedirs(BASE_UPLOAD_DIR, exist_ok=True) os.makedirs(BASE_RESULT_DIR, exist_ok=True) # Print public access information server_ip = get_server_ip() print("="*60) print("Relit-LiVE: Relighting Model Interactive Inference Tool") print(f"📡 Public Access URL: http://{server_ip}:7861") print("🔧 All original functions are preserved") print("🖼️ Demo images displayed in carousel (slider) layout") print("⚠️ No authentication - use with caution") print("="*60) demo.launch( server_name="0.0.0.0", # Listen on all interfaces (public access) server_port=7861, # Fixed port (can be modified if needed) share=False, show_error=True )