Upload example/ls_eend_step_model.py

example/ls_eend_step_model.py

@@ -0,0 +1,384 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+if TYPE_CHECKING:
+    from ls_eend_runtime import LSEENDInferenceEngine
+
+
+@dataclass(frozen=True)
+class StepStateLayout:
+    input_dim: int
+    full_output_dim: int
+    real_output_dim: int
+    encoder_layers: int
+    decoder_layers: int
+    encoder_dim: int
+    num_heads: int
+    key_dim: int
+    head_dim: int
+    encoder_conv_cache_len: int
+    top_buffer_len: int
+    conv_delay: int
+    max_nspks: int
+
+
+def build_state_layout(engine: Any) -> StepStateLayout:
+    model = engine.model
+    params = engine.config["model"]["params"]
+    n_units = int(params["n_units"])
+    n_heads = int(params["n_heads"])
+    max_nspks = int(engine.decode_max_nspks)
+    encoder_conv_cache_len = int(params["conv_kernel_size"]) - 1
+    top_buffer_len = 2 * int(params["conv_delay"]) + 1
+    return StepStateLayout(
+        input_dim=(2 * engine.config["data"]["context_recp"] + 1) * engine.config["data"]["feat"]["n_mels"],
+        full_output_dim=max_nspks,
+        real_output_dim=max(0, max_nspks - 2),
+        encoder_layers=int(params["enc_n_layers"]),
+        decoder_layers=int(params["dec_n_layers"]),
+        encoder_dim=n_units,
+        num_heads=n_heads,
+        key_dim=n_units // n_heads,
+        head_dim=n_units // n_heads,
+        encoder_conv_cache_len=encoder_conv_cache_len,
+        top_buffer_len=top_buffer_len,
+        conv_delay=int(params["conv_delay"]),
+        max_nspks=max_nspks,
+    )
+
+
+def initial_state_tensors(layout: StepStateLayout, dtype: np.dtype = np.float32) -> dict[str, np.ndarray]:
+    return {
+        "enc_ret_kv": np.zeros(
+            (layout.encoder_layers, 1, layout.num_heads, layout.key_dim, layout.head_dim),
+            dtype=dtype,
+        ),
+        "enc_ret_scale": np.zeros((layout.encoder_layers, 1, layout.num_heads), dtype=dtype),
+        "enc_conv_cache": np.zeros(
+            (layout.encoder_layers, 1, layout.encoder_conv_cache_len, layout.encoder_dim),
+            dtype=dtype,
+        ),
+        "dec_ret_kv": np.zeros(
+            (layout.decoder_layers, layout.max_nspks, layout.num_heads, layout.key_dim, layout.head_dim),
+            dtype=dtype,
+        ),
+        "dec_ret_scale": np.zeros(
+            (layout.decoder_layers, layout.max_nspks, layout.num_heads),
+            dtype=dtype,
+        ),
+        "top_buffer": np.zeros((1, layout.top_buffer_len, layout.encoder_dim), dtype=dtype),
+    }
+
+
+def _as_rank3_scalar(value: torch.Tensor, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
+    return value.to(device=device, dtype=dtype).reshape(1, 1, 1)
+
+
+def _safe_l2_normalize(x: torch.Tensor, dim: int) -> torch.Tensor:
+    # 1e-12 underflows to zero in fp16 CoreML execution and can produce NaNs
+    # during warmup frames when an embedding or attractor vector is exactly zero.
+    return x / torch.norm(x, dim=dim, keepdim=True).clamp_min(1e-4)
+
+
+class OnlineStepModule(torch.nn.Module):
+    """Single online LS-EEND step with explicit state tensors for export/runtime backends."""
+
+    def __init__(self, model: torch.nn.Module, layout: StepStateLayout) -> None:
+        super().__init__()
+        self.model = model
+        self.layout = layout
+        self.encoder_decay = torch.exp(
+            self.model.enc.encoder.layers[0].sequential[1].module.ret_pos.decay
+        ).float()
+        self.decoder_decay = torch.exp(
+            self.model.dec.attractor_decoder.layers[0].ret_pos1.decay
+        ).float()
+
+    def forward(
+        self,
+        frame: torch.Tensor,
+        enc_ret_kv: torch.Tensor,
+        enc_ret_scale: torch.Tensor,
+        enc_conv_cache: torch.Tensor,
+        dec_ret_kv: torch.Tensor,
+        dec_ret_scale: torch.Tensor,
+        top_buffer: torch.Tensor,
+        ingest: torch.Tensor,
+        decode: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        dtype = frame.dtype
+        device = frame.device
+        ingest_scalar = _as_rank3_scalar(ingest, dtype, device)
+        decode_scalar = _as_rank3_scalar(decode, dtype, device)
+        ingest_vec = ingest.to(device=device, dtype=dtype).reshape(1, 1)
+        decode_vec = decode.to(device=device, dtype=dtype).reshape(1, 1)
+
+        x = self.model.enc.encoder.input_projection(frame)
+        x = self.model.enc.encoder.layer_norm(x)
+
+        new_enc_ret_kv = []
+        new_enc_ret_scale = []
+        new_enc_conv_cache = []
+
+        for layer_index, layer in enumerate(self.model.enc.encoder.layers):
+            old_kv = enc_ret_kv[layer_index]
+            old_scale = enc_ret_scale[layer_index]
+            old_conv = enc_conv_cache[layer_index]
+            x, candidate_kv, candidate_scale, candidate_conv = self._encoder_layer_step(
+                layer=layer,
+                x=x,
+                old_kv=old_kv,
+                old_scale=old_scale,
+                old_conv_cache=old_conv,
+            )
+            blended_kv = old_kv + (candidate_kv - old_kv) * ingest_scalar.unsqueeze(-1)
+            blended_scale = old_scale + (candidate_scale - old_scale) * ingest_vec
+            blended_conv = old_conv + (candidate_conv - old_conv) * ingest_scalar
+            new_enc_ret_kv.append(blended_kv)
+            new_enc_ret_scale.append(blended_scale)
+            new_enc_conv_cache.append(blended_conv)
+
+        appended_encoder_frame = x * ingest_scalar
+        top_buffer = torch.cat([top_buffer[:, 1:, :], appended_encoder_frame], dim=1)
+
+        emb = F.conv1d(
+            top_buffer.transpose(1, 2),
+            self.model.cnn.weight,
+            self.model.cnn.bias,
+        ).transpose(1, 2)
+        emb = _safe_l2_normalize(emb, dim=-1)
+
+        logits, candidate_dec_ret_kv, candidate_dec_ret_scale = self._decoder_step(
+            emb=emb,
+            dec_ret_kv=dec_ret_kv,
+            dec_ret_scale=dec_ret_scale,
+        )
+
+        new_dec_ret_kv = dec_ret_kv + (candidate_dec_ret_kv - dec_ret_kv) * decode_scalar.unsqueeze(-1)
+        new_dec_ret_scale = dec_ret_scale + (candidate_dec_ret_scale - dec_ret_scale) * decode_vec.unsqueeze(-1)
+
+        logits = logits * decode_scalar
+
+        return (
+            logits,
+            torch.stack(new_enc_ret_kv, dim=0),
+            torch.stack(new_enc_ret_scale, dim=0),
+            torch.stack(new_enc_conv_cache, dim=0),
+            new_dec_ret_kv,
+            new_dec_ret_scale,
+            top_buffer,
+        )
+
+    def _encoder_layer_step(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        old_kv: torch.Tensor,
+        old_scale: torch.Tensor,
+        old_conv_cache: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        ff1 = layer.sequential[0]
+        attn = layer.sequential[1].module
+        conv = layer.sequential[2].module
+        ff2 = layer.sequential[3]
+        final_norm = layer.sequential[4]
+
+        x = ff1.module(x) * ff1.module_factor + x * ff1.input_factor
+        attn_input = attn.layer_norm(x)
+        attn_output, candidate_kv, candidate_scale = self._retention_recurrent(
+            retention_module=attn.self_attn,
+            x=attn_input,
+            old_kv=old_kv,
+            old_scale=old_scale,
+            decay=self.encoder_decay,
+        )
+        x = x + attn.dropout(attn_output)
+        conv_output, candidate_conv = self._conformer_conv_step(conv, x, old_conv_cache)
+        x = x + conv_output
+        x = ff2.module(x) * ff2.module_factor + x * ff2.input_factor
+        return final_norm(x), candidate_kv, candidate_scale, candidate_conv
+
+    def _conformer_conv_step(
+        self,
+        conv_module: torch.nn.Module,
+        x: torch.Tensor,
+        old_cache: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        modules = conv_module.sequential
+
+        current = modules[0](x)
+        current = modules[1](current)
+        current = modules[2](current)
+        current = modules[3](current)
+
+        cache = old_cache.transpose(1, 2)
+        depthwise_window = torch.cat([cache, current], dim=2)
+        depthwise_conv = modules[4].conv
+        depthwise = F.conv1d(
+            depthwise_window,
+            depthwise_conv.weight,
+            depthwise_conv.bias,
+            stride=depthwise_conv.stride,
+            padding=0,
+            dilation=depthwise_conv.dilation,
+            groups=depthwise_conv.groups,
+        )
+        candidate_cache = depthwise_window[:, :, -self.layout.encoder_conv_cache_len :].transpose(1, 2)
+
+        depthwise = modules[5](depthwise)
+        depthwise = modules[6](depthwise)
+        depthwise = modules[7](depthwise)
+        depthwise = modules[8](depthwise)
+        return depthwise.transpose(1, 2), candidate_cache
+
+    def _decoder_step(
+        self,
+        emb: torch.Tensor,
+        dec_ret_kv: torch.Tensor,
+        dec_ret_scale: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        pos_enc = self.model.dec.pos_enc(emb, self.layout.max_nspks)
+        repeated_emb = emb.unsqueeze(dim=2).repeat(1, 1, self.layout.max_nspks, 1)
+        attractors = self.model.dec.convert(torch.cat([repeated_emb, pos_enc], dim=-1))
+
+        new_dec_ret_kv = []
+        new_dec_ret_scale = []
+        for layer_index, layer in enumerate(self.model.dec.attractor_decoder.layers):
+            attractors, candidate_kv, candidate_scale = self._fusion_layer_step(
+                layer=layer,
+                src=attractors,
+                old_kv=dec_ret_kv[layer_index],
+                old_scale=dec_ret_scale[layer_index],
+            )
+            new_dec_ret_kv.append(candidate_kv)
+            new_dec_ret_scale.append(candidate_scale)
+
+        if self.model.dec.attractor_decoder.norm is not None:
+            attractors = self.model.dec.attractor_decoder.norm(attractors)
+        attractors = _safe_l2_normalize(attractors, dim=-1)
+        logits = torch.matmul(emb.unsqueeze(dim=-2), attractors.transpose(-1, -2)).squeeze(dim=-2)
+        return logits, torch.stack(new_dec_ret_kv, dim=0), torch.stack(new_dec_ret_scale, dim=0)
+
+    def _fusion_layer_step(
+        self,
+        layer: torch.nn.Module,
+        src: torch.Tensor,
+        old_kv: torch.Tensor,
+        old_scale: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        batch_size, time_steps, speaker_count, feat_dim = src.shape
+        x = src.transpose(1, 2).reshape(batch_size * speaker_count, time_steps, feat_dim)
+
+        if layer.norm_first:
+            time_input = layer.norm11(x)
+            time_output, candidate_kv, candidate_scale = self._retention_recurrent(
+                retention_module=layer.self_attn1,
+                x=time_input,
+                old_kv=old_kv,
+                old_scale=old_scale,
+                decay=self.decoder_decay,
+            )
+            x = x + layer.dropout11(time_output)
+        else:
+            time_output, candidate_kv, candidate_scale = self._retention_recurrent(
+                retention_module=layer.self_attn1,
+                x=x,
+                old_kv=old_kv,
+                old_scale=old_scale,
+                decay=self.decoder_decay,
+            )
+            x = layer.norm11(x + layer.dropout11(time_output))
+
+        x = x.reshape(batch_size, speaker_count, time_steps, feat_dim).transpose(1, 2)
+        x = x.reshape(batch_size * time_steps, speaker_count, feat_dim)
+
+        if layer.norm_first:
+            x = x + self._speaker_attention(layer.self_attn2, layer.norm21(x))
+            x = x + layer._ff_block(layer.norm22(x))
+        else:
+            x = layer.norm21(x + self._speaker_attention(layer.self_attn2, x))
+            x = layer.norm22(x + layer._ff_block(x))
+
+        return x.reshape(batch_size, time_steps, speaker_count, feat_dim), candidate_kv, candidate_scale
+
+    def _retention_recurrent(
+        self,
+        retention_module: torch.nn.Module,
+        x: torch.Tensor,
+        old_kv: torch.Tensor,
+        old_scale: torch.Tensor,
+        decay: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        batch_size, target_length, _ = x.shape
+        q = retention_module.q_proj(x)
+        k = retention_module.k_proj(x)
+        v = retention_module.v_proj(x)
+        g = retention_module.g_proj(x)
+
+        k = k * retention_module.scaling
+        q = q.view(batch_size, target_length, retention_module.num_heads, retention_module.key_dim).transpose(1, 2)
+        k = k.view(batch_size, target_length, retention_module.num_heads, retention_module.key_dim).transpose(1, 2)
+        v = v.view(batch_size, retention_module.num_heads, retention_module.head_dim, 1)
+
+        qr = q
+        kr = k
+        kv = kr * v
+
+        decay = decay.to(device=x.device, dtype=x.dtype).reshape(1, retention_module.num_heads)
+        candidate_scale = old_scale * decay + 1.0
+        blend = (old_scale.sqrt() * decay / candidate_scale.sqrt()).unsqueeze(-1).unsqueeze(-1)
+        candidate_kv = old_kv * blend + kv / candidate_scale.sqrt().unsqueeze(-1).unsqueeze(-1)
+
+        output = torch.sum(qr * candidate_kv, dim=3)
+        output = retention_module.group_norm(output).reshape(
+            batch_size, target_length, retention_module.head_dim * retention_module.num_heads
+        )
+        output = retention_module.gate_fn(g) * output
+        output = retention_module.out_proj(output)
+        return output, candidate_kv, candidate_scale
+
+    def _speaker_attention(self, attention: torch.nn.MultiheadAttention, x: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, embed_dim = x.shape
+        head_dim = embed_dim // attention.num_heads
+        q_weight, k_weight, v_weight = attention.in_proj_weight.chunk(3, dim=0)
+        q_bias, k_bias, v_bias = attention.in_proj_bias.chunk(3, dim=0)
+
+        q = F.linear(x, q_weight, q_bias)
+        k = F.linear(x, k_weight, k_bias)
+        v = F.linear(x, v_weight, v_bias)
+
+        q = q.view(batch_size, seq_len, attention.num_heads, head_dim).transpose(1, 2)
+        k = k.view(batch_size, seq_len, attention.num_heads, head_dim).transpose(1, 2)
+        v = v.view(batch_size, seq_len, attention.num_heads, head_dim).transpose(1, 2)
+
+        attn = torch.matmul(q, k.transpose(-2, -1)) / (head_dim**0.5)
+        attn = torch.softmax(attn, dim=-1)
+        out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).reshape(batch_size, seq_len, embed_dim)
+        return F.linear(out, attention.out_proj.weight, attention.out_proj.bias)
+
+
+def load_step_module(
+    checkpoint_path: Path,
+    config_path: Path,
+    device: str = "cpu",
+) -> tuple[OnlineStepModule, StepStateLayout, "LSEENDInferenceEngine"]:
+    from ls_eend_runtime import LSEENDInferenceEngine
+
+    engine = LSEENDInferenceEngine(
+        checkpoint_path=checkpoint_path,
+        config_path=config_path,
+        device=device,
+    )
+    engine.model = engine.model.float().to(torch.device(device))
+    engine.model.eval()
+    layout = build_state_layout(engine)
+    module = OnlineStepModule(engine.model, layout).to(torch.device(device)).eval()
+    return module, layout, engine