fix

app.py

@@ -19,225 +19,401 @@ 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
 from models.SpaTrackV2.models.predictor import Predictor
 from models.SpaTrackV2.models.utils import get_points_on_a_grid
 
+
+# --- TTM SPECIFIC IMPORTS ---
+from diffusers.utils import export_to_video, load_image
+# Note: Ensure pipelines/wan_pipeline.py and pipelines/utils.py are in your working directory
+from pipelines.wan_pipeline import WanImageToVideoTTMPipeline
+from pipelines.utils import compute_hw_from_area, validate_inputs
+
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # Constants
-MAX_FRAMES = 81
-OUTPUT_FPS = 16
+MAX_FRAMES = 80
+OUTPUT_FPS = 24
+RENDER_WIDTH = 512
+RENDER_HEIGHT = 384
 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")
+# Camera movement types
+CAMERA_MOVEMENTS = [
+    "static",
+    "move_forward",
+    "move_backward",
+    "move_left",
+    "move_right",
+    "move_up",
+    "move_down"
+]
 
-tracker_model = Predictor.from_pretrained("Yuxihenry/SpatialTrackerV2-Offline")
-tracker_model.eval()
+# Thread pool for delayed deletion
+thread_pool_executor = ThreadPoolExecutor(max_workers=2)
 
-# 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 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)
-        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)
+        _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)
+    delete_later(temp_dir, delay=600)
     return temp_dir
 
-def generate_camera_trajectory(num_frames, movement_type, base_intrinsics, scene_scale=1.0):
+
+# Global model initialization for Spatial Tracker
+print("🚀 Initializing tracking 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()
+
+# Lazy loading for Wan to save memory until needed
+wan_pipeline = None
+
+
+def get_wan_pipeline():
+    global wan_pipeline
+    if wan_pipeline is None:
+        print("🚀 Loading Wan TTM Pipeline (14B)...")
+        wan_pipeline = WanImageToVideoTTMPipeline.from_pretrained(
+            WAN_MODEL_ID,
+            torch_dtype=torch.bfloat16
+        )
+        wan_pipeline.vae.enable_tiling()
+        wan_pipeline.vae.enable_slicing()
+        wan_pipeline.to("cuda")
+    return wan_pipeline
+
+
+print("✅ Tracking models loaded successfully!")
+
+gr.set_static_paths(paths=[Path.cwd().absolute()/"_viz"])
+
+# --- YOUR ORIGINAL FUNCTIONS (generate_camera_trajectory, render_from_pointcloud, run_spatial_tracker) ---
+# [Keeping these exactly as provided in your snippet]
+
+
+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
+        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, generate_ttm_inputs=True):
-    T, H, W, _ = rgb_frames.shape
-    out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), OUTPUT_FPS, (W, H))
 
-    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))
+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()
+        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(intrinsics[t])
+        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_frames[t].reshape(-1, 1)
+        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 = (intrinsics[t] @ points_new_cam.T).T
-
-        uv_new = points_proj[:, :2] / np.clip(points_proj[:, 2:3], 1e-6, None)
+        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_buf = np.full((H, W), np.inf)
+        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 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]
+            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
-
-        # Hole filling for motion signal
-        motion_frame = rendered.copy()
-        hole_mask = (motion_frame.sum(axis=-1) == 0).astype(np.uint8)
+        motion_signal_frame = rendered.copy()
+        hole_mask = (motion_signal_frame.sum(axis=-1) == 0).astype(np.uint8)
         if hole_mask.sum() > 0:
-            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()
+            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}
 
-# --- 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"
+@spaces.GPU
+def run_spatial_tracker(video_tensor):
+    video_input = preprocess_image(video_tensor)[None].cuda()
+    with torch.no_grad():
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+            predictions = vggt4track_model(video_input / 255)
+            extrinsic, intrinsic = predictions["poses_pred"], predictions["intrs"]
+            depth_map, depth_conf = predictions["points_map"][...,
+                                                              2], predictions["unc_metric"]
+
+    depth_tensor, extrs, intrs = depth_map.squeeze().cpu().numpy(
+    ), extrinsic.squeeze().cpu().numpy(), intrinsic.squeeze().cpu().numpy()
+    unc_metric = depth_conf.squeeze().cpu().numpy() > 0.5
+    tracker_model.spatrack.track_num = 512
+    tracker_model.to("cuda")
+    grid_pts = get_points_on_a_grid(
+        30, (video_input.shape[3], video_input.shape[4]), 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, _, _, _, _, video_out = tracker_model.forward(
+            video_input.squeeze(), depth=depth_tensor, intrs=intrs, extrs=extrs, queries=query_xyt,
+            fps=1, unc_metric=unc_metric, support_frame=len(video_input.squeeze())-1
         )
-    return result.frames[0]
+    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_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"]
+# --- TTM WAN INFERENCE FUNCTION ---
 
-    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]
+@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."
 
-    # Preprocess for VGGT
-    v_in = preprocess_image(vt)[None].cuda()
-    with torch.no_grad():
-        preds = vggt4track_model(v_in / 255)
+    progress(0, desc="Loading Wan TTM Pipeline...")
+    pipe = get_wan_pipeline()
 
-    # 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()
+    progress(0.2, desc="Preparing inputs...")
+    image = load_image(first_frame_path)
 
-    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
+    # Standard Wan Negative Prompt
+    negative_prompt = (
+        "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，"
+        "低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，"
+        "毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
     )
 
-    # 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]))
+    # Match resolution logic from run_wan.py
+    max_area = 480 * 832
+    mod_value = pipe.vae_scale_factor_spatial * \
+        pipe.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="Generating Video (this may take a few minutes)...")
+    generator = torch.Generator(device="cuda").manual_seed(0)
+
+    with torch.inference_mode():
+        result = pipe(
+            image=image,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=81,  # Wan default
+            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)
 
-    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)
+    return output_path, "✅ TTM Video generated successfully with Wan 2.2!"
 
-    first_frame_path = os.path.join(res_dir, "first.png")
-    cv2.imwrite(first_frame_path, cv2.cvtColor(rgb[0], cv2.COLOR_RGB2BGR))
+# --- MODIFIED PROCESS VIDEO TO RETURN FILE PATHS ---
 
-    # 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)
 
-    return rend_path, wan_video_path, rend_res['motion_signal'], rend_res['mask'], first_frame_path, "✅ Generated successfully!"
+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"
 
-# --- 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")
+    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")
+
+    # Shared state for TTM files
+    first_frame_file = gr.State()
+    motion_signal_file = gr.State()
+    mask_file = gr.State()
 
     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.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():
-                v_msig = gr.Video(label="Motion Signal")
-                v_mask = gr.Video(label="Mask")
+                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")
+
+    # Event Handlers
+    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(
+        # Link output files to state for the next step
+        fn=lambda a, b, c, d: (b, c, d),
+        inputs=[output_video, motion_signal_output,
+                mask_output, first_frame_output],
+        outputs=[motion_signal_file, mask_file, first_frame_file]
+    )
 
-    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()])
+    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]
+    )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch(share=False)