add wan ttm

app.py

@@ -19,6 +19,14 @@ from concurrent.futures import ThreadPoolExecutor
 import atexit
 import uuid
 import decord
+from PIL import Image
+
+try:
+    from pipelines.wan_pipeline import WanImageToVideoTTMPipeline
+    from pipelines.utils import compute_hw_from_area, validate_inputs
+    from diffusers.utils import export_to_video, load_image
+except ImportError:
+    print("Warning: TTM pipelines not found. Ensure the /pipelines folder is in your path.")
 
 from models.SpaTrackV2.models.vggt4track.models.vggt_moe import VGGT4Track
 from models.SpaTrackV2.models.vggt4track.utils.load_fn import preprocess_image
@@ -30,548 +38,206 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # Constants
-MAX_FRAMES = 80
-OUTPUT_FPS = 24
-RENDER_WIDTH = 512
-RENDER_HEIGHT = 384
-
-# Camera movement types
-CAMERA_MOVEMENTS = [
-    "static",
-    "move_forward",
-    "move_backward",
-    "move_left",
-    "move_right",
-    "move_up",
-    "move_down"
-]
-
-# Thread pool for delayed deletion
-thread_pool_executor = ThreadPoolExecutor(max_workers=2)
-
-def delete_later(path: Union[str, os.PathLike], delay: int = 600):
-    """Delete file or directory after specified delay"""
-    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}")
+MAX_FRAMES = 81
+OUTPUT_FPS = 16
+WAN_MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
+DTYPE = torch.bfloat16
+
+# --- Global Model Initialization ---
+print("🚀 Initializing models...")
+vggt4track_model = VGGT4Track.from_pretrained("Yuxihenry/SpatialTrackerV2_Front")
+vggt4track_model.eval().to("cuda")
 
+tracker_model = Predictor.from_pretrained("Yuxihenry/SpatialTrackerV2-Offline")
+tracker_model.eval()
+
+# Lazy loading for Wan to save VRAM initially
+wan_pipe = None
+
+def get_wan_pipeline():
+    global wan_pipe
+    if wan_pipe is None:
+        print("🚀 Initializing Wan 2.2 TTM Pipeline...")
+        wan_pipe = WanImageToVideoTTMPipeline.from_pretrained(WAN_MODEL_ID, torch_dtype=DTYPE)
+        wan_pipe.vae.enable_tiling()
+        wan_pipe.vae.enable_slicing()
+        wan_pipe.to("cuda")
+    return wan_pipe
+
+# --- Utility Functions ---
+def delete_later(path, delay=600):
     def _wait_and_delete():
         time.sleep(delay)
-        _delete()
-
-    thread_pool_executor.submit(_wait_and_delete)
-    atexit.register(_delete)
+        try:
+            if os.path.isfile(path): os.remove(path)
+            elif os.path.isdir(path): shutil.rmtree(path)
+        except: pass
+    ThreadPoolExecutor(max_workers=1).submit(_wait_and_delete)
 
 def create_user_temp_dir():
-    """Create a unique temporary directory for each user session"""
     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)
+    delete_later(temp_dir)
     return temp_dir
 
-# Global model initialization
-print("🚀 Initializing models...")
-vggt4track_model = VGGT4Track.from_pretrained("Yuxihenry/SpatialTrackerV2_Front")
-vggt4track_model.eval()
-vggt4track_model = vggt4track_model.to("cuda")
-
-tracker_model = Predictor.from_pretrained("Yuxihenry/SpatialTrackerV2-Offline")
-tracker_model.eval()
-print("✅ 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:
-    """
-    Generate camera extrinsics for different movement types.
-
-    Returns:
-        extrinsics: (T, 4, 4) camera-to-world matrices
-    """
-    # Movement speed (adjust based on scene scale)
+def generate_camera_trajectory(num_frames, movement_type, base_intrinsics, scene_scale=1.0):
     speed = scene_scale * 0.02
-
     extrinsics = np.zeros((num_frames, 4, 4), dtype=np.float32)
-
     for t in range(num_frames):
-        # Start with identity matrix
         ext = np.eye(4, dtype=np.float32)
-
-        progress = t / max(num_frames - 1, 1)
-
-        if movement_type == "static":
-            pass  # Keep identity
-        elif movement_type == "move_forward":
-            ext[2, 3] = -speed * t  # Move along -Z (forward in OpenGL convention)
-        elif movement_type == "move_backward":
-            ext[2, 3] = speed * t  # Move along +Z
-        elif movement_type == "move_left":
-            ext[0, 3] = -speed * t  # Move along -X
-        elif movement_type == "move_right":
-            ext[0, 3] = speed * t  # Move along +X
-        elif movement_type == "move_up":
-            ext[1, 3] = -speed * t  # Move along -Y (up in OpenGL)
-        elif movement_type == "move_down":
-            ext[1, 3] = speed * t  # Move along +Y
-
+        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: np.ndarray,
-                           depth_frames: np.ndarray,
-                           intrinsics: np.ndarray,
-                           original_extrinsics: np.ndarray,
-                           new_extrinsics: np.ndarray,
-                           output_path: str,
-                           fps: int = 24,
-                           generate_ttm_inputs: bool = False) -> dict:
-    """
-    Render video from point cloud with new camera trajectory.
-
-    Args:
-        rgb_frames: (T, H, W, 3) RGB frames
-        depth_frames: (T, H, W) depth maps
-        intrinsics: (T, 3, 3) camera intrinsics
-        original_extrinsics: (T, 4, 4) original camera extrinsics (world-to-camera)
-        new_extrinsics: (T, 4, 4) new camera extrinsics for rendering
-        output_path: path to save rendered video
-        fps: output video fps
-        generate_ttm_inputs: if True, also generate motion_signal.mp4 and mask.mp4 for TTM
-
-    Returns:
-        dict with paths: {'rendered': path, 'motion_signal': path or None, 'mask': path or None}
-    """
+def render_from_pointcloud(rgb_frames, depth_frames, intrinsics, original_extrinsics, new_extrinsics, output_path, generate_ttm_inputs=True):
     T, H, W, _ = rgb_frames.shape
+    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), OUTPUT_FPS, (W, H))
 
-    # Setup video writers
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, fps, (W, H))
-
-    # TTM outputs: motion_signal (warped with NN inpainting) and mask (valid pixels before inpainting)
-    motion_signal_path = None
-    mask_path = None
-    out_motion_signal = None
-    out_mask = None
+    motion_signal_path = os.path.join(os.path.dirname(output_path), "motion_signal.mp4")
+    mask_path = os.path.join(os.path.dirname(output_path), "mask.mp4")
+    out_motion = cv2.VideoWriter(motion_signal_path, cv2.VideoWriter_fourcc(*'mp4v'), OUTPUT_FPS, (W, H))
+    out_mask = cv2.VideoWriter(mask_path, cv2.VideoWriter_fourcc(*'mp4v'), OUTPUT_FPS, (W, H))
 
-    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))
-
-    # Create meshgrid for pixel coordinates
     u, v = np.meshgrid(np.arange(W), np.arange(H))
-    ones = np.ones_like(u)
-
     for t in range(T):
-        # Get current frame data
-        rgb = rgb_frames[t]
-        depth = depth_frames[t]
-        K = intrinsics[t]
-
-        # Original camera pose (camera-to-world)
         orig_c2w = np.linalg.inv(original_extrinsics[t])
-
-        # New camera pose (camera-to-world for the new viewpoint)
-        # Apply the new extrinsics relative to the first frame
-        if t == 0:
-            base_c2w = orig_c2w.copy()
-
-        # New camera is: base_c2w @ new_extrinsics[t]
+        if t == 0: base_c2w = orig_c2w.copy()
         new_c2w = base_c2w @ new_extrinsics[t]
         new_w2c = np.linalg.inv(new_c2w)
 
-        # Unproject pixels to 3D points
-        K_inv = np.linalg.inv(K)
-
-        # Pixel coordinates to normalized camera coordinates
-        pixels = np.stack([u, v, ones], axis=-1).reshape(-1, 3)  # (H*W, 3)
-        rays_cam = (K_inv @ pixels.T).T  # (H*W, 3)
-
-        # Scale by depth to get 3D points in original camera frame
-        depth_flat = depth.reshape(-1, 1)
-        points_cam = rays_cam * depth_flat  # (H*W, 3)
-
-        # Transform to world coordinates
+        K_inv = np.linalg.inv(intrinsics[t])
+        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_frames[t].reshape(-1, 1)
         points_world = (orig_c2w[:3, :3] @ points_cam.T).T + orig_c2w[:3, 3]
-
-        # Transform to new camera coordinates
         points_new_cam = (new_w2c[:3, :3] @ points_world.T).T + new_w2c[:3, 3]
+        points_proj = (intrinsics[t] @ points_new_cam.T).T
 
-        # Project to new image
-        points_proj = (K @ points_new_cam.T).T
-
-        # Get pixel coordinates
-        z = points_proj[:, 2:3]
-        z = np.clip(z, 1e-6, None)  # Avoid division by zero
-        uv_new = points_proj[:, :2] / z
-
-        # Create output image using forward warping with z-buffer
+        uv_new = points_proj[:, :2] / np.clip(points_proj[:, 2:3], 1e-6, None)
         rendered = np.zeros((H, W, 3), dtype=np.uint8)
-        z_buffer = np.full((H, W), np.inf, dtype=np.float32)
-
-        colors = rgb.reshape(-1, 3)
-        depths_new = points_new_cam[:, 2]
+        z_buf = np.full((H, W), np.inf)
 
         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]
+            if 0 <= uu < W and 0 <= vv < H and points_new_cam[i, 2] > 0:
+                if points_new_cam[i, 2] < z_buf[vv, uu]:
+                    z_buf[vv, uu] = points_new_cam[i, 2]
+                    rendered[vv, uu] = rgb_frames[t].reshape(-1, 3)[i]
 
-        # Create valid pixel mask BEFORE hole filling (for TTM)
-        # Valid pixels are those that received projected colors
         valid_mask = (rendered.sum(axis=-1) > 0).astype(np.uint8) * 255
 
-        # Nearest-neighbor hole filling using dilation
-        # This is the inpainting method described in TTM: "Missing regions are inpainted by nearest-neighbor color assignment"
-        motion_signal_frame = rendered.copy()
-        hole_mask = (motion_signal_frame.sum(axis=-1) == 0).astype(np.uint8)
+        # Hole filling for motion signal
+        motion_frame = rendered.copy()
+        hole_mask = (motion_frame.sum(axis=-1) == 0).astype(np.uint8)
         if hole_mask.sum() > 0:
-            kernel = np.ones((3, 3), np.uint8)
-            # Iteratively dilate to fill holes with nearest neighbor colors
-            max_iterations = max(H, W)  # Ensure all holes can be filled
-            for _ in range(max_iterations):
-                if hole_mask.sum() == 0:
-                    break
-                dilated = cv2.dilate(motion_signal_frame, kernel, iterations=1)
-                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)
-
-        # Write TTM outputs if enabled
-        if generate_ttm_inputs:
-            # Motion signal: warped frame with NN inpainting
-            motion_signal_bgr = cv2.cvtColor(motion_signal_frame, cv2.COLOR_RGB2BGR)
-            out_motion_signal.write(motion_signal_bgr)
-
-            # Mask: binary mask of valid (projected) pixels - white where valid, black where holes
-            mask_frame = np.stack([valid_mask, valid_mask, valid_mask], axis=-1)
-            out_mask.write(mask_frame)
-
-        # For the rendered output, use the same inpainted result
-        rendered_bgr = cv2.cvtColor(motion_signal_frame, cv2.COLOR_RGB2BGR)
-        out.write(rendered_bgr)
-
-    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
-    """
-    # Run VGGT to get depth and camera poses
-    video_input = preprocess_image(video_tensor)[None].cuda()
+            for _ in range(10): # Iterative dilation for NN inpainting
+                dilated = cv2.dilate(motion_frame, np.ones((3,3), np.uint8))
+                motion_frame = np.where(hole_mask[:, :, None] > 0, dilated, motion_frame)
+                hole_mask = (motion_frame.sum(axis=-1) == 0).astype(np.uint8)
+                if hole_mask.sum() == 0: break
+
+        out_motion.write(cv2.cvtColor(motion_frame, cv2.COLOR_RGB2BGR))
+        out_mask.write(cv2.merge([valid_mask, valid_mask, valid_mask]))
+        out.write(cv2.cvtColor(motion_frame, cv2.COLOR_RGB2BGR))
+
+    out.release(); out_motion.release(); out_mask.release()
+    return {'rendered': output_path, 'motion_signal': motion_signal_path, 'mask': mask_path}
+
+# --- Main Processing Logic ---
+def run_ttm_wan_inference(image_path, motion_path, mask_path, prompt, tweak_idx, tstrong_idx, guidance_scale, seed=0):
+    pipe = get_wan_pipeline()
+    image = load_image(image_path)
+    max_area = 480 * 832
+    mod_val = pipe.vae_scale_factor_spatial * pipe.transformer.config.patch_size[1]
+    h, w = compute_hw_from_area(image.height, image.width, max_area, mod_val)
+    image = image.resize((w, h))
+
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    with torch.inference_mode():
+        result = pipe(
+            image=image, prompt=prompt, height=h, width=w, num_frames=81,
+            guidance_scale=guidance_scale, num_inference_steps=50, generator=generator,
+            motion_signal_video_path=motion_path, motion_signal_mask_path=mask_path,
+            tweak_index=tweak_idx, tstrong_index=tstrong_idx, negative_prompt="blurry, static, low quality"
+        )
+    return result.frames[0]
+
+def process_video_full_pipeline(video_path, camera_movement, prompt, tweak_idx, tstrong_idx, guidance_scale, progress=gr.Progress()):
+    if not video_path or not prompt: return [None]*5 + ["❌ Missing video or prompt"]
 
+    temp_dir = create_user_temp_dir()
+    res_dir = os.path.join(temp_dir, "results"); os.makedirs(res_dir, exist_ok=True)
+
+    # 1. Spatial Tracking
+    progress(0.1, desc="3D Analysis...")
+    vr = decord.VideoReader(video_path)
+    vt = torch.from_numpy(vr.get_batch(range(len(vr))).asnumpy()).permute(0,3,1,2).float()
+    vt = vt[::max(1, len(vt)//MAX_FRAMES)][:MAX_FRAMES]
+
+    # Preprocess for VGGT
+    v_in = preprocess_image(vt)[None].cuda()
     with torch.no_grad():
-        with torch.cuda.amp.autocast(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
-
-    # Setup tracker
-    tracker_model.spatrack.track_num = 512
+        preds = vggt4track_model(v_in / 255)
+
+    # Tracker
     tracker_model.to("cuda")
+    grid = get_points_on_a_grid(30, v_in.shape[3:], device="cpu")
+    queries = torch.cat([torch.zeros_like(grid[:,:,:1]), grid], dim=2)[0].numpy()
 
-    # Get grid points for tracking
-    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()
-
-    # Run tracker
-    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
-        )
+    c2w, intrs, p_map, c_depth, _, _, _, _, v_out = tracker_model.forward(
+        v_in.squeeze(), depth=preds["points_map"][...,2].squeeze().cpu().numpy(),
+        intrs=preds["intrs"].squeeze().cpu().numpy(), extrs=preds["poses_pred"].squeeze().cpu().numpy(),
+        queries=queries, fps=1, iters_track=4, fixed_cam=False
+    )
 
-    # Resize outputs for rendering
-    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
-
-    # Move results to CPU and return
-    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(),
-    }
-
-
-def process_video(video_path: str, camera_movement: str, generate_ttm: bool = True, progress=gr.Progress()):
-    """Main processing function
-
-    Args:
-        video_path: Path to input video
-        camera_movement: Type of camera movement
-        generate_ttm: If True, generate TTM-compatible outputs (motion_signal.mp4, mask.mp4, first_frame.png)
-        progress: Gradio progress tracker
-    """
-    if video_path is None:
-        return None, None, None, None, "❌ Please upload a video first"
-
-    progress(0, desc="Initializing...")
-
-    # Create temp directory
-    temp_dir = create_user_temp_dir()
-    out_dir = os.path.join(temp_dir, "results")
-    os.makedirs(out_dir, exist_ok=True)
-
-    try:
-        # Load video
-        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()
-
-        # Subsample frames if too many
-        fps_skip = max(1, len(video_tensor) // MAX_FRAMES)
-        video_tensor = video_tensor[::fps_skip][:MAX_FRAMES]
-
-        # Resize to have minimum side 336
-        h, w = video_tensor.shape[2:]
-        scale = 336 / min(h, w)
-        if scale < 1:
-            new_h, new_w = int(h * scale) // 2 * 2, int(w * scale) // 2 * 2
-            video_tensor = T.Resize((new_h, new_w))(video_tensor)
-
-        progress(0.2, desc="Estimating depth and camera poses...")
-
-        # Run GPU-intensive spatial tracking
-        progress(0.4, desc="Running 3D tracking...")
-        tracking_results = run_spatial_tracker(video_tensor)
-
-        progress(0.6, desc="Preparing point cloud...")
-
-        # Extract results from tracking
-        video_out = tracking_results['video_out']
-        point_map = tracking_results['point_map']
-        conf_depth = tracking_results['conf_depth']
-        intrs_out = tracking_results['intrs_out']
-        c2w_traj = tracking_results['c2w_traj']
-
-        # Get RGB frames and depth
-        rgb_frames = rearrange(video_out.numpy(), "T C H W -> T H W C").astype(np.uint8)
-        depth_frames = point_map[:, 2].numpy()
-        depth_conf_np = conf_depth.numpy()
-
-        # Mask out unreliable depth
-        depth_frames[depth_conf_np < 0.5] = 0
-
-        # Get camera parameters
-        intrs_np = intrs_out.numpy()
-        extrs_np = torch.inverse(c2w_traj).numpy()  # world-to-camera
-
-        progress(0.7, desc=f"Generating {camera_movement} camera trajectory...")
-
-        # Calculate scene scale from depth
-        valid_depth = depth_frames[depth_frames > 0]
-        scene_scale = np.median(valid_depth) if len(valid_depth) > 0 else 1.0
-
-        # Generate new camera trajectory
-        num_frames = len(rgb_frames)
-        new_extrinsics = generate_camera_trajectory(
-            num_frames, camera_movement, intrs_np, scene_scale
-        )
+    # 2. Rendering
+    progress(0.6, desc="Rendering Point Cloud...")
+    rgb = rearrange(v_out.cpu().numpy(), "T C H W -> T H W C").astype(np.uint8)
+    depth = p_map[0, 2].cpu().numpy() # Simplified for single view context
+    new_ext = generate_camera_trajectory(len(rgb), camera_movement, intrs.cpu().numpy(), np.median(depth[depth>0]))
 
-        progress(0.8, desc="Rendering video from new viewpoint...")
+    rend_path = os.path.join(res_dir, "warp.mp4")
+    rend_res = render_from_pointcloud(rgb, p_map[:,2].cpu().numpy(), intrs.cpu().numpy(), torch.inverse(c2w).cpu().numpy(), new_ext, rend_path)
 
-        # Render video (CPU-based, no GPU needed)
-        output_video_path = os.path.join(out_dir, "rendered_video.mp4")
-        render_results = render_from_pointcloud(
-            rgb_frames, depth_frames, intrs_np, extrs_np,
-            new_extrinsics, output_video_path, fps=OUTPUT_FPS,
-            generate_ttm_inputs=generate_ttm
-        )
+    first_frame_path = os.path.join(res_dir, "first.png")
+    cv2.imwrite(first_frame_path, cv2.cvtColor(rgb[0], cv2.COLOR_RGB2BGR))
 
-        # Save first frame for TTM
-        first_frame_path = None
-        motion_signal_path = None
-        mask_path = None
-
-        if generate_ttm:
-            first_frame_path = os.path.join(out_dir, "first_frame.png")
-            # Save original first frame (before warping) as PNG
-            first_frame_rgb = rgb_frames[0]
-            first_frame_bgr = cv2.cvtColor(first_frame_rgb, cv2.COLOR_RGB2BGR)
-            cv2.imwrite(first_frame_path, first_frame_bgr)
-
-            motion_signal_path = render_results['motion_signal']
-            mask_path = render_results['mask']
-
-        progress(1.0, desc="Done!")
-
-        status_msg = f"✅ Video rendered successfully with '{camera_movement}' camera movement!"
-        if generate_ttm:
-            status_msg += "\n\n📁 **TTM outputs generated:**\n"
-            status_msg += f"- `first_frame.png`: Input frame for TTM\n"
-            status_msg += f"- `motion_signal.mp4`: Warped video with NN inpainting\n"
-            status_msg += f"- `mask.mp4`: Valid pixel mask (white=valid, black=hole)"
-
-        return render_results['rendered'], motion_signal_path, mask_path, first_frame_path, status_msg
-
-    except Exception as e:
-        logger.error(f"Error processing video: {e}")
-        import traceback
-        traceback.print_exc()
-        return None, None, None, None, f"❌ Error: {str(e)}"
-
-
-# Create Gradio interface
-print("🎨 Creating Gradio interface...")
-
-with gr.Blocks(
-    theme=gr.themes.Soft(),
-    title="🎬 Video to Point Cloud Renderer",
-    css="""
-    .gradio-container {
-        max-width: 1200px !important;
-        margin: auto !important;
-    }
-    """
-) as demo:
-    gr.Markdown("""
-    # 🎬 Video to Point Cloud Renderer (TTM Compatible)
-
-    Upload a video to generate a 3D point cloud and render it from a new camera perspective.
-    Generates outputs compatible with **Time-to-Move (TTM)** motion-controlled video generation.
-
-    **How it works:**
-    1. Upload a video
-    2. Select a camera movement type
-    3. Click "Generate" to create the rendered video and TTM inputs
-
-    **TTM Inputs:**
-    - `first_frame.png`: The first frame of the original video
-    - `motion_signal.mp4`: Warped video with nearest-neighbor inpainting
-    - `mask.mp4`: Binary mask showing valid projected pixels (white) vs holes (black)
-    """)
+    # 3. Wan TTM Inference
+    progress(0.8, desc="Wan 2.2 Realistic Generation...")
+    wan_video_path = os.path.join(res_dir, "final_wan.mp4")
+    wan_frames = run_ttm_wan_inference(first_frame_path, rend_res['motion_signal'], rend_res['mask'], prompt, tweak_idx, tstrong_idx, guidance_scale)
+    export_to_video(wan_frames, wan_video_path, fps=16)
 
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### 📥 Input")
-            video_input = gr.Video(
-                label="Upload Video",
-                format="mp4",
-                height=300
-            )
-
-            camera_movement = gr.Dropdown(
-                choices=CAMERA_MOVEMENTS,
-                value="static",
-                label="🎥 Camera Movement",
-                info="Select how the camera should move in the rendered video"
-            )
-
-            generate_ttm = gr.Checkbox(
-                label="🎯 Generate TTM Inputs",
-                value=True,
-                info="Generate motion_signal.mp4 and mask.mp4 for Time-to-Move"
-            )
-
-            generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")
-
-        with gr.Column(scale=1):
-            gr.Markdown("### 📤 Rendered Output")
-            output_video = gr.Video(
-                label="Rendered Video",
-                height=250
-            )
-            first_frame_output = gr.Image(
-                label="First Frame (first_frame.png)",
-                height=150
-            )
+    return rend_path, wan_video_path, rend_res['motion_signal'], rend_res['mask'], first_frame_path, "✅ Generated successfully!"
 
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### 🎯 TTM: Motion Signal")
-            motion_signal_output = gr.Video(
-                label="Motion Signal Video (motion_signal.mp4)",
-                height=250
-            )
-        with gr.Column(scale=1):
-            gr.Markdown("### 🎭 TTM: Mask")
-            mask_output = gr.Video(
-                label="Mask Video (mask.mp4)",
-                height=250
-            )
-
-    status_text = gr.Markdown("Ready to process...")
-
-    # Event handlers
-    generate_btn.click(
-        fn=process_video,
-        inputs=[video_input, camera_movement, generate_ttm],
-        outputs=[output_video, motion_signal_output, mask_output, first_frame_output, status_text]
-    )
+# --- Gradio UI ---
+with gr.Blocks(theme=gr.themes.Soft(), title="Wan 2.2 TTM Video Generator") as demo:
+    gr.Markdown("# 🎬 Time-to-Move (TTM) with Wan 2.2")
 
-    # Examples
-    gr.Markdown("### 📁 Examples")
-    if os.path.exists("./examples"):
-        example_videos = [f for f in os.listdir("./examples") if f.endswith(".mp4")][:4]
-        if example_videos:
-            gr.Examples(
-                examples=[[f"./examples/{v}", "move_forward", True] for v in example_videos],
-                inputs=[video_input, camera_movement, generate_ttm],
-                outputs=[output_video, motion_signal_output, mask_output, first_frame_output, status_text],
-                fn=process_video,
-                cache_examples=False
-            )
-
-# Launch
-if __name__ == "__main__":
-    demo.launch(share=False)
+    with gr.Row():
+        with gr.Column():
+            v_in = gr.Video(label="Source Video")
+            p_in = gr.Textbox(label="Prompt", placeholder="Describe the action...")
+            c_in = gr.Dropdown(choices=["move_forward", "move_backward", "move_left", "move_right", "move_up", "move_down", "static"], value="move_forward", label="Camera Movement")
+            with gr.Accordion("TTM Settings", open=False):
+                twk = gr.Slider(0, 15, value=3, label="Tweak Index")
+                strng = gr.Slider(0, 20, value=7, label="Tstrong Index")
+                cfg = gr.Slider(1, 10, value=5.0, label="CFG Scale")
+            btn = gr.Button("Generate Realistic Video", variant="primary")
+
+        with gr.Column():
+            v_final = gr.Video(label="Final Realistic Result")
+            v_warp = gr.Video(label="Point Cloud Warp (Guide)")
+            with gr.Row():
+                v_msig = gr.Video(label="Motion Signal")
+                v_mask = gr.Video(label="Mask")
+
+    btn.click(process_video_full_pipeline, [v_in, c_in, p_in, twk, strng, cfg], [v_warp, v_final, v_msig, v_mask, gr.Image(visible=False), gr.Markdown()])
+
+demo.launch()