import gradio as gr import os import numpy as np import cv2 import time import shutil from pathlib import Path from einops import rearrange from typing import Union try: import spaces except ImportError: def spaces(func): return func import torch import torchvision.transforms as T import logging from concurrent.futures import ThreadPoolExecutor import atexit import uuid import decord from models.SpaTrackV2.models.vggt4track.models.vggt_moe import VGGT4Track from models.SpaTrackV2.models.vggt4track.utils.load_fn import preprocess_image from models.SpaTrackV2.models.predictor import Predictor from models.SpaTrackV2.models.utils import get_points_on_a_grid from diffusers.utils import export_to_video, load_image from pipelines.wan_pipeline import WanImageToVideoTTMPipeline from pipelines.utils import compute_hw_from_area logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) MAX_FRAMES = 81 OUTPUT_FPS = 24 RENDER_WIDTH = 512 RENDER_HEIGHT = 384 WAN_MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers" CAMERA_MOVEMENTS = [ "static", "move_forward", "move_backward", "move_left", "move_right", "move_up", "move_down" ] thread_pool_executor = ThreadPoolExecutor(max_workers=2) def delete_later(path: Union[str, os.PathLike], delay: int = 600): def _delete(): try: if os.path.isfile(path): os.remove(path) elif os.path.isdir(path): shutil.rmtree(path) except Exception as e: logger.warning(f"Failed to delete {path}: {e}") def _wait_and_delete(): time.sleep(delay) _delete() thread_pool_executor.submit(_wait_and_delete) atexit.register(_delete) def create_user_temp_dir(): session_id = str(uuid.uuid4())[:8] temp_dir = os.path.join("temp_local", f"session_{session_id}") os.makedirs(temp_dir, exist_ok=True) delete_later(temp_dir, delay=600) return temp_dir print("🚀 Initializing tracking models...") vggt4track_model = VGGT4Track.from_pretrained( "Yuxihenry/SpatialTrackerV2_Front") vggt4track_model.eval() if not hasattr(vggt4track_model, 'infer'): vggt4track_model.infer = vggt4track_model.forward tracker_model = Predictor.from_pretrained("Yuxihenry/SpatialTrackerV2-Offline") tracker_model.eval() wan_pipeline = WanImageToVideoTTMPipeline.from_pretrained( WAN_MODEL_ID, torch_dtype=torch.bfloat16 ) wan_pipeline.vae.enable_tiling() wan_pipeline.vae.enable_slicing() print("✅ Tracking models loaded successfully!") gr.set_static_paths(paths=[Path.cwd().absolute()/"_viz"]) def generate_camera_trajectory(num_frames: int, movement_type: str, base_intrinsics: np.ndarray, scene_scale: float = 1.0) -> tuple: speed = scene_scale * 0.02 extrinsics = np.zeros((num_frames, 4, 4), dtype=np.float32) for t in range(num_frames): ext = np.eye(4, dtype=np.float32) if movement_type == "move_forward": ext[2, 3] = -speed * t elif movement_type == "move_backward": ext[2, 3] = speed * t elif movement_type == "move_left": ext[0, 3] = -speed * t elif movement_type == "move_right": ext[0, 3] = speed * t elif movement_type == "move_up": ext[1, 3] = -speed * t elif movement_type == "move_down": ext[1, 3] = speed * t extrinsics[t] = ext return extrinsics def render_from_pointcloud(rgb_frames, depth_frames, intrinsics, original_extrinsics, new_extrinsics, output_path, fps=24, generate_ttm_inputs=False): T, H, W, _ = rgb_frames.shape fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(output_path, fourcc, fps, (W, H)) motion_signal_path = mask_path = out_motion_signal = out_mask = None if generate_ttm_inputs: base_dir = os.path.dirname(output_path) motion_signal_path = os.path.join(base_dir, "motion_signal.mp4") mask_path = os.path.join(base_dir, "mask.mp4") out_motion_signal = cv2.VideoWriter( motion_signal_path, fourcc, fps, (W, H)) out_mask = cv2.VideoWriter(mask_path, fourcc, fps, (W, H)) u, v = np.meshgrid(np.arange(W), np.arange(H)) for t in range(T): rgb, depth, K = rgb_frames[t], depth_frames[t], intrinsics[t] orig_c2w = np.linalg.inv(original_extrinsics[t]) if t == 0: base_c2w = orig_c2w.copy() new_c2w = base_c2w @ new_extrinsics[t] new_w2c = np.linalg.inv(new_c2w) K_inv = np.linalg.inv(K) pixels = np.stack([u, v, np.ones_like(u)], axis=-1).reshape(-1, 3) rays_cam = (K_inv @ pixels.T).T points_cam = rays_cam * depth.reshape(-1, 1) points_world = (orig_c2w[:3, :3] @ points_cam.T).T + orig_c2w[:3, 3] points_new_cam = (new_w2c[:3, :3] @ points_world.T).T + new_w2c[:3, 3] points_proj = (K @ points_new_cam.T).T z = np.clip(points_proj[:, 2:3], 1e-6, None) uv_new = points_proj[:, :2] / z rendered = np.zeros((H, W, 3), dtype=np.uint8) z_buffer = np.full((H, W), np.inf, dtype=np.float32) colors, depths_new = rgb.reshape(-1, 3), points_new_cam[:, 2] for i in range(len(uv_new)): uu, vv = int(round(uv_new[i, 0])), int(round(uv_new[i, 1])) if 0 <= uu < W and 0 <= vv < H and depths_new[i] > 0: if depths_new[i] < z_buffer[vv, uu]: z_buffer[vv, uu] = depths_new[i] rendered[vv, uu] = colors[i] valid_mask = (rendered.sum(axis=-1) > 0).astype(np.uint8) * 255 motion_signal_frame = rendered.copy() hole_mask = (motion_signal_frame.sum(axis=-1) == 0).astype(np.uint8) if hole_mask.sum() > 0: kernel = np.ones((3, 3), np.uint8) for _ in range(10): # Iterative fill if hole_mask.sum() == 0: break dilated = cv2.dilate(motion_signal_frame, kernel) motion_signal_frame = np.where( hole_mask[:, :, None] > 0, dilated, motion_signal_frame) hole_mask = (motion_signal_frame.sum( axis=-1) == 0).astype(np.uint8) if generate_ttm_inputs: out_motion_signal.write(cv2.cvtColor( motion_signal_frame, cv2.COLOR_RGB2BGR)) out_mask.write(np.stack([valid_mask]*3, axis=-1)) out.write(cv2.cvtColor(motion_signal_frame, cv2.COLOR_RGB2BGR)) out.release() if generate_ttm_inputs: out_motion_signal.release() out_mask.release() return {'rendered': output_path, 'motion_signal': motion_signal_path, 'mask': mask_path} @spaces.GPU def run_spatial_tracker(video_tensor: torch.Tensor): """ GPU-intensive spatial tracking function. Args: video_tensor: Preprocessed video tensor (T, C, H, W) Returns: Dictionary containing tracking results """ global vggt4track_model global tracker_model global wan_pipeline video_input = preprocess_image(video_tensor)[None].cuda() vggt4track_model = vggt4track_model.to("cuda") with torch.no_grad(): with torch.amp.autocast('cuda', dtype=torch.bfloat16): predictions = vggt4track_model(video_input / 255) extrinsic = predictions["poses_pred"] intrinsic = predictions["intrs"] depth_map = predictions["points_map"][..., 2] depth_conf = predictions["unc_metric"] depth_tensor = depth_map.squeeze().cpu().numpy() extrs = extrinsic.squeeze().cpu().numpy() intrs = intrinsic.squeeze().cpu().numpy() video_tensor_gpu = video_input.squeeze() unc_metric = depth_conf.squeeze().cpu().numpy() > 0.5 tracker_model.spatrack.track_num = 512 tracker_model.to("cuda") frame_H, frame_W = video_tensor_gpu.shape[2:] grid_pts = get_points_on_a_grid(30, (frame_H, frame_W), device="cpu") query_xyt = torch.cat([torch.zeros_like(grid_pts[:, :, :1]), grid_pts], dim=2)[ 0].numpy() with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): ( c2w_traj, intrs_out, point_map, conf_depth, track3d_pred, track2d_pred, vis_pred, conf_pred, video_out ) = tracker_model.forward( video_tensor_gpu, depth=depth_tensor, intrs=intrs, extrs=extrs, queries=query_xyt, fps=1, full_point=False, iters_track=4, query_no_BA=True, fixed_cam=False, stage=1, unc_metric=unc_metric, support_frame=len(video_tensor_gpu)-1, replace_ratio=0.2 ) max_size = 384 h, w = video_out.shape[2:] scale = min(max_size / h, max_size / w) if scale < 1: new_h, new_w = int(h * scale), int(w * scale) video_out = T.Resize((new_h, new_w))(video_out) point_map = T.Resize((new_h, new_w))(point_map) conf_depth = T.Resize((new_h, new_w))(conf_depth) intrs_out[:, :2, :] = intrs_out[:, :2, :] * scale return { 'video_out': video_out.cpu(), 'point_map': point_map.cpu(), 'conf_depth': conf_depth.cpu(), 'intrs_out': intrs_out.cpu(), 'c2w_traj': c2w_traj.cpu(), } @spaces.GPU def run_wan_ttm_generation(prompt, tweak_index, tstrong_index, first_frame_path, motion_video_path, mask_video_path, progress=gr.Progress()): if not first_frame_path or not motion_video_path or not mask_video_path: return None, "❌ TTM Inputs missing. Please run 3D tracking first." progress(0, desc="Loading Wan TTM Pipeline...") import decord vr = decord.VideoReader(motion_video_path) actual_frame_count = len(vr) target_num_frames = ((actual_frame_count - 1) // 4) * 4 + 1 if target_num_frames < 5: return None, f"❌ Video too short. Only {actual_frame_count} frames tracked." logger.info(f"Setting Wan num_frames to {target_num_frames} based on tracking output.") progress(0.2, desc="Preparing inputs...") image = load_image(first_frame_path) negative_prompt = ( "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，" "低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，" "毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" ) wan_pipeline.to("cuda") max_area = 480 * 832 mod_value = wan_pipeline.vae_scale_factor_spatial * \ wan_pipeline.transformer.config.patch_size[1] height, width = compute_hw_from_area( image.height, image.width, max_area, mod_value) image = image.resize((width, height)) progress(0.4, desc=f"Generating {target_num_frames} frames (this may take a few minutes)...") generator = torch.Generator(device="cuda").manual_seed(0) with torch.inference_mode(): result = wan_pipeline( image=image, prompt=prompt, negative_prompt=negative_prompt, height=height, width=width, num_frames=target_num_frames, guidance_scale=3.5, num_inference_steps=50, generator=generator, motion_signal_video_path=motion_video_path, motion_signal_mask_path=mask_video_path, tweak_index=int(tweak_index), tstrong_index=int(tstrong_index), ) output_path = os.path.join(os.path.dirname( first_frame_path), "wan_ttm_output.mp4") export_to_video(result.frames[0], output_path, fps=16) return output_path, f"✅ TTM Video ({target_num_frames} frames) generated successfully!" # --- MODIFIED PROCESS VIDEO TO RETURN FILE PATHS --- def process_video(video_path, camera_movement, generate_ttm=True, progress=gr.Progress()): if video_path is None: return None, None, None, None, "❌ Please upload a video first" progress(0, desc="Initializing...") temp_dir = create_user_temp_dir() out_dir = os.path.join(temp_dir, "results") os.makedirs(out_dir, exist_ok=True) try: progress(0.1, desc="Loading video...") video_reader = decord.VideoReader(video_path) video_tensor = torch.from_numpy(video_reader.get_batch( range(len(video_reader))).asnumpy()).permute(0, 3, 1, 2).float() video_tensor = video_tensor[::max( 1, len(video_tensor)//MAX_FRAMES)][:MAX_FRAMES] h, w = video_tensor.shape[2:] scale = 336 / min(h, w) if scale < 1: video_tensor = T.Resize( (int(h*scale)//2*2, int(w*scale)//2*2))(video_tensor) progress(0.4, desc="Running 3D tracking...") tracking_results = run_spatial_tracker(video_tensor) rgb_frames = rearrange( tracking_results['video_out'].numpy(), "T C H W -> T H W C").astype(np.uint8) depth_frames = tracking_results['point_map'][:, 2].numpy() depth_frames[tracking_results['conf_depth'].numpy() < 0.5] = 0 scene_scale = np.median(depth_frames[depth_frames > 0]) if np.any( depth_frames > 0) else 1.0 new_exts = generate_camera_trajectory(len( rgb_frames), camera_movement, tracking_results['intrs_out'].numpy(), scene_scale) progress(0.8, desc="Rendering viewpoint...") output_video_path = os.path.join(out_dir, "rendered_video.mp4") render_results = render_from_pointcloud(rgb_frames, depth_frames, tracking_results['intrs_out'].numpy(), torch.inverse( tracking_results['c2w_traj']).numpy(), new_exts, output_video_path, fps=OUTPUT_FPS, generate_ttm_inputs=generate_ttm) first_frame_path = os.path.join(out_dir, "first_frame.png") cv2.imwrite(first_frame_path, cv2.cvtColor( rgb_frames[0], cv2.COLOR_RGB2BGR)) status_msg = f"✅ 3D results ready! You can now use the prompt below to generate a high-quality TTM video." return render_results['rendered'], render_results['motion_signal'], render_results['mask'], first_frame_path, status_msg except Exception as e: logger.error(f"Error: {e}") return None, None, None, None, f"❌ Error: {str(e)}" # --- GRADIO INTERFACE --- with gr.Blocks(theme=gr.themes.Soft(), title="🎬 TTM Wan Video Generator") as demo: gr.Markdown("# 🎬 Video to Point Cloud & TTM Wan Generator") gr.Markdown( "Transform standard videos into 3D-aware motion signals for Time-to-Move (TTM) generation.") first_frame_file = gr.State("") motion_signal_file = gr.State("") mask_file = gr.State("") with gr.Row(): with gr.Column(scale=1): gr.Markdown("### 1. Tracking & Viewpoint") video_input = gr.Video(label="Upload Video") camera_movement = gr.Dropdown( choices=CAMERA_MOVEMENTS, value="static", label="Camera Movement" ) generate_btn = gr.Button( "🚀 1. Run Spatial Tracker", variant="primary") output_video = gr.Video(label="Point Cloud Render (Draft)") status_text = gr.Markdown("Ready...") with gr.Column(scale=1): gr.Markdown("### 2. Time-to-Move (Wan 2.2)") ttm_prompt = gr.Textbox( label="Prompt", placeholder="Describe the scene (e.g., 'A monkey walking in the forest, high quality')" ) with gr.Row(): tweak_idx = gr.Number( label="Tweak Index", value=3, precision=0) tstrong_idx = gr.Number( label="Tstrong Index", value=6, precision=0) wan_generate_btn = gr.Button( "✨ 2. Generate TTM Video (Wan)", variant="secondary") wan_output_video = gr.Video(label="Final High-Quality TTM Video") wan_status = gr.Markdown("Awaiting 3D inputs...") with gr.Accordion("Debug: TTM Intermediate Inputs", open=False): with gr.Row(): motion_signal_output = gr.Video(label="motion_signal.mp4") mask_output = gr.Video(label="mask.mp4") first_frame_output = gr.Image( label="first_frame.png", type="filepath") generate_btn.click( fn=process_video, inputs=[video_input, camera_movement], outputs=[ output_video, motion_signal_output, mask_output, first_frame_output, status_text ] ).then( fn=lambda a, b, c, d, e: (b, c, d), inputs=[ output_video, motion_signal_output, mask_output, first_frame_output, status_text ], outputs=[motion_signal_file, mask_file, first_frame_file] ) wan_generate_btn.click( fn=run_wan_ttm_generation, inputs=[ ttm_prompt, tweak_idx, tstrong_idx, first_frame_file, motion_signal_file, mask_file ], outputs=[wan_output_video, wan_status] ) if __name__ == "__main__": demo.launch(share=False)