Make EgoSpeak Mamba dependency optional

third_party/EgoSpeak/MiniROAD/model/__init__.py

@@ -1,4 +1,3 @@
-from .mamba import OADMamba
 from .model_builder import build_model
-from .rnn import MROAD
+from .rnn import MROAD, MROADA
 from .transformer_models import ViTEnc

third_party/EgoSpeak/MiniROAD/model/model_builder.py

@@ -1,9 +1,37 @@
 __all__ = ['build_model']
 
+import importlib
+
 from utils import Registry
 
 META_ARCHITECTURES = Registry()
+_MAMBA_MODELS = frozenset({"Mamba", "MambaA"})
+
+
+def _load_optional_architecture(model_name):
+    if model_name not in _MAMBA_MODELS or model_name in META_ARCHITECTURES:
+        return
+
+    try:
+        importlib.import_module(".mamba", package=__package__)
+    except ModuleNotFoundError as error:
+        missing_package = error.name or "mamba_ssm"
+        raise ModuleNotFoundError(
+            f"Model {model_name!r} requires the optional Mamba backend "
+            f"(missing package: {missing_package!r}). Install "
+            "requirements-mamba.txt, or use MiniROAD/MiniROADA without it."
+        ) from error
+
 
 def build_model(cfg, device=None):
-    model = META_ARCHITECTURES[cfg["model"]](cfg)
-    return model.to(device)
+    model_name = cfg["model"]
+    _load_optional_architecture(model_name)
+
+    if model_name not in META_ARCHITECTURES:
+        available = ", ".join(sorted(META_ARCHITECTURES))
+        raise KeyError(
+            f"Unknown model {model_name!r}. Available architectures: {available}"
+        )
+
+    model = META_ARCHITECTURES[model_name](cfg)
+    return model.to(device)

third_party/EgoSpeak/requirements-mamba.txt

@@ -0,0 +1,4 @@
+-r requirements.txt
+
+# Optional: only required by the MiniROAD Mamba/MambaA baselines.
+mamba-ssm==2.2.4

third_party/EgoSpeak/requirements.txt

@@ -16,7 +16,6 @@ huggingface-hub==0.28.1
 idna==3.10
 Jinja2==3.1.5
 joblib==1.4.2
-mamba-ssm==2.2.4
 Markdown==3.7
 MarkupSafe==3.0.2
 -e git+ssh://git@github.com/jun297/EgoSpeak-CAM.git@dcbfd115abd8ce6ecfe37cf4a40e71d0089c98a1#egg=MiniROAD&subdirectory=MiniROAD