"""变体3 hidden2cv 映射层: ST LM hidden 直接映射进 CV3 LLM 文本 embedding 位。 与 bridge (bridge/model.py TextBridge) 的关键区别: 不做 e_cv 蒸馏, 训练信号只有 codec CE (teacher-forcing 穿过冻结 CV3 LLM), 允许表示离开文本嵌入流形。 设计: 零初始化残差 —— out 层零初始化, 初始 e_hat == e (恰在 CV3 文本嵌入流形上, val_ce 起点应等于冻结 LLM 自身地板 ~2.95), 训练逐步学习 off-manifold 偏移 delta。 结构: LN(h)->proj + LN(e)->proj + 可学习位置 -> pre-norm 双向 Transformer -> 零初始化输出层。 """ import torch from torch import nn class _CausalTransformerBlock(nn.Module): """Small causal Transformer block with an explicit KV cache. ``nn.TransformerEncoder`` can apply a causal mask for a complete sequence, but it does not expose the key/value tensors needed by a long-lived streaming session. This block keeps those tensors in the mapper state and therefore has exactly the same computation when a sequence is supplied in one call or in several chunks. """ def __init__(self, d_model, n_head, d_ff, dropout): super().__init__() if d_model % n_head: raise ValueError(f'd_model ({d_model}) must be divisible by n_head ({n_head})') self.d_model = d_model self.n_head = n_head self.head_dim = d_model // n_head self.norm1 = nn.LayerNorm(d_model) self.qkv = nn.Linear(d_model, 3 * d_model) self.proj = nn.Linear(d_model, d_model) self.attn_drop = nn.Dropout(dropout) self.norm2 = nn.LayerNorm(d_model) self.ff = nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model), nn.Dropout(dropout), ) def _split_heads(self, x): # (B,T,D) -> (B,H,T,D/H) b, t, _ = x.shape return x.view(b, t, self.n_head, self.head_dim).transpose(1, 2) def forward(self, x, cache=None, key_padding_mask=None): """Run a chunk and return ``(output, new_cache)``. ``cache`` is ``(key, value, key_padding_mask)``. The mask component is optional and is only used when callers feed padded chunks. A regular streaming caller can leave it as ``None``. """ b, t, _ = x.shape qkv = self.qkv(self.norm1(x)).chunk(3, dim=-1) q, k_new, v_new = (self._split_heads(v) for v in qkv) past_k = past_v = past_kpm = None if cache is not None: # Keep accepting a two-tensor cache for callers that construct one # themselves; the third item was added for padded training batches. past_k, past_v = cache[:2] if len(cache) > 2: past_kpm = cache[2] past_len = 0 if past_k is None else past_k.shape[2] k = k_new if past_k is None else torch.cat([past_k, k_new], dim=2) v = v_new if past_v is None else torch.cat([past_v, v_new], dim=2) # Query positions are [past_len, past_len + t). A chunk must not # attend to a future token from that same chunk. total_len = k.shape[2] q_pos = torch.arange(past_len, past_len + t, device=x.device) k_pos = torch.arange(total_len, device=x.device) future = k_pos.unsqueeze(0) > q_pos.unsqueeze(1) # (T,total_len) cur_kpm = None if key_padding_mask is not None: cur_kpm = key_padding_mask.to(torch.bool) if cur_kpm.ndim != 2 or cur_kpm.shape != (b, t): raise ValueError( 'key_padding_mask must have shape (batch, chunk_length)') if past_len: # A cache created without a padding mask still represents valid # (non-padding) prefix rows. Materialize those zeros when a later # chunk supplies a mask so the mask always spans all keys. if past_kpm is None: past_kpm = torch.zeros( b, past_len, dtype=torch.bool, device=x.device) if cur_kpm is None: cur_kpm = torch.zeros(b, t, dtype=torch.bool, device=x.device) all_kpm = torch.cat([past_kpm, cur_kpm], dim=1) else: all_kpm = cur_kpm mask = future.unsqueeze(0).expand(b, -1, -1) if all_kpm is not None: mask = mask | all_kpm.unsqueeze(1) scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5) # ``-inf`` yields NaNs for a chunk consisting entirely of padding. A # large finite value keeps that corner case well-defined; its output is # masked back to the residual input by the mapper below. scores = scores.masked_fill(mask.unsqueeze(1), -1e4) attn = torch.softmax(scores, dim=-1) attn = self.attn_drop(attn) attn = torch.matmul(attn, v).transpose(1, 2).contiguous() attn = attn.view(b, t, self.d_model) x = x + self.proj(attn) x = x + self.ff(self.norm2(x)) new_kpm = all_kpm return x, (k, v, new_kpm) class Hidden2CVMapper(nn.Module): """(h (B,m,2048), e (B,m,896), kpm (B,m)) -> e_hat (B,m,896)""" def __init__(self, d_h=2048, d_e=896, n_layer=3, n_head=8, d_ff=2048, dropout=0.1, max_len=512, n_lang=0): super().__init__() self.d_e = d_e self.n_lang = n_lang if n_lang > 0: # 语种条件 (v3): 零初始化, 不破坏 sanity 地板 (e_hat==e) self.lang_emb = nn.Embedding(n_lang, d_e) nn.init.zeros_(self.lang_emb.weight) self.h_norm = nn.LayerNorm(d_h) self.h_proj = nn.Linear(d_h, d_e) self.e_norm = nn.LayerNorm(d_e) self.e_proj = nn.Linear(d_e, d_e) self.pos = nn.Parameter(torch.zeros(1, max_len, d_e)) nn.init.trunc_normal_(self.pos, std=0.02) layer = nn.TransformerEncoderLayer( d_e, n_head, d_ff, dropout, activation='gelu', batch_first=True, norm_first=True) self.encoder = nn.TransformerEncoder(layer, n_layer) self.out = nn.Linear(d_e, d_e) nn.init.zeros_(self.out.weight) nn.init.zeros_(self.out.bias) def forward(self, h, e, kpm=None, lang=None): """h (B,m,d_h) float, e (B,m,d_e) float, kpm (B,m) bool True=pad, lang (B,) long 语种 id (n_lang>0 时必传) -> e_hat (B,m,d_e)""" m = h.shape[1] x = self.h_proj(self.h_norm(h)) + self.e_proj(self.e_norm(e)) + self.pos[:, :m] if self.n_lang > 0: assert lang is not None, 'n_lang>0 requires lang ids' x = x + self.lang_emb(lang).unsqueeze(1) x = self.encoder(x, src_key_padding_mask=kpm) return e + self.out(x) class CausalHidden2CVMapper(nn.Module): """Causal, stateful hidden-to-CosyVoice embedding mapper. The public input/output contract matches :class:`Hidden2CVMapper`, while ``forward_incremental`` and ``step`` keep a per-layer KV cache for a persistent streaming session:: y_offline = mapper(h, e) y_a, state = mapper.forward_incremental(h[:, :8], e[:, :8]) y_b, state = mapper.forward_incremental(h[:, 8:], e[:, 8:], state) assert torch.allclose(torch.cat([y_a, y_b], 1), y_offline) In residual mode the final projection is zero-initialized, so a new mapper is an identity adapter (``e_hat == e``), matching the initialization and checkpoint sanity behavior of ``Hidden2CVMapper``. ``forward`` and the incremental path share the same block implementation, so there is no train/inference-only causal mask discrepancy. For exact numerical parity between full and chunked calls, use ``eval()`` or ``dropout=0``. ``state`` is an opaque dictionary containing ``position`` and one ``(key, value, key_padding_mask)`` cache tuple per layer. Callers should pass it back unchanged to the next invocation; ``reset_state`` starts a fresh sequence. ``input_mode='h_only'`` removes the native e projections and residual. It uses only h, language and position to predict the CV embedding; passing e then supplies shape/dtype compatibility but never its values. """ def __init__(self, d_h=2048, d_e=896, n_layer=3, n_head=8, d_ff=2048, dropout=0.1, max_len=512, n_lang=0, input_mode='residual'): super().__init__() if input_mode not in ('residual', 'h_only'): raise ValueError("input_mode must be 'residual' or 'h_only'") self.input_mode = input_mode self.d_h = d_h self.d_e = d_e self.n_layer = n_layer self.n_head = n_head self.d_ff = d_ff self.dropout = dropout self.max_len = max_len self.n_lang = n_lang if n_lang > 0: # Keep language conditioning zero-initialized so adding it cannot # alter the e_hat == e initialization sanity check. self.lang_emb = nn.Embedding(n_lang, d_e) nn.init.zeros_(self.lang_emb.weight) self.h_norm = nn.LayerNorm(d_h) self.h_proj = nn.Linear(d_h, d_e) if input_mode == 'residual': self.e_norm = nn.LayerNorm(d_e) self.e_proj = nn.Linear(d_e, d_e) self.pos = nn.Parameter(torch.zeros(1, max_len, d_e)) nn.init.trunc_normal_(self.pos, std=0.02) self.layers = nn.ModuleList([ _CausalTransformerBlock(d_e, n_head, d_ff, dropout) for _ in range(n_layer) ]) self.out = nn.Linear(d_e, d_e) if input_mode == 'residual': nn.init.zeros_(self.out.weight) else: # h-only arms must not have an embedding lookup/residual shortcut. nn.init.xavier_uniform_(self.out.weight) nn.init.zeros_(self.out.bias) def _position_slice(self, start, length, device, dtype): """Return deterministic position features, including beyond max_len. Existing checkpoints are trained with ``max_len`` learned positions. For a longer streaming sequence, learned positions are used where available and zeros are appended. The recurrent KV context still distinguishes all later tokens and, crucially, 60-second inputs do not fail merely because they exceed the training position table. """ if length == 0: return self.pos.new_zeros((1, 0, self.d_e), device=device, dtype=dtype) end = start + length if start >= self.max_len: return self.pos.new_zeros((1, length, self.d_e), device=device, dtype=dtype) learned_end = min(end, self.max_len) learned = self.pos[:, start:learned_end].to(device=device, dtype=dtype) if learned_end == end: return learned pad = self.pos.new_zeros( (1, end - learned_end, self.d_e), device=device, dtype=dtype) return torch.cat([learned, pad], dim=1) def reset_state(self, batch_size=1, device=None, dtype=None): """Create an empty streaming state for ``batch_size`` sequences.""" if device is None: device = self.pos.device if dtype is None: dtype = self.pos.dtype # KV tensors are created by the first call, so dtype is retained here # only for API symmetry and future callers that inspect the state. del dtype return { 'position': 0, 'kv': [None for _ in self.layers], 'batch_size': int(batch_size), 'device': torch.device(device), } # Alias used by a few streaming wrappers. init_state = reset_state def _validate_inputs(self, h, e): if h.ndim == 2: h = h.unsqueeze(1) if e is None: if self.input_mode != 'h_only': raise ValueError('residual mapper requires native embeddings e') e = h.new_zeros((*h.shape[:2], self.d_e)) if e.ndim == 2: e = e.unsqueeze(1) if h.ndim != 3 or e.ndim != 3: raise ValueError('h and e must have shape (B,T,D) or (B,D)') if h.shape[:2] != e.shape[:2]: raise ValueError('h and e must have the same batch and sequence dimensions') if h.shape[-1] != self.d_h: raise ValueError(f'h last dimension must be {self.d_h}, got {h.shape[-1]}') if e.shape[-1] != self.d_e: raise ValueError(f'e last dimension must be {self.d_e}, got {e.shape[-1]}') return h, e def _prepare_state(self, state, batch_size, device): if state is None: return self.reset_state(batch_size, device=device) if not isinstance(state, dict): raise TypeError('state must be a state dictionary returned by the mapper') kv = state.get('kv') if kv is None or len(kv) != len(self.layers): raise ValueError('state has an invalid per-layer KV cache') if state.get('batch_size', batch_size) != batch_size: raise ValueError('state batch size does not match h/e') if state.get('position', 0) < 0: raise ValueError('state position must be non-negative') # Do not mutate the caller's dictionary. The cache tensors themselves # are immutable references until a layer concatenates a new chunk. return { 'position': int(state.get('position', 0)), 'kv': list(kv), 'batch_size': batch_size, 'device': torch.device(device), } def forward_incremental(self, h, e=None, state=None, lang=None, kpm=None): """Map one sequence/chunk and return ``(e_hat, new_state)``. Args: h: ST hidden states, ``(B,T,d_h)`` or one-token ``(B,d_h)``. e: CV3 text embeddings, ``(B,T,d_e)`` or ``(B,d_e)``. state: State returned by an earlier call, or ``None`` to reset. lang: Optional ``(B,)`` language ids when ``n_lang > 0``. kpm: Optional ``(B,T)`` padding mask (``True`` means padding). """ h, e = self._validate_inputs(h, e) b, t, _ = h.shape if kpm is not None: if kpm.shape != (b, t): raise ValueError('kpm must have shape (batch, sequence_length)') kpm = kpm.to(dtype=torch.bool, device=h.device) if self.n_lang > 0: if lang is None: raise AssertionError('n_lang>0 requires lang ids') if lang.ndim == 0: lang = lang.reshape(1).expand(b) if lang.shape != (b,): raise ValueError('lang must have shape (batch,)') lang = lang.to(device=h.device, dtype=torch.long) st = self._prepare_state(state, b, h.device) if t == 0: return e, st pos = self._position_slice(st['position'], t, h.device, e.dtype) x = self.h_proj(self.h_norm(h)) + pos if self.input_mode == 'residual': x = x + self.e_proj(self.e_norm(e)) if self.n_lang > 0: x = x + self.lang_emb(lang).unsqueeze(1).to(dtype=x.dtype) new_kv = [] for layer, cache in zip(self.layers, st['kv']): x, cache = layer(x, cache=cache, key_padding_mask=kpm) new_kv.append(cache) delta = self.out(x) if kpm is not None: # A padded position must not update either its output or the # semantic prefix represented by the state. delta = delta.masked_fill(kpm.unsqueeze(-1), 0) e_hat = e + delta if self.input_mode == 'residual' else delta st['position'] += t st['kv'] = new_kv return e_hat, st def step(self, h, e=None, state=None, lang=None, kpm=None): """Process one token and return ``(e_hat_t, new_state)``. ``h``/``e`` may be either ``(B,D)`` or ``(B,1,D)``. This is a thin wrapper around ``forward_incremental`` so both entry points have identical cache semantics. """ if h.ndim == 3 and h.shape[1] != 1: raise ValueError('step requires exactly one token') h_hat, st = self.forward_incremental(h, e, state=state, lang=lang, kpm=kpm) return h_hat[:, 0], st def forward(self, h, e=None, kpm=None, lang=None, chunk_sizes=None): """Offline causal forward; equivalent to one incremental call.""" if chunk_sizes is not None: h, e = self._validate_inputs(h, e) if any(size <= 0 for size in chunk_sizes) or sum(chunk_sizes) != h.shape[1]: raise ValueError('chunk_sizes must be positive and cover the complete sequence') state, outputs, start = None, [], 0 for size in chunk_sizes: end = start + size out, state = self.forward_incremental( h[:, start:end], e[:, start:end], state=state, lang=lang, kpm=None if kpm is None else kpm[:, start:end]) outputs.append(out) start = end return torch.cat(outputs, dim=1) if outputs else e e_hat, _ = self.forward_incremental(h, e, state=None, lang=lang, kpm=kpm) return e_hat __all__ = ['Hidden2CVMapper', 'CausalHidden2CVMapper']