Add Apple Silicon (MPS) support with bf16 autocast

demo.py

@@ -219,7 +219,7 @@ def _warm_streaming(model, images, scale_frames, warm_stream_n, dtype,
     for _ in range(passes):
         model.clean_kv_cache()
         torch.compiler.cudagraph_mark_step_begin()
-        with torch.no_grad(), torch.amp.autocast("cuda", dtype=dtype):
+        with torch.no_grad(), torch.amp.autocast(images.device.type, dtype=dtype, enabled=dtype != torch.float32):
             model.forward(
                 warm_scale,
                 num_frame_for_scale=scale_frames,
@@ -231,7 +231,7 @@ def _warm_streaming(model, images, scale_frames, warm_stream_n, dtype,
             if not is_keyframe:
                 model._set_skip_append(True)
             torch.compiler.cudagraph_mark_step_begin()
-            with torch.no_grad(), torch.amp.autocast("cuda", dtype=dtype):
+            with torch.no_grad(), torch.amp.autocast(images.device.type, dtype=dtype, enabled=dtype != torch.float32):
                 model.forward(
                     warm_stream[:, i:i + 1],
                     num_frame_for_scale=scale_frames,
@@ -418,7 +418,12 @@ def main():
     assert args.image_folder or args.video_path, \
         "Provide --image_folder or --video_path"
 
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
 
     # ── Load images & model ──────────────────────────────────────────────────
     t0 = time.time()
@@ -447,6 +452,8 @@ def main():
     # Pick inference dtype; autocast still runs for the ops that need fp32 (e.g. LayerNorm).
     if torch.cuda.is_available():
         dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
+    elif device.type == "mps":
+        dtype = torch.bfloat16
     else:
         dtype = torch.float32
 
@@ -542,7 +549,8 @@ def main():
 
     output_device = torch.device("cpu") if args.offload_to_cpu else None
 
-    with torch.no_grad(), torch.amp.autocast("cuda", dtype=dtype):
+    autocast_ctx = torch.amp.autocast(device.type, dtype=dtype, enabled=dtype != torch.float32)
+    with torch.no_grad(), autocast_ctx:
         if args.mode == "streaming":
             predictions = model.inference_streaming(
                 images,

lingbot_map/layers/rope.py

@@ -362,6 +362,10 @@ def forward(self, ppf, pph, ppw, patch_start_idx, device: torch.device, f_start:
         """
 
         # 步骤1：将预计算的频率移到目标设备，并分割成三个维度
+        # MPS has no float64/complex128 — downcast the precomputed freqs to
+        # complex64 before moving them (precompute itself stays fp64 on CPU).
+        if torch.device(device).type == "mps" and self.freqs.dtype == torch.complex128:
+            self.freqs = self.freqs.to(torch.complex64)
         self.freqs = self.freqs.to(device)
         # 获取实际的维度分配
         if hasattr(self, 'fhw_dim') and self.fhw_dim is not None:

lingbot_map/models/gct_base.py

@@ -166,7 +166,7 @@ def _predict_camera(
 
         camera_sliding_window = sliding_window_size if self.enable_camera_sliding_window else -1
 
-        with torch.amp.autocast('cuda', enabled=False):
+        with torch.amp.autocast(aggregated_tokens_list_fp32[0].device.type, enabled=False):
             pose_enc_list = self.camera_head(
                 aggregated_tokens_list_fp32,
                 mask=mask,
@@ -194,7 +194,7 @@ def _predict_depth(
         aggregated_tokens_list_fp32 = [t.float() for t in aggregated_tokens_list]
         images_fp32 = images.float()
 
-        with torch.amp.autocast('cuda', enabled=False):
+        with torch.amp.autocast(aggregated_tokens_list_fp32[0].device.type, enabled=False):
             depth, depth_conf = self.depth_head(
                 aggregated_tokens_list_fp32,
                 images=images_fp32,
@@ -216,7 +216,7 @@ def _predict_points(
         aggregated_tokens_list_fp32 = [t.float() for t in aggregated_tokens_list]
         images_fp32 = images.float()
 
-        with torch.amp.autocast('cuda', enabled=False):
+        with torch.amp.autocast(aggregated_tokens_list_fp32[0].device.type, enabled=False):
             pts3d, pts3d_conf = self.point_head(
                 aggregated_tokens_list_fp32,
                 images=images_fp32,
@@ -238,7 +238,7 @@ def _predict_local_points(
         aggregated_tokens_list_fp32 = [t.float() for t in aggregated_tokens_list]
         images_fp32 = images.float()
 
-        with torch.amp.autocast('cuda', enabled=False):
+        with torch.amp.autocast(aggregated_tokens_list_fp32[0].device.type, enabled=False):
             pts3d, pts3d_conf = self.local_point_head(
                 aggregated_tokens_list_fp32,
                 images=images_fp32,

lingbot_map/vis/point_cloud_viewer.py

@@ -22,6 +22,7 @@
 import numpy as np
 import torch
 import cv2
+import matplotlib
 import matplotlib.cm as cm
 from tqdm.auto import tqdm
 
@@ -1045,7 +1046,7 @@ def read_data(self, pc_list, color_list, conf_list, edge_color_list=None):
             normalized_indices = np.array(list(range(num_cameras))) / (num_cameras - 1)
         else:
             normalized_indices = np.array([0.0])
-        cmap = cm.get_cmap('viridis')
+        cmap = matplotlib.colormaps.get_cmap('viridis')
         self.camera_colors = cmap(normalized_indices)
         return pcs, step_list
 

lingbot_map/vis/utils.py

@@ -14,6 +14,7 @@
 import numpy as np
 import torch
 import cv2
+import matplotlib
 import matplotlib.cm as cm
 
 
@@ -67,7 +68,7 @@ def get_vertical_colorbar(
     canvas = FigureCanvasAgg(fig)
 
     ax = fig.add_subplot(111)
-    cmap = cm.get_cmap(cmap_name)
+    cmap = matplotlib.colormaps.get_cmap(cmap_name)
     norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
 
     tick_cnt = 6
@@ -135,7 +136,7 @@ def colorize_np(
     x = np.clip(x, vmin, vmax)
     x = (x - vmin) / (vmax - vmin)
 
-    cmap = cm.get_cmap(cmap_name)
+    cmap = matplotlib.colormaps.get_cmap(cmap_name)
     x_new = cmap(x)[:, :, :3]
 
     if mask is not None: