app.py

import os
import sys
import gradio as gr
import torch
import argparse
from PIL import Image
import numpy as np
import torchvision.transforms as transforms
from moviepy.editor import VideoFileClip
from diffusers.utils import load_image, load_video
import spaces

project_root = os.path.dirname(os.path.abspath(__file__))
os.environ["GRADIO_TEMP_DIR"] = os.path.join(project_root, "tmp", "gradio")
sys.path.append(project_root)

try:
    sys.path.append(os.path.join(project_root, "submodules/MoGe"))
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
except:
    print("Warning: MoGe not found, motion transfer will not be applied")

HERE_PATH = os.path.normpath(os.path.dirname(__file__))
sys.path.insert(0, HERE_PATH)
from huggingface_hub import hf_hub_download
hf_hub_download(repo_id="EXCAI/Diffusion-As-Shader", filename='spatracker/spaT_final.pth', local_dir=f'{HERE_PATH}/checkpoints/')

from models.pipelines import DiffusionAsShaderPipeline, FirstFrameRepainter, CameraMotionGenerator, ObjectMotionGenerator
from submodules.MoGe.moge.model import MoGeModel

# Parse command line arguments
parser = argparse.ArgumentParser(description="Diffusion as Shader Web UI")
parser.add_argument("--port", type=int, default=7860, help="Port to run the web UI on")
parser.add_argument("--share", action="store_true", help="Share the web UI")
parser.add_argument("--gpu", type=int, default=0, help="GPU device ID")
parser.add_argument("--model_path", type=str, default="EXCAI/Diffusion-As-Shader", help="Path to model checkpoint")
parser.add_argument("--output_dir", type=str, default="tmp", help="Output directory")
args = parser.parse_args()

# Use the original GPU ID throughout the entire code for consistency
GPU_ID = args.gpu
DEFAULT_MODEL_PATH = args.model_path
OUTPUT_DIR = args.output_dir

# Create necessary directories
os.makedirs("outputs", exist_ok=True)
# Create project tmp directory instead of using system temp
os.makedirs(os.path.join(project_root, "tmp"), exist_ok=True)
os.makedirs(os.path.join(project_root, "tmp", "gradio"), exist_ok=True)
def load_media(media_path, max_frames=49, transform=None):
    """Load video or image frames and convert to tensor
    
    Args:
        media_path (str): Path to video or image file
        max_frames (int): Maximum number of frames to load
        transform (callable): Transform to apply to frames
        
    Returns:
        Tuple[torch.Tensor, float, bool]: Video tensor [T,C,H,W], FPS, and is_video flag
    """
    if transform is None:
        transform = transforms.Compose([
            transforms.Resize((480, 720)),
            transforms.ToTensor()
        ])
    
    # Determine if input is video or image based on extension
    ext = os.path.splitext(media_path)[1].lower()
    is_video = ext in ['.mp4', '.avi', '.mov']
    
    if is_video:
        frames = load_video(media_path)
        fps = len(frames) / VideoFileClip(media_path).duration
    else:
        # Handle image as single frame
        image = load_image(media_path)
        frames = [image]
        fps = 8  # Default fps for images
    
    # Ensure we have exactly max_frames
    if len(frames) > max_frames:
        frames = frames[:max_frames]
    elif len(frames) < max_frames:
        last_frame = frames[-1]
        while len(frames) < max_frames:
            frames.append(last_frame.copy())
            
    # Convert frames to tensor
    video_tensor = torch.stack([transform(frame) for frame in frames])
    
    return video_tensor, fps, is_video

def save_uploaded_file(file):
    if file is None:
        return None
        
    # Use project tmp directory instead of system temp
    temp_dir = os.path.join(project_root, "tmp")
    
    if hasattr(file, 'name'):
        filename = file.name
    else:
        # Generate a unique filename if name attribute is missing
        import uuid
        ext = ".tmp"
        if hasattr(file, 'content_type'):
            if "image" in file.content_type:
                ext = ".png"
            elif "video" in file.content_type:
                ext = ".mp4"
        filename = f"{uuid.uuid4()}{ext}"
    
    temp_path = os.path.join(temp_dir, filename)
    
    try:
        # Check if file is a FileStorage object or already a path
        if hasattr(file, 'save'):
            file.save(temp_path)
        elif isinstance(file, str):
            # It's already a path
            return file
        else:
            # Try to read and save the file
            with open(temp_path, 'wb') as f:
                f.write(file.read() if hasattr(file, 'read') else file)
    except Exception as e:
        print(f"Error saving file: {e}")
        return None
        
    return temp_path

das_pipeline = None
moge_model = None

@spaces.GPU
def get_das_pipeline():
    global das_pipeline
    if das_pipeline is None:
        das_pipeline = DiffusionAsShaderPipeline(gpu_id=GPU_ID, output_dir=OUTPUT_DIR)
    return das_pipeline

@spaces.GPU
def get_moge_model():
    global moge_model
    if moge_model is None:
        das = get_das_pipeline()
        moge_model = MoGeModel.from_pretrained("Ruicheng/moge-vitl").to(das.device)
    return moge_model


def process_motion_transfer(source, prompt, mt_repaint_option, mt_repaint_image):
    """Process video motion transfer task"""
    try:
        # Save uploaded files
        input_video_path = save_uploaded_file(source)
        if input_video_path is None:
            return None
        
        print(f"DEBUG: Repaint option: {mt_repaint_option}")
        print(f"DEBUG: Repaint image: {mt_repaint_image}")
        

        das = get_das_pipeline()
        video_tensor, fps, is_video = load_media(input_video_path)
        if not is_video:
            tracking_method = "moge"
            print("Image input detected, using MoGe for tracking video generation.")
        else:
            tracking_method = "spatracker"
        
        repaint_img_tensor = None
        if mt_repaint_image is not None:
            repaint_path = save_uploaded_file(mt_repaint_image)
            repaint_img_tensor, _, _ = load_media(repaint_path)
            repaint_img_tensor = repaint_img_tensor[0]
        elif mt_repaint_option == "Yes":
            repainter = FirstFrameRepainter(gpu_id=GPU_ID, output_dir=OUTPUT_DIR)
            repaint_img_tensor = repainter.repaint(
                video_tensor[0], 
                prompt=prompt,
                depth_path=None
            )
        tracking_tensor = None
        if tracking_method == "moge":
            moge = get_moge_model()
            infer_result = moge.infer(video_tensor[0].to(das.device))  # [C, H, W] in range [0,1]
            H, W = infer_result["points"].shape[0:2]
            pred_tracks = infer_result["points"].unsqueeze(0).repeat(49, 1, 1, 1) #[T, H, W, 3]
            poses = torch.eye(4).unsqueeze(0).repeat(49, 1, 1)

            pred_tracks_flatten = pred_tracks.reshape(video_tensor.shape[0], H*W, 3)
  
            cam_motion = CameraMotionGenerator(None)
            cam_motion.set_intr(infer_result["intrinsics"])
  
            pred_tracks = cam_motion.w2s(pred_tracks_flatten, poses).reshape([video_tensor.shape[0], H, W, 3]) # [T, H, W, 3]

            _, tracking_tensor = das.visualize_tracking_moge(
                pred_tracks.cpu().numpy(), 
                infer_result["mask"].cpu().numpy()
            )
            print('Export tracking video via MoGe')
        else:
            pred_tracks, pred_visibility, T_Firsts = das.generate_tracking_spatracker(video_tensor)

            _, tracking_tensor = das.visualize_tracking_spatracker(video_tensor, pred_tracks, pred_visibility, T_Firsts)
            print('Export tracking video via SpaTracker')

        output_path = das.apply_tracking(
            video_tensor=video_tensor,
            fps=8,
            tracking_tensor=tracking_tensor,
            img_cond_tensor=repaint_img_tensor,
            prompt=prompt,
            checkpoint_path=DEFAULT_MODEL_PATH
        )
        
        return output_path
    except Exception as e:
        import traceback
        print(f"Processing failed: {str(e)}\n{traceback.format_exc()}")
        return None


def process_camera_control(source, prompt, camera_motion, tracking_method):
    """Process camera control task"""
    try:
        # Save uploaded files
        input_media_path = save_uploaded_file(source)
        if input_media_path is None:
            return None
        
        print(f"DEBUG: Camera motion: '{camera_motion}'")
        print(f"DEBUG: Tracking method: '{tracking_method}'")
        
        das = get_das_pipeline()

        video_tensor, fps, is_video = load_media(input_media_path)
        if not is_video and tracking_method == "spatracker":
            tracking_method = "moge"
            print("Image input detected with spatracker selected, switching to MoGe")

        cam_motion = CameraMotionGenerator(camera_motion)
        repaint_img_tensor = None
        tracking_tensor = None

        if tracking_method == "moge":
            moge = get_moge_model()
            
            infer_result = moge.infer(video_tensor[0].to(das.device))  # [C, H, W] in range [0,1]
            H, W = infer_result["points"].shape[0:2]
            pred_tracks = infer_result["points"].unsqueeze(0).repeat(49, 1, 1, 1) #[T, H, W, 3]
            cam_motion.set_intr(infer_result["intrinsics"])

            if camera_motion:
                poses = cam_motion.get_default_motion() # shape: [49, 4, 4]
                print("Camera motion applied")
            else:
                poses = torch.eye(4).unsqueeze(0).repeat(49, 1, 1)

            pred_tracks_flatten = pred_tracks.reshape(video_tensor.shape[0], H*W, 3)
            pred_tracks = cam_motion.w2s(pred_tracks_flatten, poses).reshape([video_tensor.shape[0], H, W, 3]) # [T, H, W, 3]

            _, tracking_tensor = das.visualize_tracking_moge(
                pred_tracks.cpu().numpy(), 
                infer_result["mask"].cpu().numpy()
            )
            print('Export tracking video via MoGe')
        else:
           
            pred_tracks, pred_visibility, T_Firsts = das.generate_tracking_spatracker(video_tensor)
            if camera_motion:
                poses = cam_motion.get_default_motion() # shape: [49, 4, 4]
                pred_tracks = cam_motion.apply_motion_on_pts(pred_tracks, poses)
                print("Camera motion applied")

            _, tracking_tensor = das.visualize_tracking_spatracker(video_tensor, pred_tracks, pred_visibility, T_Firsts)
            print('Export tracking video via SpaTracker')
        
        
        output_path = das.apply_tracking(
            video_tensor=video_tensor,
            fps=8,
            tracking_tensor=tracking_tensor,
            img_cond_tensor=repaint_img_tensor,
            prompt=prompt,
            checkpoint_path=DEFAULT_MODEL_PATH
        )
        
        return output_path
    except Exception as e:
        import traceback
        print(f"Processing failed: {str(e)}\n{traceback.format_exc()}")
        return None


def process_object_manipulation(source, prompt, object_motion, object_mask, tracking_method):
    """Process object manipulation task"""
    try:
        # Save uploaded files
        input_image_path = save_uploaded_file(source)
        if input_image_path is None:
            return None
            
        object_mask_path = save_uploaded_file(object_mask)
        if object_mask_path is None:
            print("Object mask not provided")
            return None
        
        
        das = get_das_pipeline()
        video_tensor, fps, is_video = load_media(input_image_path)
        if not is_video and tracking_method == "spatracker":
            tracking_method = "moge"
            print("Image input detected with spatracker selected, switching to MoGe")

        
        mask_image = Image.open(object_mask_path).convert('L')
        mask_image = transforms.Resize((480, 720))(mask_image)
        mask = torch.from_numpy(np.array(mask_image) > 127)

        motion_generator = ObjectMotionGenerator(device=das.device)
        repaint_img_tensor = None
        tracking_tensor = None
        if tracking_method == "moge":
            moge = get_moge_model()
            
            
            infer_result = moge.infer(video_tensor[0].to(das.device))  # [C, H, W] in range [0,1]
            H, W = infer_result["points"].shape[0:2]
            pred_tracks = infer_result["points"].unsqueeze(0).repeat(49, 1, 1, 1) #[T, H, W, 3]

            pred_tracks = motion_generator.apply_motion(
                pred_tracks=pred_tracks,
                mask=mask,
                motion_type=object_motion,
                distance=50,
                num_frames=49,
                tracking_method="moge"
            )
            print(f"Object motion '{object_motion}' applied using provided mask")
            poses = torch.eye(4).unsqueeze(0).repeat(49, 1, 1)
            pred_tracks_flatten = pred_tracks.reshape(video_tensor.shape[0], H*W, 3)
            
            
            cam_motion = CameraMotionGenerator(None)
            cam_motion.set_intr(infer_result["intrinsics"])
            pred_tracks = cam_motion.w2s(pred_tracks_flatten, poses).reshape([video_tensor.shape[0], H, W, 3]) # [T, H, W, 3]

            _, tracking_tensor = das.visualize_tracking_moge(
                pred_tracks.cpu().numpy(), 
                infer_result["mask"].cpu().numpy()
            )
            print('Export tracking video via MoGe')
        else:
            
            pred_tracks, pred_visibility, T_Firsts = das.generate_tracking_spatracker(video_tensor)
            
            
            pred_tracks = motion_generator.apply_motion(
                pred_tracks=pred_tracks.squeeze(),
                mask=mask,
                motion_type=object_motion,
                distance=50,
                num_frames=49,
                tracking_method="spatracker"
            ).unsqueeze(0)
            print(f"Object motion '{object_motion}' applied using provided mask")
            
            
            _, tracking_tensor = das.visualize_tracking_spatracker(video_tensor, pred_tracks, pred_visibility, T_Firsts)
            print('Export tracking video via SpaTracker')
        
        
        output_path = das.apply_tracking(
            video_tensor=video_tensor,
            fps=8,
            tracking_tensor=tracking_tensor,
            img_cond_tensor=repaint_img_tensor,
            prompt=prompt,
            checkpoint_path=DEFAULT_MODEL_PATH
        )
        
        return output_path
    except Exception as e:
        import traceback
        print(f"Processing failed: {str(e)}\n{traceback.format_exc()}")
        return None


def process_mesh_animation(source, prompt, tracking_video, ma_repaint_option, ma_repaint_image):
    """Process mesh animation task"""
    try:
        # Save uploaded files
        input_video_path = save_uploaded_file(source)
        if input_video_path is None:
            return None
            
        tracking_video_path = save_uploaded_file(tracking_video)
        if tracking_video_path is None:
            return None
        
        
        das = get_das_pipeline()   
        video_tensor, fps, is_video = load_media(input_video_path)
        tracking_tensor, tracking_fps, _ = load_media(tracking_video_path)
        repaint_img_tensor = None
        if ma_repaint_image is not None:
            repaint_path = save_uploaded_file(ma_repaint_image)
            repaint_img_tensor, _, _ = load_media(repaint_path)
            repaint_img_tensor = repaint_img_tensor[0]  # 获取第一帧
        elif ma_repaint_option == "Yes":

            repainter = FirstFrameRepainter(gpu_id=GPU_ID, output_dir=OUTPUT_DIR)
            repaint_img_tensor = repainter.repaint(
                video_tensor[0], 
                prompt=prompt,
                depth_path=None
            )

        output_path = das.apply_tracking(
            video_tensor=video_tensor,
            fps=8,
            tracking_tensor=tracking_tensor,
            img_cond_tensor=repaint_img_tensor,
            prompt=prompt,
            checkpoint_path=DEFAULT_MODEL_PATH
        )
        
        return output_path
    except Exception as e:
        import traceback
        print(f"Processing failed: {str(e)}\n{traceback.format_exc()}")
        return None

# Create Gradio interface with updated layout
with gr.Blocks(title="Diffusion as Shader") as demo:
    gr.Markdown("# Diffusion as Shader Web UI")
    gr.Markdown("### [Project Page](https://igl-hkust.github.io/das/) | [GitHub](https://github.com/IGL-HKUST/DiffusionAsShader)")
    
    with gr.Row():
        left_column = gr.Column(scale=1)
        right_column = gr.Column(scale=1)

    with right_column:
        output_video = gr.Video(label="Generated Video")

    with left_column:
        source = gr.File(label="Source", file_types=["image", "video"])
        common_prompt = gr.Textbox(label="Prompt", lines=2)
        gr.Markdown(f"**Using GPU: {GPU_ID}**")
        
        with gr.Tabs() as task_tabs:
            # Motion Transfer tab
            with gr.TabItem("Motion Transfer"):
                gr.Markdown("## Motion Transfer")
                
                # Simplified controls - Radio buttons for Yes/No and separate file upload
                with gr.Row():
                    mt_repaint_option = gr.Radio(
                        label="Repaint First Frame",
                        choices=["No", "Yes"],
                        value="No"
                    )
                gr.Markdown("### Note: If you want to use your own image as repainted first frame, please upload the image in below.")
                # Custom image uploader (always visible)
                mt_repaint_image = gr.File(
                    label="Custom Repaint Image", 
                    file_types=["image"]
                )
                
                # Add run button for Motion Transfer tab
                mt_run_btn = gr.Button("Run Motion Transfer", variant="primary", size="lg")
                
                # Connect to process function
                mt_run_btn.click(
                    fn=process_motion_transfer,
                    inputs=[
                        source, common_prompt,
                        mt_repaint_option, mt_repaint_image
                    ],
                    outputs=[output_video]
                )
            
            # Camera Control tab
            with gr.TabItem("Camera Control"):
                gr.Markdown("## Camera Control")
                
                cc_camera_motion = gr.Textbox(
                    label="Current Camera Motion Sequence",
                    placeholder="Your camera motion sequence will appear here...",
                    interactive=False
                )
                
                # Use tabs for different motion types
                with gr.Tabs() as cc_motion_tabs:
                    # Translation tab
                    with gr.TabItem("Translation (trans)"):
                        with gr.Row():
                            cc_trans_x = gr.Slider(minimum=-1.0, maximum=1.0, value=0.0, step=0.05, label="X-axis Movement")
                            cc_trans_y = gr.Slider(minimum=-1.0, maximum=1.0, value=0.0, step=0.05, label="Y-axis Movement")
                            cc_trans_z = gr.Slider(minimum=-1.0, maximum=1.0, value=0.0, step=0.05, label="Z-axis Movement (depth)")
                        
                        with gr.Row():
                            cc_trans_start = gr.Number(minimum=0, maximum=48, value=0, step=1, label="Start Frame", precision=0)
                            cc_trans_end = gr.Number(minimum=0, maximum=48, value=48, step=1, label="End Frame", precision=0)
                        
                        cc_trans_note = gr.Markdown("""
                        **Translation Notes:**
                        - Positive X: Move right, Negative X: Move left
                        - Positive Y: Move down, Negative Y: Move up
                        - Positive Z: Zoom in, Negative Z: Zoom out
                        """)
                        
                        # Add translation button in the Translation tab
                        cc_add_trans = gr.Button("Add Camera Translation", variant="secondary")
                        
                        # Function to add translation motion
                        def add_translation_motion(current_motion, trans_x, trans_y, trans_z, trans_start, trans_end):
                            # Format: trans dx dy dz [start_frame end_frame]
                            frame_range = f" {int(trans_start)} {int(trans_end)}" if trans_start != 0 or trans_end != 48 else ""
                            new_motion = f"trans {trans_x:.2f} {trans_y:.2f} {trans_z:.2f}{frame_range}"
                            
                            # Append to existing motion string with semicolon separator if needed
                            if current_motion and current_motion.strip():
                                updated_motion = f"{current_motion}; {new_motion}"
                            else:
                                updated_motion = new_motion
                            
                            return updated_motion
                        
                        # Connect translation button
                        cc_add_trans.click(
                            fn=add_translation_motion,
                            inputs=[
                                cc_camera_motion,
                                cc_trans_x, cc_trans_y, cc_trans_z, cc_trans_start, cc_trans_end
                            ],
                            outputs=[cc_camera_motion]
                        )
                    
                    # Rotation tab
                    with gr.TabItem("Rotation (rot)"):
                        with gr.Row():
                            cc_rot_axis = gr.Dropdown(choices=["x", "y", "z"], value="y", label="Rotation Axis")
                            cc_rot_angle = gr.Slider(minimum=-30, maximum=30, value=5, step=1, label="Rotation Angle (degrees)")
                        
                        with gr.Row():
                            cc_rot_start = gr.Number(minimum=0, maximum=48, value=0, step=1, label="Start Frame", precision=0)
                            cc_rot_end = gr.Number(minimum=0, maximum=48, value=48, step=1, label="End Frame", precision=0)
                        
                        cc_rot_note = gr.Markdown("""
                        **Rotation Notes:**
                        - X-axis rotation: Tilt camera up/down
                        - Y-axis rotation: Pan camera left/right
                        - Z-axis rotation: Roll camera
                        """)
                        
                        # Add rotation button in the Rotation tab
                        cc_add_rot = gr.Button("Add Camera Rotation", variant="secondary")
                        
                        # Function to add rotation motion
                        def add_rotation_motion(current_motion, rot_axis, rot_angle, rot_start, rot_end):
                            # Format: rot axis angle [start_frame end_frame]
                            frame_range = f" {int(rot_start)} {int(rot_end)}" if rot_start != 0 or rot_end != 48 else ""
                            new_motion = f"rot {rot_axis} {rot_angle}{frame_range}"
                            
                            # Append to existing motion string with semicolon separator if needed
                            if current_motion and current_motion.strip():
                                updated_motion = f"{current_motion}; {new_motion}"
                            else:
                                updated_motion = new_motion
                            
                            return updated_motion
                        
                        # Connect rotation button
                        cc_add_rot.click(
                            fn=add_rotation_motion,
                            inputs=[
                                cc_camera_motion,
                                cc_rot_axis, cc_rot_angle, cc_rot_start, cc_rot_end
                            ],
                            outputs=[cc_camera_motion]
                        )
                
                # Add a clear button to reset the motion sequence
                cc_clear_motion = gr.Button("Clear All Motions", variant="stop")
                
                def clear_camera_motion():
                    return ""
                
                cc_clear_motion.click(
                    fn=clear_camera_motion,
                    inputs=[],
                    outputs=[cc_camera_motion]
                )

                cc_tracking_method = gr.Radio(
                    label="Tracking Method",
                    choices=["spatracker", "moge"],
                    value="moge"
                )
                
                # Add run button for Camera Control tab
                cc_run_btn = gr.Button("Run Camera Control", variant="primary", size="lg")
                
                # Connect to process function
                cc_run_btn.click(
                    fn=process_camera_control,
                    inputs=[
                        source, common_prompt,
                        cc_camera_motion, cc_tracking_method
                    ],
                    outputs=[output_video]
                )
            
            # Object Manipulation tab
            with gr.TabItem("Object Manipulation"):
                gr.Markdown("## Object Manipulation")
                om_object_mask = gr.File(
                    label="Object Mask Image", 
                    file_types=["image"]
                )
                gr.Markdown("Upload a binary mask image, white areas indicate the object to manipulate")
                om_object_motion = gr.Dropdown(
                    label="Object Motion Type",
                    choices=["up", "down", "left", "right", "front", "back", "rot"],
                    value="up"
                )
                om_tracking_method = gr.Radio(
                    label="Tracking Method",
                    choices=["spatracker", "moge"],
                    value="moge"
                )
                
                # Add run button for Object Manipulation tab
                om_run_btn = gr.Button("Run Object Manipulation", variant="primary", size="lg")
                
                # Connect to process function
                om_run_btn.click(
                    fn=process_object_manipulation,
                    inputs=[
                        source, common_prompt,
                        om_object_motion, om_object_mask, om_tracking_method
                    ],
                    outputs=[output_video]
                )
            
            # Animating meshes to video tab
            with gr.TabItem("Animating meshes to video"):
                gr.Markdown("## Mesh Animation to Video")
                gr.Markdown("""
                    Note: Currently only supports tracking videos generated with Blender (version > 4.0).
                    Please run the script `scripts/blender.py` in your Blender project to generate tracking videos.
                """)
                ma_tracking_video = gr.File(
                    label="Tracking Video",
                    file_types=["video"]
                )
                gr.Markdown("Tracking video needs to be generated from Blender")
                
                # Simplified controls - Radio buttons for Yes/No and separate file upload
                with gr.Row():
                    ma_repaint_option = gr.Radio(
                        label="Repaint First Frame",
                        choices=["No", "Yes"],
                        value="No"
                    )
                gr.Markdown("### Note: If you want to use your own image as repainted first frame, please upload the image in below.")
                # Custom image uploader (always visible)
                ma_repaint_image = gr.File(
                    label="Custom Repaint Image", 
                    file_types=["image"]
                )
                
                # Add run button for Mesh Animation tab
                ma_run_btn = gr.Button("Run Mesh Animation", variant="primary", size="lg")
                
                # Connect to process function
                ma_run_btn.click(
                    fn=process_mesh_animation,
                    inputs=[
                        source, common_prompt,
                        ma_tracking_video, ma_repaint_option, ma_repaint_image
                    ],
                    outputs=[output_video]
                )

# Launch interface
if __name__ == "__main__":
    print(f"Using GPU: {GPU_ID}")
    print(f"Web UI will start on port {args.port}")
    if args.share:
        print("Creating public link for remote access")
    
    # Launch interface
    demo.launch(share=args.share, server_port=args.port)