import os import sys import spaces 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 from tqdm import tqdm from image_gen_aux import DepthPreprocessor 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")) sys.path.append(os.path.join(project_root, "submodules/vggt")) 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 from submodules.vggt.vggt.utils.pose_enc import pose_encoding_to_extri_intri from submodules.vggt.vggt.models.vggt import VGGT # 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: # Load video file info video_clip = VideoFileClip(media_path) duration = video_clip.duration original_fps = video_clip.fps # Case 1: Video longer than 6 seconds, sample first 6 seconds + 1 frame if duration > 6.0: sampling_fps = 8 # 8 frames per second frames = load_video(media_path, sampling_fps=sampling_fps, max_frames=max_frames) fps = sampling_fps # Cases 2 and 3: Video shorter than 6 seconds else: # Load all frames frames = load_video(media_path) # Case 2: Total frames less than max_frames, need interpolation if len(frames) < max_frames: fps = len(frames) / duration # Keep original fps # Evenly interpolate to max_frames indices = np.linspace(0, len(frames) - 1, max_frames) new_frames = [] for i in indices: idx = int(i) new_frames.append(frames[idx]) frames = new_frames # Case 3: Total frames more than max_frames but video less than 6 seconds else: # Evenly sample to max_frames indices = np.linspace(0, len(frames) - 1, max_frames) new_frames = [] for i in indices: idx = int(i) new_frames.append(frames[idx]) frames = new_frames fps = max_frames / duration # New fps to maintain duration else: # Handle image as single frame image = load_image(media_path) frames = [image] fps = 8 # Default fps for images # Duplicate frame to max_frames while len(frames) < max_frames: frames.append(frames[0].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 vggt_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 @spaces.GPU def get_vggt_model(): global vggt_model if vggt_model is None: das = get_das_pipeline() vggt_model = VGGT.from_pretrained("facebook/VGGT-1B").to(das.device) return vggt_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, 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) das.fps = fps # 设置 das.fps 为 load_media 返回的 fps if not is_video: tracking_method = "moge" print("Image input detected, using MoGe for tracking video generation.") else: tracking_method = "cotracker" 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 tracking_path = 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_path, tracking_tensor = das.visualize_tracking_moge( pred_tracks.cpu().numpy(), infer_result["mask"].cpu().numpy() ) print('Export tracking video via MoGe') else: # 使用 cotracker pred_tracks, pred_visibility = generate_tracking_cotracker(video_tensor) tracking_path, tracking_tensor = das.visualize_tracking_cotracker(pred_tracks, pred_visibility) print('Export tracking video via cotracker') output_path = das.apply_tracking( video_tensor=video_tensor, fps=fps, # 使用 load_media 返回的 fps tracking_tensor=tracking_tensor, img_cond_tensor=repaint_img_tensor, prompt=prompt, checkpoint_path=DEFAULT_MODEL_PATH ) return tracking_path, output_path except Exception as e: import traceback print(f"Processing failed: {str(e)}\n{traceback.format_exc()}") return None, 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, 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) das.fps = fps # 设置 das.fps 为 load_media 返回的 fps if not is_video: tracking_method = "moge" print("Image input detected, 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: # 使用在CPU上运行的cotracker pred_tracks, pred_visibility = generate_tracking_cotracker(video_tensor) t, c, h, w = video_tensor.shape new_width = 518 new_height = round(h * (new_width / w) / 14) * 14 resize_transform = transforms.Resize((new_height, new_width), interpolation=Image.BICUBIC) video_vggt = resize_transform(video_tensor) # [T, C, H, W] if new_height > 518: start_y = (new_height - 518) // 2 video_vggt = video_vggt[:, :, start_y:start_y + 518, :] vggt_model = get_vggt_model() with torch.no_grad(): with torch.cuda.amp.autocast(dtype=das.dtype): video_vggt = video_vggt.unsqueeze(0) # [1, T, C, H, W] aggregated_tokens_list, ps_idx = vggt_model.aggregator(video_vggt.to(das.device)) extr, intr = pose_encoding_to_extri_intri(vggt_model.camera_head(aggregated_tokens_list)[-1], video_vggt.shape[-2:]) cam_motion.set_intr(intr) cam_motion.set_extr(extr) if camera_motion: poses = cam_motion.get_default_motion() # shape: [49, 4, 4] pred_tracks_world = cam_motion.s2w_vggt(pred_tracks, extr, intr) pred_tracks = cam_motion.w2s_vggt(pred_tracks_world, extr, intr, poses) # [T, N, 3] print("Camera motion applied") tracking_path, tracking_tensor = das.visualize_tracking_cotracker(pred_tracks, None) print('Export tracking video via cotracker') output_path = das.apply_tracking( video_tensor=video_tensor, fps=fps, # 使用 load_media 返回的 fps tracking_tensor=tracking_tensor, img_cond_tensor=repaint_img_tensor, prompt=prompt, checkpoint_path=DEFAULT_MODEL_PATH ) return tracking_path, output_path except Exception as e: import traceback print(f"Processing failed: {str(e)}\n{traceback.format_exc()}") return None, 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, None object_mask_path = save_uploaded_file(object_mask) if object_mask_path is None: print("Object mask not provided") return None, None das = get_das_pipeline() video_tensor, fps, is_video = load_media(input_image_path) das.fps = fps # 设置 das.fps 为 load_media 返回的 fps if not is_video: tracking_method = "moge" print("Image input detected, 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: # 使用在CPU上运行的cotracker pred_tracks, pred_visibility = generate_tracking_cotracker(video_tensor) t, c, h, w = video_tensor.shape new_width = 518 new_height = round(h * (new_width / w) / 14) * 14 resize_transform = transforms.Resize((new_height, new_width), interpolation=Image.BICUBIC) video_vggt = resize_transform(video_tensor) # [T, C, H, W] if new_height > 518: start_y = (new_height - 518) // 2 video_vggt = video_vggt[:, :, start_y:start_y + 518, :] vggt_model = get_vggt_model() with torch.no_grad(): with torch.cuda.amp.autocast(dtype=das.dtype): video_vggt = video_vggt.unsqueeze(0) # [1, T, C, H, W] aggregated_tokens_list, ps_idx = vggt_model.aggregator(video_vggt.to(das.device)) extr, intr = pose_encoding_to_extri_intri(vggt_model.camera_head(aggregated_tokens_list)[-1], video_vggt.shape[-2:]) pred_tracks = motion_generator.apply_motion( pred_tracks=pred_tracks.squeeze(), mask=mask, motion_type=object_motion, distance=50, num_frames=49, tracking_method="cotracker" ) print(f"Object motion '{object_motion}' applied using provided mask") tracking_path, tracking_tensor = das.visualize_tracking_cotracker(pred_tracks.unsqueeze(0), None) print('Export tracking video via cotracker') output_path = das.apply_tracking( video_tensor=video_tensor, fps=fps, # 使用 load_media 返回的 fps tracking_tensor=tracking_tensor, img_cond_tensor=repaint_img_tensor, prompt=prompt, checkpoint_path=DEFAULT_MODEL_PATH ) return tracking_path, output_path except Exception as e: import traceback print(f"Processing failed: {str(e)}\n{traceback.format_exc()}") return None, 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, None tracking_video_path = save_uploaded_file(tracking_video) if tracking_video_path is None: return None, None das = get_das_pipeline() video_tensor, fps, is_video = load_media(input_video_path) das.fps = fps # 设置 das.fps 为 load_media 返回的 fps 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=fps, # 使用 load_media 返回的 fps tracking_tensor=tracking_tensor, img_cond_tensor=repaint_img_tensor, prompt=prompt, checkpoint_path=DEFAULT_MODEL_PATH ) return tracking_video_path, output_path except Exception as e: import traceback print(f"Processing failed: {str(e)}\n{traceback.format_exc()}") return None, None def generate_tracking_cotracker(video_tensor, density=30): """在CPU上生成跟踪视频，只使用第一帧的深度信息，使用矩阵运算提高效率 参数: video_tensor (torch.Tensor): 输入视频张量 density (int): 跟踪点的密度 返回: tuple: (pred_tracks, pred_visibility) """ cotracker = torch.hub.load("facebookresearch/co-tracker", "cotracker3_offline").to("cpu") depth_preprocessor = DepthPreprocessor.from_pretrained("Intel/zoedepth-nyu-kitti").to("cpu") video = video_tensor.unsqueeze(0).to("cpu") # 只处理第一帧以获取深度图 print("estimating depth for first frame...") frame = (video_tensor[0].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8) depth = depth_preprocessor(Image.fromarray(frame))[0] depth_tensor = transforms.ToTensor()(depth) # [1, H, W] # 获取跟踪点和可见性 print("tracking on CPU...") pred_tracks, pred_visibility = cotracker(video, grid_size=density) # B T N 2, B T N 1 # 提取维度 B, T, N, _ = pred_tracks.shape H, W = depth_tensor.shape[1], depth_tensor.shape[2] # 创建带深度的输出张量 pred_tracks_with_depth = torch.zeros((B, T, N, 3), device="cpu") pred_tracks_with_depth[:, :, :, :2] = pred_tracks # 复制x,y坐标 # 使用矩阵运算一次性处理所有帧和点 # 重塑pred_tracks为[B*T*N, 2]以便于处理 flat_tracks = pred_tracks.reshape(-1, 2) # 将坐标限制在有效图像边界内 x_coords = flat_tracks[:, 0].clamp(0, W-1).long() y_coords = flat_tracks[:, 1].clamp(0, H-1).long() # 从第一帧的深度图获取所有点的深度值 depths = depth_tensor[0, y_coords, x_coords] # 重塑回原始形状并分配给输出张量 pred_tracks_with_depth[:, :, :, 2] = depths.reshape(B, T, N) del cotracker,depth_preprocessor # 将结果返回 return pred_tracks_with_depth.squeeze(0), pred_visibility.squeeze(0) # 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") tracking_video = gr.Video(label="Tracking 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=[tracking_video, 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=["moge", "cotracker"], value="cotracker" ) # 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=[tracking_video, 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=["moge", "cotracker"], value="cotracker" ) # 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=[tracking_video, 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=[tracking_video, 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)