run_full_inference_gradio.py

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("""
    <style>
    /* Carousel container style */
    .carousel-container {
        position: relative;
        width: 100%;
        max-width: 600px;
        margin: 20px auto;
        overflow: hidden;
        border-radius: 12px;
        box-shadow: 0 4px 20px rgba(0,0,0,0.1);
    }
    /* Navigation button style */
    .carousel-btn {
        position: absolute;
        top: 50%;
        transform: translateY(-50%);
        width: 40px;
        height: 40px;
        border-radius: 50%;
        background-color: rgba(255,255,255,0.8);
        border: none;
        cursor: pointer;
        font-size: 20px;
        display: flex;
        align-items: center;
        justify-content: center;
        z-index: 10;
        transition: all 0.3s ease;
    }
    .carousel-btn:hover {
        background-color: rgba(255,255,255,1);
        box-shadow: 0 2px 10px rgba(0,0,0,0.2);
    }
    .carousel-prev {
        left: 10px;
    }
    .carousel-next {
        right: 10px;
    }
    /* Image indicator style */
    .carousel-indicators {
        display: flex;
        justify-content: center;
        gap: 8px;
        margin-top: 15px;
    }
    .indicator-dot {
        width: 10px;
        height: 10px;
        border-radius: 50%;
        background-color: #ddd;
        cursor: pointer;
        transition: all 0.3s ease;
    }
    .indicator-dot.active {
        background-color: #666;
        transform: scale(1.2);
    }
    /* Select button style */
    .select-image-btn {
        margin-top: 15px;
        width: 100%;
        max-width: 200px;
    }
    </style>
    """)
    
    # 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('<div class="carousel-container">')
                
                # 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('<div class="carousel-indicators">')
                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('</div>')
                
                # Select current image button
                select_current_btn = gr.Button(
                    "✅ Select This Image",
                    variant="primary",
                    elem_classes=["select-image-btn"]
                )
                
                gr.HTML('</div>')
                
                # 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
    )