Add pure-PyTorch mamba_ssm stand-in so the StreamMind gate loads

app.py

@@ -24,8 +24,6 @@ os.makedirs(os.environ["ONLINE_CODEC_CACHE_DIR"], exist_ok=True)
 os.makedirs(os.environ["STREAMMIND_CLIP_CACHE"], exist_ok=True)
 
 import importlib
-import math
-import shutil
 import subprocess
 import time
 import traceback
@@ -35,15 +33,9 @@ from pathlib import Path
 
 import spaces  # noqa: F401  (must precede torch)
 import torch
-import torch.nn as nn
-import torch.nn.functional as F
 import gradio as gr
-import numpy as np
 from PIL import Image
-from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file
-from transformers import AutoModelForCausalLM, AutoProcessor, Qwen3Config
-from transformers.models.qwen3 import Qwen3ForCausalLM
+from transformers import AutoModelForCausalLM, AutoProcessor
 
 MODEL_ID = "microsoft/Mage-VL"
 
@@ -80,225 +72,24 @@ ProcessorCls = type(processor)
 print(f"[init] model loaded, remote package = {_REMOTE_PKG}", flush=True)
 
 
-# ----------------------------------------------------- pure-torch StreamMind gate
-# The released gate module imports `mamba_ssm.models.mixer_seq_simple.create_block`,
-# whose CUDA kernels (selective_scan_cuda / causal_conv1d) have no wheel for the
-# ZeroGPU Blackwell (sm_120) / torch-2.11 cell and cannot be compiled inside a
-# Space build. The block below is a faithful pure-PyTorch port of Mamba-1
-# (`mamba_ssm.modules.mamba_simple.Mamba` + `mamba_ssm.modules.block.Block`,
-# fused_add_norm=False, rms_norm=False, d_intermediate=0) using the upstream
-# `selective_scan_ref` reference recurrence, with identical parameter names so the
-# released `streammind_gate.safetensors` loads with strict=True.
-
-
-def _selective_scan_ref(u, delta, A, B, C, D=None, z=None, delta_bias=None,
-                        delta_softplus=False):
-    """u,delta,z: (b, d, l) · A: (d, n) · B,C: (b, n, l) · D: (d,)"""
-    dtype_in = u.dtype
-    u = u.float()
-    delta = delta.float()
-    if delta_bias is not None:
-        delta = delta + delta_bias[..., None].float()
-    if delta_softplus:
-        delta = F.softplus(delta)
-    batch, dim, dstate = u.shape[0], A.shape[0], A.shape[1]
-    B = B.float()
-    C = C.float()
-    x = A.new_zeros((batch, dim, dstate))
-    deltaA = torch.exp(torch.einsum("bdl,dn->bdln", delta, A))
-    deltaB_u = torch.einsum("bdl,bnl,bdl->bdln", delta, B, u)
-    ys = []
-    for i in range(u.shape[2]):
-        x = deltaA[:, :, i] * x + deltaB_u[:, :, i]
-        ys.append(torch.einsum("bdn,bn->bd", x, C[:, :, i]))
-    y = torch.stack(ys, dim=2)
-    out = y if D is None else y + u * D.unsqueeze(-1)
-    if z is not None:
-        out = out * F.silu(z)
-    return out.to(dtype=dtype_in)
-
-
-class _Mamba(nn.Module):
-    def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
-        super().__init__()
-        self.d_model = d_model
-        self.d_state = d_state
-        self.d_conv = d_conv
-        self.d_inner = int(expand * d_model)
-        self.dt_rank = math.ceil(d_model / 16)
-        self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False)
-        self.conv1d = nn.Conv1d(self.d_inner, self.d_inner, bias=True,
-                                kernel_size=d_conv, groups=self.d_inner,
-                                padding=d_conv - 1)
-        self.act = nn.SiLU()
-        self.x_proj = nn.Linear(self.d_inner, self.dt_rank + 2 * self.d_state, bias=False)
-        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)
-        self.A_log = nn.Parameter(torch.zeros(self.d_inner, self.d_state))
-        self.D = nn.Parameter(torch.ones(self.d_inner))
-        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
-
-    def forward(self, hidden_states, inference_params=None, **kwargs):
-        b, l, _ = hidden_states.shape
-        xz = self.in_proj(hidden_states).transpose(1, 2)          # (b, 2*d_inner, l)
-        A = -torch.exp(self.A_log.float())                        # (d_inner, d_state)
-        x, z = xz.chunk(2, dim=1)
-        x = self.act(self.conv1d(x)[..., :l])
-        x_dbl = self.x_proj(x.transpose(1, 2).reshape(b * l, self.d_inner))
-        dt, B, C = torch.split(x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=-1)
-        dt = (self.dt_proj.weight @ dt.t()).view(self.d_inner, b, l).permute(1, 0, 2)
-        B = B.view(b, l, self.d_state).transpose(1, 2).contiguous()
-        C = C.view(b, l, self.d_state).transpose(1, 2).contiguous()
-        y = _selective_scan_ref(
-            x, dt, A, B, C, self.D.float(), z=z,
-            delta_bias=self.dt_proj.bias.float(), delta_softplus=True,
-        )
-        return self.out_proj(y.transpose(1, 2))
-
-
-class _Block(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.mixer = _Mamba(dim)
-
-    def forward(self, hidden_states, residual=None, inference_params=None):
-        residual = (hidden_states + residual) if residual is not None else hidden_states
-        hidden_states = self.norm(residual.to(dtype=self.norm.weight.dtype))
-        hidden_states = self.mixer(hidden_states, inference_params=inference_params)
-        return hidden_states, residual
-
-
-# The remaining gate modules below are verbatim ports of the released
-# `streammind_gate.py` (PreNet / PostNet / VideoMamba / Qwen3ForCausalLMCls /
-# ClsNet / StreamMindGate); only `create_block` is swapped for `_Block` above.
-# The released module cannot simply be imported because line 6 of it is a
-# module-level `from mamba_ssm... import create_block`.
-
-
-class _PreNet(nn.Module):
-    def __init__(self, d_code, d_model):
-        super().__init__()
-        self.fc3 = nn.Linear(d_code, d_model)
-
-    def forward(self, x):
-        return F.leaky_relu(self.fc3(x))
-
-
-class _PostNet(nn.Module):
-    def __init__(self, d_model, n_class):
-        super().__init__()
-        self.fc3 = nn.Linear(d_model, n_class)
-
-    def forward(self, x):
-        return self.fc3(F.leaky_relu(x))
-
-
-class _VideoMamba(nn.Module):
-    def __init__(self, d_model, n_ssm=1):
-        super().__init__()
-        self.ssms = nn.ModuleList([_Block(d_model) for _ in range(n_ssm)])
-        self.norm_fn = nn.LayerNorm(d_model)
-
-    def forward(self, embeds, inference_params=None):
-        hidden_states, residual = embeds, None
-        for ssm in self.ssms:
-            hidden_states, residual = ssm(
-                hidden_states, residual, inference_params=inference_params
-            )
-        residual = hidden_states + residual if residual is not None else hidden_states
-        return self.norm_fn(residual.to(dtype=self.norm_fn.weight.dtype))
-
-
-class _Qwen3ForCausalLMCls(Qwen3ForCausalLM):
-    def forward(self, inputs_embeds=None, labels=None, attention_mask=None, **kwargs):
-        outputs = self.model(inputs_embeds=inputs_embeds, attention_mask=attention_mask)
-        logits = self.lm_head(outputs.last_hidden_state).float()
-        return {"loss": None, "logits": logits}
-
-
-class _ClsNet(nn.Module):
-    def __init__(self, hidden_size=2560, num_layers=4):
-        super().__init__()
-        config = Qwen3Config(
-            vocab_size=2,
-            hidden_size=hidden_size,
-            num_hidden_layers=num_layers,
-            num_attention_heads=32,
-            num_key_value_heads=8,
-            intermediate_size=12288,
-            head_dim=128,
-            max_position_embeddings=8192,
-            rms_norm_eps=1e-6,
-            tie_word_embeddings=False,
-            attention_bias=False,
-        )
-        self.cls_model = _Qwen3ForCausalLMCls(config)
-
-    def forward(self, x, labels=None, attention_mask=None):
-        return self.cls_model(inputs_embeds=x, labels=labels, attention_mask=attention_mask)
-
-
-class StreamMindGate(nn.Module):
-    def __init__(self, hidden_size=2560):
-        super().__init__()
-        self.pre_net = _PreNet(hidden_size, hidden_size)
-        self.mamba_model = _VideoMamba(hidden_size)
-        self.post_net = _PostNet(hidden_size, hidden_size)
-        self.cls_net = _ClsNet(hidden_size=hidden_size, num_layers=4)
-
-    def perception_tokens(self, vision_tokens):
-        """[B,T,P,D] visual patches -> one EPFE token per time step."""
-        x = vision_tokens.mean(dim=2)
-        batch, time_, dim = x.shape
-        x = self.pre_net(x.reshape(batch * time_, dim)).reshape(batch, time_, dim)
-        x = self.mamba_model(x)
-        x = self.post_net(x.reshape(batch * time_, dim)).reshape(batch, time_, dim)
-        return x
-
-    def forward(self, vision_tokens, response_positions=None):
-        """Return [B,T,2] silence/speak logits for every EPFE time step."""
-        tokens = self.perception_tokens(vision_tokens)
-        batch, time_, dim = tokens.shape
-        target_ids = torch.zeros(batch, time_, dtype=torch.long, device=tokens.device)
-        if response_positions is not None:
-            target_ids[:, torch.as_tensor(response_positions, device=tokens.device) - 1] = 1
-        targets = self.cls_net.cls_model.model.embed_tokens(
-            target_ids.reshape(batch * time_)
-        )
-        pair = torch.stack((tokens.reshape(batch * time_, dim), targets), dim=1)
-        rotary = self.cls_net.cls_model.model.rotary_emb
-        saved_inv_freq = rotary.inv_freq
-        try:
-            # Match the training checkpoint, where the full model (including
-            # non-persistent Qwen3 RoPE buffers) was cast to BF16.
-            rotary.inv_freq = rotary.inv_freq.to(pair.dtype)
-            output = self.cls_net(
-                pair,
-                attention_mask=torch.ones(pair.shape[:2], device=pair.device),
-            )
-        finally:
-            rotary.inv_freq = saved_inv_freq
-        return output["logits"][:, 0].reshape(batch, time_, 2)
-
-
-def _build_gate():
-    """StreamMindGate with the pure-torch Mamba block, released weights."""
-    gate = StreamMindGate(model.config.text_config.hidden_size)
-    state = load_file(hf_hub_download(MODEL_ID, "streammind_gate.safetensors"))
-    gate.load_state_dict(state, strict=True)
-    return gate.to("cuda", dtype=torch.bfloat16).eval()
-
+# ----------------------------------------------------- StreamMind cognition gate
+# `streammind_gate.py` in the model repo imports
+# `mamba_ssm.models.mixer_seq_simple.create_block`; the real mamba-ssm ships CUDA
+# extensions with no wheel for this runtime, so `./mamba_ssm/` in this Space is a
+# pure-PyTorch stand-in with identical parameter names/shapes (see its docstring).
+# With that in place the released gate code + checkpoint load unmodified.
 
 GATE_ERROR = ""
 try:
     _t0 = time.perf_counter()
-    model.model.streammind_gate = _build_gate()
+    model.model._load_streammind_gate()
     print(f"[gate] loaded in {time.perf_counter() - _t0:.1f}s", flush=True)
 except Exception as exc:  # surfaced in the streaming tab, never kills the boot
     GATE_ERROR = f"{type(exc).__name__}: {exc}"
     print("[gate] FAILED:\n" + traceback.format_exc(), flush=True)
 
 
+
 # --------------------------------------------------------------------------- utils
 
 

mamba_ssm/__init__.py

@@ -0,0 +1,26 @@
+"""Pure-PyTorch stand-in for the parts of `mamba-ssm` that Mage-VL needs.
+
+`microsoft/Mage-VL`'s remote code (`streammind_gate.py`) does
+
+    from mamba_ssm.models.mixer_seq_simple import create_block
+
+at module level, and `transformers.dynamic_module_utils.check_imports` imports
+every top-level dependency of every remote file before it will load the model —
+so `mamba_ssm` must be importable even though only the StreamMind cognition gate
+uses it.
+
+The real `mamba-ssm` ships CUDA extensions (`selective_scan_cuda`,
+`causal_conv1d_cuda`) with no wheel for the ZeroGPU Blackwell (sm_120) /
+torch-2.11 / cp312 runtime, and compiling them from source is not viable inside
+a Space build. This package therefore provides a faithful pure-PyTorch port of
+Mamba-1 (`mamba_ssm.modules.mamba_simple.Mamba` + `mamba_ssm.modules.block.Block`
+with the defaults `create_block` uses: rms_norm=False, fused_add_norm=False,
+residual_in_fp32=False, d_intermediate=0) built on the upstream
+`selective_scan_ref` reference recurrence, with identical parameter names and
+shapes so `streammind_gate.safetensors` loads with `strict=True`.
+
+The gate runs over a handful of EPFE tokens (one per codec canvas), so the slow
+sequential scan costs milliseconds — the CUDA kernel buys nothing here.
+"""
+
+__version__ = "2.2.6.mage-vl-pure-torch"

mamba_ssm/models/mixer_seq_simple.py

@@ -0,0 +1,173 @@
+"""Pure-PyTorch `create_block` — Mamba-1 mixer + residual block, no CUDA kernels.
+
+Mirrors `mamba_ssm.modules.mamba_simple.Mamba` and `mamba_ssm.modules.block.Block`
+for the argument set `create_block()` is called with by Mage-VL's StreamMind gate
+(`create_block(d_model, d_intermediate=0, layer_idx=i)`), i.e. the upstream
+defaults: ssm_cfg={} -> d_state=16, d_conv=4, expand=2, dt_rank=ceil(d_model/16),
+rms_norm=False -> nn.LayerNorm, fused_add_norm=False, residual_in_fp32=False.
+
+Parameter names/shapes match upstream exactly, so the released
+`streammind_gate.safetensors` loads with strict=True:
+
+    mixer.in_proj.weight    (2*d_inner, d_model)
+    mixer.conv1d.{weight,bias}  (d_inner, 1, d_conv) / (d_inner,)
+    mixer.x_proj.weight     (dt_rank + 2*d_state, d_inner)
+    mixer.dt_proj.{weight,bias}  (d_inner, dt_rank) / (d_inner,)
+    mixer.A_log             (d_inner, d_state)
+    mixer.D                 (d_inner,)
+    mixer.out_proj.weight   (d_model, d_inner)
+    norm.{weight,bias}      (d_model,)
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def selective_scan_ref(u, delta, A, B, C, D=None, z=None, delta_bias=None,
+                       delta_softplus=False, return_last_state=False):
+    """Upstream reference implementation (single-group, non-complex path).
+
+    u, delta, z: (b, d, l) · A: (d, n) · B, C: (b, n, l) · D: (d,)
+    """
+    dtype_in = u.dtype
+    u = u.float()
+    delta = delta.float()
+    if delta_bias is not None:
+        delta = delta + delta_bias[..., None].float()
+    if delta_softplus:
+        delta = F.softplus(delta)
+    batch, dim = u.shape[0], A.shape[0]
+    dstate = A.shape[1]
+    B = B.float()
+    C = C.float()
+    x = A.new_zeros((batch, dim, dstate))
+    deltaA = torch.exp(torch.einsum("bdl,dn->bdln", delta, A))
+    deltaB_u = torch.einsum("bdl,bnl,bdl->bdln", delta, B, u)
+    ys = []
+    last_state = None
+    for i in range(u.shape[2]):
+        x = deltaA[:, :, i] * x + deltaB_u[:, :, i]
+        ys.append(torch.einsum("bdn,bn->bd", x, C[:, :, i]))
+        if i == u.shape[2] - 1:
+            last_state = x
+    y = torch.stack(ys, dim=2)
+    out = y if D is None else y + u * D.unsqueeze(-1)
+    if z is not None:
+        out = out * F.silu(z)
+    out = out.to(dtype=dtype_in)
+    return (out, last_state) if return_last_state else out
+
+
+class Mamba(nn.Module):
+    def __init__(self, d_model, d_state=16, d_conv=4, expand=2, dt_rank="auto",
+                 conv_bias=True, bias=False, layer_idx=None, **kwargs):
+        super().__init__()
+        self.d_model = d_model
+        self.d_state = d_state
+        self.d_conv = d_conv
+        self.expand = expand
+        self.d_inner = int(expand * d_model)
+        self.dt_rank = math.ceil(d_model / 16) if dt_rank == "auto" else dt_rank
+        self.layer_idx = layer_idx
+
+        self.in_proj = nn.Linear(self.d_model, self.d_inner * 2, bias=bias)
+        self.conv1d = nn.Conv1d(
+            in_channels=self.d_inner, out_channels=self.d_inner, bias=conv_bias,
+            kernel_size=d_conv, groups=self.d_inner, padding=d_conv - 1,
+        )
+        self.activation = "silu"
+        self.act = nn.SiLU()
+        self.x_proj = nn.Linear(self.d_inner, self.dt_rank + self.d_state * 2, bias=False)
+        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)
+        self.A_log = nn.Parameter(torch.zeros(self.d_inner, self.d_state))
+        self.D = nn.Parameter(torch.ones(self.d_inner))
+        self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias)
+
+    def forward(self, hidden_states, inference_params=None, **kwargs):
+        batch, seqlen, _ = hidden_states.shape
+        xz = self.in_proj(hidden_states).transpose(1, 2)          # (b, 2*d_inner, l)
+        A = -torch.exp(self.A_log.float())                        # (d_inner, d_state)
+        x, z = xz.chunk(2, dim=1)
+        x = self.act(self.conv1d(x)[..., :seqlen])
+        x_dbl = self.x_proj(x.transpose(1, 2).reshape(batch * seqlen, self.d_inner))
+        dt, B, C = torch.split(
+            x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=-1
+        )
+        dt = (self.dt_proj.weight @ dt.t()).view(self.d_inner, batch, seqlen)
+        dt = dt.permute(1, 0, 2).contiguous()                     # (b, d_inner, l)
+        B = B.view(batch, seqlen, self.d_state).transpose(1, 2).contiguous()
+        C = C.view(batch, seqlen, self.d_state).transpose(1, 2).contiguous()
+        y = selective_scan_ref(
+            x, dt, A, B, C, self.D.float(), z=z,
+            delta_bias=self.dt_proj.bias.float(), delta_softplus=True,
+        )
+        return self.out_proj(y.transpose(1, 2))
+
+
+class Block(nn.Module):
+    """`mamba_ssm.modules.block.Block` with fused_add_norm=False, mlp=Identity."""
+
+    def __init__(self, dim, mixer_cls, norm_cls=nn.LayerNorm, mlp_cls=nn.Identity,
+                 fused_add_norm=False, residual_in_fp32=False):
+        super().__init__()
+        self.residual_in_fp32 = residual_in_fp32
+        self.fused_add_norm = fused_add_norm
+        self.norm = norm_cls(dim)
+        self.mixer = mixer_cls(dim)
+        if mlp_cls is not nn.Identity:
+            self.norm2 = norm_cls(dim)
+            self.mlp = mlp_cls(dim)
+        else:
+            self.mlp = None
+
+    def forward(self, hidden_states, residual=None, inference_params=None, **kwargs):
+        residual = (hidden_states + residual) if residual is not None else hidden_states
+        hidden_states = self.norm(residual.to(dtype=self.norm.weight.dtype))
+        if self.residual_in_fp32:
+            residual = residual.to(torch.float32)
+        hidden_states = self.mixer(hidden_states, inference_params=inference_params)
+        if self.mlp is not None:
+            residual = hidden_states + residual
+            hidden_states = self.norm2(residual.to(dtype=self.norm2.weight.dtype))
+            hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+def create_block(d_model, d_intermediate=0, ssm_cfg=None, attn_layer_idx=None,
+                 attn_cfg=None, norm_epsilon=1e-5, rms_norm=False,
+                 residual_in_fp32=False, fused_add_norm=False, layer_idx=None,
+                 device=None, dtype=None, **kwargs):
+    if d_intermediate:
+        raise NotImplementedError(
+            "This pure-PyTorch mamba_ssm stand-in only supports d_intermediate=0 "
+            "(the configuration used by Mage-VL's StreamMind gate)."
+        )
+    if attn_layer_idx and layer_idx in attn_layer_idx:
+        raise NotImplementedError(
+            "Attention blocks are not supported by this mamba_ssm stand-in."
+        )
+    if rms_norm:
+        raise NotImplementedError(
+            "rms_norm=True is not supported by this mamba_ssm stand-in."
+        )
+    factory_kwargs = {"device": device, "dtype": dtype}
+    ssm_cfg = dict(ssm_cfg or {})
+    ssm_cfg.pop("layer", None)
+
+    def mixer_cls(dim):
+        return Mamba(dim, layer_idx=layer_idx, **ssm_cfg)
+
+    def norm_cls(dim):
+        return nn.LayerNorm(dim, eps=norm_epsilon)
+
+    block = Block(
+        d_model, mixer_cls, norm_cls=norm_cls, mlp_cls=nn.Identity,
+        fused_add_norm=fused_add_norm, residual_in_fp32=residual_in_fp32,
+    )
+    block.layer_idx = layer_idx
+    if factory_kwargs["device"] is not None or factory_kwargs["dtype"] is not None:
+        block = block.to(**{k: v for k, v in factory_kwargs.items() if v is not None})
+    return block