Update basicsr/archs/mair_arch.py

basicsr/archs/mair_arch.py

@@ -7,7 +7,12 @@ import torch.nn.functional as F
 from functools import partial
 from typing import Optional, Callable
 from timm.models.layers import DropPath, to_2tuple, trunc_normal_
-from mamba_ssm.ops.selective_scan_interface import selective_scan_fn, selective_scan_ref
+try:
+    from mamba_ssm.ops.selective_scan_interface import selective_scan_fn, selective_scan_ref
+except Exception:
+    # CPU fallback (slow)
+    from mamba_cpu_fallback.selective_scan_interface import selective_scan_ref
+    selective_scan_fn = selective_scan_ref
 from einops import rearrange, repeat
 import time
 import sys