deploy cotracker on cpu

app.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import spaces
 import gradio as gr
 import torch
 import argparse
@@ -8,7 +9,8 @@ import numpy as np
 import torchvision.transforms as transforms
 from moviepy.editor import VideoFileClip
 from diffusers.utils import load_image, load_video
-import spaces
+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")
@@ -247,7 +249,7 @@ def process_motion_transfer(source, prompt, mt_repaint_option, mt_repaint_image)
             print('Export tracking video via MoGe')
         else:
             # 使用 cotracker
-            pred_tracks, pred_visibility = das.generate_tracking_cotracker(video_tensor)
+            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')
 
@@ -312,8 +314,8 @@ def process_camera_control(source, prompt, camera_motion, tracking_method):
             )
             print('Export tracking video via MoGe')
         else:
-            # 使用 cotracker
-            pred_tracks, pred_visibility = das.generate_tracking_cotracker(video_tensor)
+            # 使用在CPU上运行的cotracker
+            pred_tracks, pred_visibility = generate_tracking_cotracker(video_tensor)
             
             t, c, h, w = video_tensor.shape
             new_width = 518
@@ -419,8 +421,8 @@ def process_object_manipulation(source, prompt, object_motion, object_mask, trac
             )
             print('Export tracking video via MoGe')
         else:
-            # 使用 cotracker
-            pred_tracks, pred_visibility = das.generate_tracking_cotracker(video_tensor)
+            # 使用在CPU上运行的cotracker
+            pred_tracks, pred_visibility = generate_tracking_cotracker(video_tensor)
             
             t, c, h, w = video_tensor.shape
             new_width = 518
@@ -515,6 +517,58 @@ def process_mesh_animation(source, prompt, tracking_video, ma_repaint_option, ma
         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")

models/pipelines.py

@@ -470,7 +470,7 @@ class DiffusionAsShaderPipeline:
             del cotracker
             torch.cuda.empty_cache()
 
-    def visualize_tracking_cotracker(self, points, vis_mask=None, save_tracking=True, point_wise=4, video_size=(480, 720)):
+    def visualize_tracking_cotracker(self, points, vis_mask=None, save_tracking=True, point_wise=10, video_size=(480, 720)):
         """Visualize tracking results from CoTracker
         
         Args:

submodules/vggt

@@ -0,0 +1 @@
+Subproject commit b02cc03ceee70821ed1231a530c1992507ef9862