from __future__ import annotations import math from dataclasses import asdict import torch import torch.nn as nn import torch.nn.functional as F from .config import ModelConfig DURATION_SPEAKER_FUSIONS = { "concat", "adarn", "adarn_zero", "speaker_cross_attn", "text_cross_attn", } DURATION_ARCHITECTURES = {"pooled", "token_sum_adarn_zero_no_aux"} def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> torch.Tensor: freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) t = torch.arange(end, dtype=torch.float32) freqs = torch.outer(t, freqs) return torch.complex(torch.cos(freqs), torch.sin(freqs)) def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: # x: (B, S, H, Dh), Dh must be even. x_ = torch.view_as_complex(x.float().reshape(*x.shape[:3], -1, 2)) x_ = x_ * freqs_cis[None, :, None, :] x_ = torch.view_as_real(x_).reshape_as(x) return x_.type_as(x) def get_timestep_embedding(timestep: torch.Tensor, dim: int) -> torch.Tensor: assert dim % 2 == 0 half = dim // 2 freqs = 1000.0 * torch.exp( -torch.log(torch.tensor(10000.0, device=timestep.device, dtype=torch.float32)) * torch.arange(half, device=timestep.device, dtype=torch.float32) / half ) args = timestep[:, None].float() * freqs[None, :] return torch.cat([torch.cos(args), torch.sin(args)], dim=-1).to(timestep.dtype) class RMSNorm(nn.Module): def __init__(self, dim: int | tuple[int, ...], eps: float = 1e-6): super().__init__() if isinstance(dim, int): dim = (dim,) self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x: torch.Tensor) -> torch.Tensor: x_dtype = x.dtype x = x.float() x = x * torch.rsqrt((x * x).mean(dim=-1, keepdim=True) + self.eps) return (x * self.weight).to(x_dtype) class LowRankAdaLN(nn.Module): """ Echo-style low-rank AdaLN that returns both modulated activations and a residual gate. """ def __init__(self, model_dim: int, rank: int, eps: float): super().__init__() rank = max(1, min(int(rank), int(model_dim))) self.eps = eps self.shift_down = nn.Linear(model_dim, rank, bias=False) self.scale_down = nn.Linear(model_dim, rank, bias=False) self.gate_down = nn.Linear(model_dim, rank, bias=False) self.shift_up = nn.Linear(rank, model_dim, bias=True) self.scale_up = nn.Linear(rank, model_dim, bias=True) self.gate_up = nn.Linear(rank, model_dim, bias=True) # Match Echo/JAX AdaLN behavior: zero-init output projections. nn.init.zeros_(self.shift_up.weight) nn.init.zeros_(self.scale_up.weight) nn.init.zeros_(self.gate_up.weight) if self.shift_up.bias is not None: nn.init.zeros_(self.shift_up.bias) if self.scale_up.bias is not None: nn.init.zeros_(self.scale_up.bias) if self.gate_up.bias is not None: nn.init.zeros_(self.gate_up.bias) def forward( self, x: torch.Tensor, cond_embed: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: shift, scale, gate = cond_embed.chunk(3, dim=-1) shift = self.shift_up(self.shift_down(F.silu(shift))) + shift scale = self.scale_up(self.scale_down(F.silu(scale))) + scale gate = self.gate_up(self.gate_down(F.silu(gate))) + gate x_dtype = x.dtype x = x.float() x = x * torch.rsqrt((x * x).mean(dim=-1, keepdim=True) + self.eps) x = x * (1.0 + scale) + shift gate = torch.tanh(gate) return x.to(x_dtype), gate def patch_sequence_with_mask( seq: torch.Tensor, mask: torch.Tensor, patch_size: int, ) -> tuple[torch.Tensor, torch.Tensor]: """ Patch along sequence axis: seq: (B, S, D) -> (B, S//patch, D*patch) mask: (B, S) -> (B, S//patch) with all() over patch window. Note: For speaker conditioning in this project, `seq` is already in latent-patched space (D = latent_dim * latent_patch_size). This helper applies an additional sequence patching for `speaker_patch_size`. """ if patch_size <= 1: return seq, mask if seq.ndim != 3 or mask.ndim != 2: raise ValueError( f"Expected seq=(B,S,D), mask=(B,S), got seq={tuple(seq.shape)} mask={tuple(mask.shape)}" ) if seq.shape[0] != mask.shape[0] or seq.shape[1] != mask.shape[1]: raise ValueError( f"Sequence/mask shape mismatch: seq={tuple(seq.shape)}, mask={tuple(mask.shape)}. " "Expected matching (B,S)." ) bsz, seq_len, dim = seq.shape usable = (seq_len // patch_size) * patch_size if usable <= 0: raise ValueError( f"Reference sequence too short for speaker_patch_size={patch_size}: seq_len={seq_len}" ) seq = seq[:, :usable].reshape(bsz, usable // patch_size, dim * patch_size) mask = mask[:, :usable].reshape(bsz, usable // patch_size, patch_size).all(dim=-1) return seq, mask class SelfAttention(nn.Module): def __init__(self, dim: int, heads: int, norm_eps: float): super().__init__() if dim % heads != 0: raise ValueError(f"dim={dim} must be divisible by heads={heads}") if (dim // heads) % 2 != 0: raise ValueError("head_dim must be even for RoPE") self.dim = dim self.heads = heads self.head_dim = dim // heads self.wq = nn.Linear(dim, dim, bias=False) self.wk = nn.Linear(dim, dim, bias=False) self.wv = nn.Linear(dim, dim, bias=False) self.wo = nn.Linear(dim, dim, bias=False) self.gate = nn.Linear(dim, dim, bias=False) self.q_norm = RMSNorm((self.heads, self.head_dim), eps=norm_eps) self.k_norm = RMSNorm((self.heads, self.head_dim), eps=norm_eps) def forward( self, x: torch.Tensor, key_mask: torch.Tensor | None, freqs_cis: torch.Tensor, ) -> torch.Tensor: bsz, seq_len, _ = x.shape q = self.wq(x).reshape(bsz, seq_len, self.heads, self.head_dim) k = self.wk(x).reshape(bsz, seq_len, self.heads, self.head_dim) v = self.wv(x).reshape(bsz, seq_len, self.heads, self.head_dim) gate = self.gate(x) q = self.q_norm(q) k = self.k_norm(k) q = apply_rotary_emb(q, freqs_cis[:seq_len]) k = apply_rotary_emb(k, freqs_cis[:seq_len]) attn_mask = None if key_mask is not None: attn_mask = key_mask[:, None, None, :] y = F.scaled_dot_product_attention( q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_mask=attn_mask, is_causal=False, ).transpose(1, 2) y = y.reshape(bsz, seq_len, self.dim) y = y * torch.sigmoid(gate) return self.wo(y) class JointAttention(nn.Module): """ Echo-style joint attention over latent self tokens + conditioning contexts. """ def __init__( self, dim: int, heads: int, text_ctx_dim: int, speaker_ctx_dim: int | None, caption_ctx_dim: int | None, norm_eps: float, ): super().__init__() if dim % heads != 0: raise ValueError(f"dim={dim} must be divisible by heads={heads}") if (dim // heads) % 2 != 0: raise ValueError("head_dim must be even for RoPE") self.dim = dim self.heads = heads self.head_dim = dim // heads self.wq = nn.Linear(dim, dim, bias=False) self.wk = nn.Linear(dim, dim, bias=False) self.wv = nn.Linear(dim, dim, bias=False) self.wk_text = nn.Linear(text_ctx_dim, dim, bias=False) self.wv_text = nn.Linear(text_ctx_dim, dim, bias=False) self.has_speaker_condition = speaker_ctx_dim is not None if self.has_speaker_condition: self.wk_speaker = nn.Linear(int(speaker_ctx_dim), dim, bias=False) self.wv_speaker = nn.Linear(int(speaker_ctx_dim), dim, bias=False) self.has_caption_condition = caption_ctx_dim is not None if self.has_caption_condition: self.wk_caption = nn.Linear(int(caption_ctx_dim), dim, bias=False) self.wv_caption = nn.Linear(int(caption_ctx_dim), dim, bias=False) self.gate = nn.Linear(dim, dim, bias=False) self.wo = nn.Linear(dim, dim, bias=False) self.q_norm = RMSNorm((self.heads, self.head_dim), eps=norm_eps) self.k_norm = RMSNorm((self.heads, self.head_dim), eps=norm_eps) def _apply_rotary_half(self, x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: x_rot, x_passthrough = x.chunk(2, dim=-2) x_rot = apply_rotary_emb(x_rot, freqs_cis) return torch.cat([x_rot, x_passthrough], dim=-2) def project_context_kv( self, text_context: torch.Tensor, speaker_context: torch.Tensor | None, caption_context: torch.Tensor | None = None, ) -> tuple[torch.Tensor, ...]: """ Precompute conditioning KV projections for static conditioning. """ bsz = text_context.shape[0] k_text = self.wk_text(text_context).reshape( bsz, text_context.shape[1], self.heads, self.head_dim ) v_text = self.wv_text(text_context).reshape( bsz, text_context.shape[1], self.heads, self.head_dim ) k_text = self.k_norm(k_text) projected: list[torch.Tensor] = [k_text, v_text] if self.has_speaker_condition: if speaker_context is None: raise ValueError( "speaker_context is required when speaker conditioning is enabled." ) if speaker_context.shape[0] != bsz: raise ValueError( "Batch mismatch for context projection: " f"text={tuple(text_context.shape)} speaker={tuple(speaker_context.shape)}" ) k_speaker = self.wk_speaker(speaker_context).reshape( bsz, speaker_context.shape[1], self.heads, self.head_dim ) v_speaker = self.wv_speaker(speaker_context).reshape( bsz, speaker_context.shape[1], self.heads, self.head_dim ) k_speaker = self.k_norm(k_speaker) projected.extend([k_speaker, v_speaker]) elif speaker_context is not None and speaker_context.shape[0] != bsz: raise ValueError( "Batch mismatch for ignored speaker context: " f"text={tuple(text_context.shape)} speaker={tuple(speaker_context.shape)}" ) if not self.has_caption_condition: return tuple(projected) if caption_context is None: raise ValueError("caption_context is required when caption conditioning is enabled.") if caption_context.shape[0] != bsz: raise ValueError( "Batch mismatch for caption context projection: " f"text={tuple(text_context.shape)} caption={tuple(caption_context.shape)}" ) k_caption = self.wk_caption(caption_context).reshape( bsz, caption_context.shape[1], self.heads, self.head_dim ) v_caption = self.wv_caption(caption_context).reshape( bsz, caption_context.shape[1], self.heads, self.head_dim ) k_caption = self.k_norm(k_caption) projected.extend([k_caption, v_caption]) return tuple(projected) def forward( self, x: torch.Tensor, text_context: torch.Tensor, text_mask: torch.Tensor | None, speaker_context: torch.Tensor | None, speaker_mask: torch.Tensor | None, caption_context: torch.Tensor | None, caption_mask: torch.Tensor | None, freqs_cis: torch.Tensor, self_mask: torch.Tensor | None = None, context_kv: tuple[torch.Tensor, ...] | None = None, ) -> torch.Tensor: bsz, seq_len, _ = x.shape q = self.wq(x).reshape(bsz, seq_len, self.heads, self.head_dim) k_self = self.wk(x).reshape(bsz, seq_len, self.heads, self.head_dim) v_self = self.wv(x).reshape(bsz, seq_len, self.heads, self.head_dim) if context_kv is None: projected = self.project_context_kv( text_context=text_context, speaker_context=speaker_context, caption_context=caption_context, ) else: projected = context_kv if projected is None: raise RuntimeError("JointAttention projected context unexpectedly missing.") offset = 0 k_text, v_text = projected[offset], projected[offset + 1] offset += 2 k_speaker = None v_speaker = None if self.has_speaker_condition: k_speaker, v_speaker = projected[offset], projected[offset + 1] offset += 2 k_caption = None v_caption = None if self.has_caption_condition: k_caption, v_caption = projected[offset], projected[offset + 1] q = self.q_norm(q) k_self = self.k_norm(k_self) q = self._apply_rotary_half(q, freqs_cis[:seq_len]) k_self = self._apply_rotary_half(k_self, freqs_cis[:seq_len]) if self_mask is None: self_mask = torch.ones((bsz, seq_len), dtype=torch.bool, device=x.device) if text_mask is None: text_mask = torch.ones( (bsz, text_context.shape[1]), dtype=torch.bool, device=x.device, ) context_k = [k_self, k_text] context_v = [v_self, v_text] context_masks = [self_mask, text_mask] if self.has_speaker_condition: if speaker_context is None or k_speaker is None or v_speaker is None: raise ValueError( "speaker_context is required when speaker conditioning is enabled." ) if speaker_mask is None: speaker_mask = torch.ones( (bsz, speaker_context.shape[1]), dtype=torch.bool, device=x.device, ) context_k.append(k_speaker) context_v.append(v_speaker) context_masks.append(speaker_mask) if self.has_caption_condition: if caption_context is None: raise ValueError( "caption_context is required when caption conditioning is enabled." ) if caption_mask is None: caption_mask = torch.ones( (bsz, caption_context.shape[1]), dtype=torch.bool, device=x.device, ) if k_caption is None or v_caption is None: raise RuntimeError( "Caption projections are missing despite enabled caption conditioning." ) context_k.append(k_caption) context_v.append(v_caption) context_masks.append(caption_mask) k = torch.cat(context_k, dim=1) v = torch.cat(context_v, dim=1) attn_mask = torch.cat(context_masks, dim=1) attn_mask = attn_mask[:, None, None, :] y = F.scaled_dot_product_attention( q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_mask=attn_mask, is_causal=False, ).transpose(1, 2) y = y.reshape(bsz, seq_len, self.dim) y = y * torch.sigmoid(self.gate(x)) return self.wo(y) class SwiGLU(nn.Module): def __init__(self, dim: int, hidden_dim: int): super().__init__() self.w1 = nn.Linear(dim, hidden_dim, bias=False) self.w2 = nn.Linear(hidden_dim, dim, bias=False) self.w3 = nn.Linear(dim, hidden_dim, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.w2(F.silu(self.w1(x)) * self.w3(x)) def _safe_attention_mask( x: torch.Tensor, mask: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: if mask.ndim != 2 or mask.shape[0] != x.shape[0] or mask.shape[1] != x.shape[1]: raise ValueError( f"mask must have shape (B, S) matching x, got x={tuple(x.shape)} " f"mask={tuple(mask.shape)}" ) mask = mask.to(device=x.device, dtype=torch.bool) has_any = mask.any(dim=1) if bool(has_any.all()): return x, mask if x.shape[1] <= 0: raise ValueError("Cannot attention-pool an empty sequence.") x = x.clone() mask = mask.clone() x[~has_any] = 0 mask[~has_any, 0] = True return x, mask class AttentionPooling(nn.Module): def __init__(self, dim: int, heads: int, norm_eps: float): super().__init__() if dim % heads != 0: raise ValueError(f"dim={dim} must be divisible by heads={heads}") self.dim = int(dim) self.heads = int(heads) self.head_dim = int(dim) // int(heads) self.query = nn.Parameter(torch.empty(1, 1, int(dim))) nn.init.normal_(self.query, mean=0.0, std=0.02) self.q_norm = RMSNorm(dim, eps=norm_eps) self.k_norm = RMSNorm(dim, eps=norm_eps) self.wq = nn.Linear(dim, dim, bias=False) self.wk = nn.Linear(dim, dim, bias=False) self.wv = nn.Linear(dim, dim, bias=False) self.wo = nn.Linear(dim, dim, bias=False) def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: if x.ndim != 3 or x.shape[-1] != self.dim: raise ValueError(f"x must have shape (B, S, {self.dim}), got {tuple(x.shape)}") x, mask = _safe_attention_mask(x, mask) bsz, seq_len, _ = x.shape q = self.query.to(dtype=x.dtype).expand(bsz, -1, -1) q = self.wq(self.q_norm(q)).reshape(bsz, 1, self.heads, self.head_dim) k = self.wk(self.k_norm(x)).reshape(bsz, seq_len, self.heads, self.head_dim) v = self.wv(x).reshape(bsz, seq_len, self.heads, self.head_dim) y = F.scaled_dot_product_attention( q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_mask=mask[:, None, None, :], is_causal=False, ) y = y.transpose(1, 2).reshape(bsz, 1, self.dim) return self.wo(y).squeeze(1) class CrossAttentionPooling(nn.Module): def __init__( self, *, query_dim: int, context_dim: int, output_dim: int, heads: int, norm_eps: float, ): super().__init__() if output_dim % heads != 0: raise ValueError(f"output_dim={output_dim} must be divisible by heads={heads}") self.query_dim = int(query_dim) self.context_dim = int(context_dim) self.output_dim = int(output_dim) self.heads = int(heads) self.head_dim = int(output_dim) // int(heads) self.q_norm = RMSNorm(query_dim, eps=norm_eps) self.k_norm = RMSNorm(context_dim, eps=norm_eps) self.wq = nn.Linear(query_dim, output_dim, bias=False) self.wk = nn.Linear(context_dim, output_dim, bias=False) self.wv = nn.Linear(context_dim, output_dim, bias=False) self.wo = nn.Linear(output_dim, output_dim, bias=False) def forward( self, query: torch.Tensor, context: torch.Tensor, context_mask: torch.Tensor, ) -> torch.Tensor: if query.ndim != 2 or query.shape[-1] != self.query_dim: raise ValueError( f"query must have shape (B, {self.query_dim}), got {tuple(query.shape)}" ) if context.ndim != 3 or context.shape[-1] != self.context_dim: raise ValueError( f"context must have shape (B, S, {self.context_dim}), got {tuple(context.shape)}" ) context, context_mask = _safe_attention_mask(context, context_mask) bsz, seq_len, _ = context.shape q = query[:, None, :] q = self.wq(self.q_norm(q)).reshape(bsz, 1, self.heads, self.head_dim) k = self.wk(self.k_norm(context)).reshape(bsz, seq_len, self.heads, self.head_dim) v = self.wv(context).reshape(bsz, seq_len, self.heads, self.head_dim) y = F.scaled_dot_product_attention( q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_mask=context_mask[:, None, None, :], is_causal=False, ) y = y.transpose(1, 2).reshape(bsz, 1, self.output_dim) return self.wo(y).squeeze(1) class DurationSwiGLUBlock(nn.Module): def __init__( self, *, dim: int, hidden_dim: int, dropout: float, norm_eps: float, cond_dim: int | None = None, ): super().__init__() self.norm = RMSNorm(dim, eps=norm_eps) self.mlp = SwiGLU(dim, hidden_dim) self.dropout = nn.Dropout(dropout) self.cond_dim = cond_dim self.modulation = None if cond_dim is not None: self.modulation = nn.Linear(cond_dim, dim * 3, bias=True) nn.init.zeros_(self.modulation.weight) nn.init.zeros_(self.modulation.bias) def forward(self, x: torch.Tensor, cond: torch.Tensor | None = None) -> torch.Tensor: h = self.norm(x) if self.modulation is not None: if cond is None: raise ValueError("cond is required for AdaRN-Zero duration blocks.") shift, scale, gate = self.modulation(F.silu(cond)).chunk(3, dim=-1) if h.ndim == 3 and shift.ndim == 2: shift = shift.unsqueeze(1) scale = scale.unsqueeze(1) gate = gate.unsqueeze(1) h = h * (1.0 + scale) + shift return x + self.dropout(torch.tanh(gate) * self.mlp(h)) return x + self.dropout(self.mlp(h)) class TextBlock(nn.Module): def __init__(self, dim: int, heads: int, mlp_ratio: float, norm_eps: float, dropout: float): super().__init__() self.attention_norm = RMSNorm(dim, eps=norm_eps) self.attention = SelfAttention(dim, heads, norm_eps=norm_eps) self.mlp_norm = RMSNorm(dim, eps=norm_eps) self.mlp = SwiGLU(dim, int(dim * mlp_ratio)) self.dropout = nn.Dropout(dropout) def forward(self, x: torch.Tensor, mask: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: x = x + self.dropout( self.attention(self.attention_norm(x), key_mask=mask, freqs_cis=freqs_cis) ) x = x + self.dropout(self.mlp(self.mlp_norm(x))) return x class TextEncoder(nn.Module): def __init__( self, *, vocab_size: int, dim: int, layers: int, heads: int, mlp_ratio: float, norm_eps: float, dropout: float, ): super().__init__() self.text_embedding = nn.Embedding(vocab_size, dim) self.blocks = nn.ModuleList( TextBlock( dim=dim, heads=heads, mlp_ratio=mlp_ratio, norm_eps=norm_eps, dropout=dropout, ) for _ in range(layers) ) self.head_dim = dim // heads self.register_buffer( "_freqs_cis_cache", torch.empty(0, 0, dtype=torch.complex64), persistent=False ) def _rope_freqs(self, seq_len: int, device: torch.device) -> torch.Tensor: cache = self._freqs_cis_cache if cache.device != device or cache.shape[0] < seq_len: cache = precompute_freqs_cis(self.head_dim, seq_len).to(device) self._freqs_cis_cache = cache return cache[:seq_len] def forward(self, input_ids: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: x = self.text_embedding(input_ids) # Hard-mask invalid tokens so fully-masked conditioning becomes truly unconditional. mask_f = mask.unsqueeze(-1).to(dtype=x.dtype) x = x * mask_f freqs = self._rope_freqs(input_ids.shape[1], x.device) for block in self.blocks: x = block(x, mask=mask, freqs_cis=freqs) x = x * mask_f return x * mask_f class ReferenceLatentEncoder(nn.Module): """ Encoder for reference latents used as speaker/style conditioning. """ def __init__(self, cfg: ModelConfig): super().__init__() self.in_proj = nn.Linear(cfg.speaker_patched_latent_dim, cfg.speaker_dim, bias=True) speaker_mlp_ratio = cfg.speaker_mlp_ratio_resolved self.blocks = nn.ModuleList( TextBlock( dim=cfg.speaker_dim, heads=cfg.speaker_heads, mlp_ratio=speaker_mlp_ratio, norm_eps=cfg.norm_eps, dropout=cfg.dropout, ) for _ in range(cfg.speaker_layers) ) self.head_dim = cfg.speaker_dim // cfg.speaker_heads self.register_buffer( "_freqs_cis_cache", torch.empty(0, 0, dtype=torch.complex64), persistent=False ) def _rope_freqs(self, seq_len: int, device: torch.device) -> torch.Tensor: cache = self._freqs_cis_cache if cache.device != device or cache.shape[0] < seq_len: cache = precompute_freqs_cis(self.head_dim, seq_len).to(device) self._freqs_cis_cache = cache return cache[:seq_len] def forward(self, latent: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: x = self.in_proj(latent) x = x / 6.0 # Keep masked reference positions strictly zero across residual/MLP paths. mask_f = mask.unsqueeze(-1).to(dtype=x.dtype) x = x * mask_f freqs = self._rope_freqs(x.shape[1], x.device) for block in self.blocks: x = block(x, mask=mask, freqs_cis=freqs) x = x * mask_f return x * mask_f class DiffusionBlock(nn.Module): def __init__(self, cfg: ModelConfig): super().__init__() self.attention = JointAttention( cfg.model_dim, cfg.num_heads, cfg.text_dim, cfg.speaker_dim if cfg.use_speaker_condition else None, cfg.caption_dim_resolved if cfg.use_caption_condition else None, norm_eps=cfg.norm_eps, ) self.mlp = SwiGLU(cfg.model_dim, int(cfg.model_dim * cfg.mlp_ratio)) adaln_rank = max(1, min(int(cfg.adaln_rank), int(cfg.model_dim))) self.attention_adaln = LowRankAdaLN( model_dim=cfg.model_dim, rank=adaln_rank, eps=cfg.norm_eps, ) self.mlp_adaln = LowRankAdaLN( model_dim=cfg.model_dim, rank=adaln_rank, eps=cfg.norm_eps, ) self.dropout = nn.Dropout(cfg.dropout) def forward( self, x: torch.Tensor, cond_embed: torch.Tensor, text_state: torch.Tensor, text_mask: torch.Tensor, speaker_state: torch.Tensor | None, speaker_mask: torch.Tensor | None, caption_state: torch.Tensor | None, caption_mask: torch.Tensor | None, freqs_cis: torch.Tensor, self_mask: torch.Tensor | None = None, context_kv: tuple[torch.Tensor, ...] | None = None, ) -> torch.Tensor: h, attention_gate = self.attention_adaln(x, cond_embed) x = x + self.dropout( attention_gate * self.attention( x=h, text_context=text_state, text_mask=text_mask, speaker_context=speaker_state, speaker_mask=speaker_mask, caption_context=caption_state, caption_mask=caption_mask, freqs_cis=freqs_cis, self_mask=self_mask, context_kv=context_kv, ) ) h, mlp_gate = self.mlp_adaln(x, cond_embed) x = x + self.dropout(mlp_gate * self.mlp(h)) return x class DurationPredictor(nn.Module): def __init__( self, *, text_dim: int, aux_dim: int, hidden_dim: int, layers: int, dropout: float, speaker_dim: int | None = None, speaker_fusion: str = "concat", attention_heads: int = 8, norm_eps: float = 1e-5, architecture: str = "pooled", token_init_frames: float = 6.3, ): super().__init__() if text_dim <= 0: raise ValueError(f"duration predictor text_dim must be > 0, got {text_dim}") if aux_dim <= 0: raise ValueError(f"duration predictor aux_dim must be > 0, got {aux_dim}") if hidden_dim <= 0: raise ValueError(f"duration predictor hidden_dim must be > 0, got {hidden_dim}") if layers <= 0: raise ValueError(f"duration predictor layers must be > 0, got {layers}") if speaker_dim is not None and speaker_dim <= 0: raise ValueError(f"duration predictor speaker_dim must be > 0, got {speaker_dim}") speaker_fusion = str(speaker_fusion).strip().lower() if speaker_fusion not in DURATION_SPEAKER_FUSIONS: raise ValueError( f"duration speaker fusion must be one of {sorted(DURATION_SPEAKER_FUSIONS)}, " f"got {speaker_fusion!r}" ) architecture = str(architecture).strip().lower() if architecture not in DURATION_ARCHITECTURES: raise ValueError( "duration architecture must be one of " f"{sorted(DURATION_ARCHITECTURES)}, got {architecture!r}" ) if attention_heads <= 0: raise ValueError( f"duration predictor attention_heads must be > 0, got {attention_heads}" ) if token_init_frames <= 0: raise ValueError( f"duration token_init_frames must be > 0, got {token_init_frames}" ) if speaker_dim is None and speaker_fusion != "concat": raise ValueError( f"duration speaker fusion {speaker_fusion!r} requires speaker_dim." ) if architecture == "token_sum_adarn_zero_no_aux" and speaker_dim is None: raise ValueError("token_sum_adarn_zero_no_aux requires speaker_dim.") if architecture == "token_sum_adarn_zero_no_aux" and speaker_fusion != "adarn_zero": raise ValueError( "token_sum_adarn_zero_no_aux uses block-level speaker AdaRN-Zero and " "requires speaker_fusion='adarn_zero'." ) self.text_dim = int(text_dim) self.aux_dim = int(aux_dim) self.hidden_dim = int(hidden_dim) self.speaker_dim = None if speaker_dim is None else int(speaker_dim) self.speaker_fusion = speaker_fusion self.duration_architecture = architecture self.text_pool = None self.null_speaker = ( nn.Parameter(torch.zeros(int(speaker_dim))) if speaker_dim is not None else None ) self.text_adarn_norm = None self.text_adarn = None self.speaker_cross_attn = None self.text_cross_attn = None self.token_input_proj = None self.token_blocks = None self.token_out_norm = None self.token_out_proj = None if architecture == "token_sum_adarn_zero_no_aux": self.token_input_proj = nn.Linear(int(text_dim), int(hidden_dim)) self.token_blocks = nn.ModuleList( DurationSwiGLUBlock( dim=int(hidden_dim), hidden_dim=int(hidden_dim), dropout=float(dropout), norm_eps=float(norm_eps), cond_dim=int(speaker_dim), ) for _ in range(int(layers)) ) self.token_out_norm = RMSNorm(int(hidden_dim), eps=float(norm_eps)) self.token_out_proj = nn.Linear(int(hidden_dim), 1) nn.init.zeros_(self.token_out_proj.weight) nn.init.constant_( self.token_out_proj.bias, float(math.log(math.expm1(float(token_init_frames)))), ) return self.text_pool = AttentionPooling( dim=int(text_dim), heads=int(attention_heads), norm_eps=float(norm_eps), ) if speaker_dim is not None: if speaker_fusion == "concat": input_dim = int(text_dim) + int(speaker_dim) + int(aux_dim) elif speaker_fusion == "adarn": input_dim = int(text_dim) + int(aux_dim) self.text_adarn_norm = RMSNorm(int(text_dim), eps=float(norm_eps)) self.text_adarn = nn.Linear(int(speaker_dim), int(text_dim) * 2) nn.init.zeros_(self.text_adarn.weight) nn.init.zeros_(self.text_adarn.bias) elif speaker_fusion == "adarn_zero": input_dim = int(text_dim) + int(aux_dim) elif speaker_fusion == "speaker_cross_attn": input_dim = int(text_dim) * 2 + int(aux_dim) self.speaker_cross_attn = CrossAttentionPooling( query_dim=int(text_dim), context_dim=int(speaker_dim), output_dim=int(text_dim), heads=int(attention_heads), norm_eps=float(norm_eps), ) elif speaker_fusion == "text_cross_attn": input_dim = int(text_dim) + int(speaker_dim) + int(aux_dim) self.text_cross_attn = CrossAttentionPooling( query_dim=int(speaker_dim), context_dim=int(text_dim), output_dim=int(text_dim), heads=int(attention_heads), norm_eps=float(norm_eps), ) else: raise RuntimeError(f"Unsupported duration speaker fusion: {speaker_fusion!r}") else: input_dim = int(text_dim) + int(aux_dim) self.input_proj = nn.Linear(int(input_dim), int(hidden_dim)) block_cond_dim = int(speaker_dim) if speaker_fusion == "adarn_zero" else None self.blocks = nn.ModuleList( DurationSwiGLUBlock( dim=int(hidden_dim), hidden_dim=int(hidden_dim), dropout=float(dropout), norm_eps=float(norm_eps), cond_dim=block_cond_dim, ) for _ in range(int(layers)) ) self.out_norm = RMSNorm(int(hidden_dim), eps=float(norm_eps)) self.out_proj = nn.Linear(int(hidden_dim), 1) def _speaker_vec( self, *, batch_size: int, device: torch.device, dtype: torch.dtype, speaker_state: torch.Tensor | None, has_speaker: torch.Tensor, ) -> torch.Tensor: if self.null_speaker is None or self.speaker_dim is None: raise RuntimeError("Duration speaker modules are missing.") null_vec = self.null_speaker.to(device=device, dtype=dtype)[None, :].expand( batch_size, -1 ) if speaker_state is None: return null_vec if speaker_state.ndim != 3 or speaker_state.shape[0] != batch_size: raise ValueError( f"speaker_state must have shape (B, S, D), got {tuple(speaker_state.shape)}" ) if speaker_state.shape[-1] != self.speaker_dim: raise ValueError( f"speaker_state last dim must be {self.speaker_dim}, got {speaker_state.shape[-1]}" ) speaker_vec = speaker_state[:, 0].to(device=device, dtype=dtype) return torch.where(has_speaker[:, None], speaker_vec, null_vec) def _speaker_sequence( self, *, batch_size: int, device: torch.device, dtype: torch.dtype, speaker_state: torch.Tensor | None, speaker_mask: torch.Tensor | None, has_speaker: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: if self.null_speaker is None or self.speaker_dim is None: raise RuntimeError("Duration speaker modules are missing.") null_token = self.null_speaker.to(device=device, dtype=dtype)[None, None, :].expand( batch_size, 1, -1 ) if speaker_state is None: return null_token, torch.ones((batch_size, 1), dtype=torch.bool, device=device) if speaker_state.ndim != 3 or speaker_state.shape[0] != batch_size: raise ValueError( f"speaker_state must have shape (B, S, D), got {tuple(speaker_state.shape)}" ) if speaker_state.shape[-1] != self.speaker_dim: raise ValueError( f"speaker_state last dim must be {self.speaker_dim}, got {speaker_state.shape[-1]}" ) speaker_state = speaker_state.to(device=device, dtype=dtype) if speaker_mask is None: speaker_mask = torch.ones( (batch_size, speaker_state.shape[1]), dtype=torch.bool, device=device ) elif speaker_mask.ndim != 2 or speaker_mask.shape[:2] != speaker_state.shape[:2]: raise ValueError( "speaker_mask must have shape matching speaker_state (B, S), " f"got speaker_state={tuple(speaker_state.shape)} mask={tuple(speaker_mask.shape)}" ) speaker_mask = speaker_mask.to(device=device, dtype=torch.bool) real_mask = speaker_mask & has_speaker[:, None] fallback_mask = ~real_mask.any(dim=1, keepdim=True) context = torch.cat([speaker_state, null_token], dim=1) context_mask = torch.cat([real_mask, fallback_mask], dim=1) return context, context_mask def forward( self, *, text_state: torch.Tensor, text_mask: torch.Tensor, aux_features: torch.Tensor, speaker_state: torch.Tensor | None = None, speaker_mask: torch.Tensor | None = None, has_speaker: torch.Tensor | None = None, ) -> torch.Tensor: if text_state.ndim != 3 or text_state.shape[-1] != self.text_dim: raise ValueError( f"text_state must have shape (B, S, {self.text_dim}), " f"got {tuple(text_state.shape)}" ) if aux_features.ndim != 2 or aux_features.shape[1] != self.aux_dim: raise ValueError( f"aux_features must have shape (B, {self.aux_dim}), " f"got {tuple(aux_features.shape)}" ) if aux_features.shape[0] != text_state.shape[0]: raise ValueError( "Batch mismatch for duration predictor: " f"text_state={tuple(text_state.shape)} aux_features={tuple(aux_features.shape)}" ) text_state, text_mask = _safe_attention_mask(text_state, text_mask) aux_features = aux_features.to(device=text_state.device, dtype=text_state.dtype) if self.duration_architecture == "token_sum_adarn_zero_no_aux": if self.speaker_dim is None: raise RuntimeError("Token-sum duration architecture requires speaker modules.") if has_speaker is None: raise ValueError( "has_speaker is required for speaker-conditioned duration prediction." ) has_speaker = has_speaker.to(device=text_state.device, dtype=torch.bool) if has_speaker.ndim != 1 or has_speaker.shape[0] != text_state.shape[0]: raise ValueError( f"has_speaker must have shape (B,), got {tuple(has_speaker.shape)}" ) speaker_vec = self._speaker_vec( batch_size=text_state.shape[0], device=text_state.device, dtype=text_state.dtype, speaker_state=speaker_state, has_speaker=has_speaker, ) if ( self.token_input_proj is None or self.token_blocks is None or self.token_out_norm is None or self.token_out_proj is None ): raise RuntimeError("Token-sum duration modules are missing.") h = self.token_input_proj(text_state) for block in self.token_blocks: h = block(h, cond=speaker_vec) token_logits = self.token_out_proj(self.token_out_norm(h)).squeeze(-1) token_frames = F.softplus(token_logits.float()) total_frames = (token_frames * text_mask.to(dtype=token_frames.dtype)).sum(dim=1) return torch.log1p(total_frames.clamp_min(0.0)) if self.text_pool is None: raise RuntimeError("Pooled duration modules are missing.") text_vec = self.text_pool(text_state, text_mask) if self.speaker_dim is None: x = torch.cat([text_vec, aux_features], dim=-1) h = self.input_proj(x) for block in self.blocks: h = block(h) return self.out_proj(self.out_norm(h)).squeeze(-1) if has_speaker is None: raise ValueError("has_speaker is required for speaker-conditioned duration prediction.") has_speaker = has_speaker.to(device=text_vec.device, dtype=torch.bool) if has_speaker.ndim != 1 or has_speaker.shape[0] != text_vec.shape[0]: raise ValueError( f"has_speaker must have shape (B,), got {tuple(has_speaker.shape)}" ) speaker_vec = self._speaker_vec( batch_size=text_vec.shape[0], device=text_vec.device, dtype=text_vec.dtype, speaker_state=speaker_state, has_speaker=has_speaker, ) if self.speaker_fusion == "concat": x = torch.cat([text_vec, speaker_vec, aux_features], dim=-1) cond = None elif self.speaker_fusion == "adarn": if self.text_adarn_norm is None or self.text_adarn is None: raise RuntimeError("AdaRN duration speaker modules are missing.") scale, shift = self.text_adarn(speaker_vec).chunk(2, dim=-1) text_vec = (self.text_adarn_norm(text_vec) * (1.0 + scale)) + shift x = torch.cat([text_vec, aux_features], dim=-1) cond = None elif self.speaker_fusion == "adarn_zero": x = torch.cat([text_vec, aux_features], dim=-1) cond = speaker_vec elif self.speaker_fusion == "speaker_cross_attn": if self.speaker_cross_attn is None: raise RuntimeError("speaker_cross_attn duration module is missing.") speaker_context, speaker_context_mask = self._speaker_sequence( batch_size=text_vec.shape[0], device=text_vec.device, dtype=text_vec.dtype, speaker_state=speaker_state, speaker_mask=speaker_mask, has_speaker=has_speaker, ) context_vec = self.speaker_cross_attn( query=text_vec, context=speaker_context, context_mask=speaker_context_mask, ) x = torch.cat([text_vec, context_vec, aux_features], dim=-1) cond = None elif self.speaker_fusion == "text_cross_attn": if self.text_cross_attn is None: raise RuntimeError("text_cross_attn duration module is missing.") context_vec = self.text_cross_attn( query=speaker_vec, context=text_state, context_mask=text_mask, ) x = torch.cat([context_vec, speaker_vec, aux_features], dim=-1) cond = None else: raise RuntimeError(f"Unsupported duration speaker fusion: {self.speaker_fusion!r}") h = self.input_proj(x) for block in self.blocks: h = block(h, cond=cond) return self.out_proj(self.out_norm(h)).squeeze(-1) class TextToLatentRFDiT(nn.Module): """ Text + reference-latent conditioned RF diffusion model over patched DACVAE latent sequences. Input x_t shape: (B, S, latent_dim * latent_patch_size) Output v_pred shape: same as input. """ def __init__(self, cfg: ModelConfig): super().__init__() self.cfg = cfg self.text_encoder = TextEncoder( vocab_size=cfg.text_vocab_size, dim=cfg.text_dim, layers=cfg.text_layers, heads=cfg.text_heads, mlp_ratio=cfg.text_mlp_ratio_resolved, norm_eps=cfg.norm_eps, dropout=cfg.dropout, ) self.caption_encoder = None self.caption_norm = None if cfg.use_caption_condition: self.caption_encoder = TextEncoder( vocab_size=cfg.caption_vocab_size_resolved, dim=cfg.caption_dim_resolved, layers=cfg.caption_layers_resolved, heads=cfg.caption_heads_resolved, mlp_ratio=cfg.caption_mlp_ratio_resolved, norm_eps=cfg.norm_eps, dropout=cfg.dropout, ) self.caption_norm = RMSNorm(cfg.caption_dim_resolved, eps=cfg.norm_eps) self.speaker_encoder = None if cfg.use_speaker_condition: self.speaker_encoder = ReferenceLatentEncoder(cfg) self.text_norm = RMSNorm(cfg.text_dim, eps=cfg.norm_eps) self.speaker_norm = None if cfg.use_speaker_condition: self.speaker_norm = RMSNorm(cfg.speaker_dim, eps=cfg.norm_eps) self.duration_predictor = None if cfg.use_duration_predictor: duration_speaker_dim = None if cfg.use_speaker_condition: duration_speaker_dim = int(cfg.speaker_dim) self.duration_predictor = DurationPredictor( text_dim=cfg.text_dim, aux_dim=cfg.duration_aux_dim, hidden_dim=cfg.duration_hidden_dim, layers=cfg.duration_layers, dropout=cfg.duration_dropout, speaker_dim=duration_speaker_dim, speaker_fusion=cfg.duration_speaker_fusion, attention_heads=cfg.duration_attention_heads, norm_eps=cfg.norm_eps, architecture=cfg.duration_architecture, token_init_frames=cfg.duration_token_init_frames, ) self.cond_module = nn.Sequential( nn.Linear(cfg.timestep_embed_dim, cfg.model_dim, bias=False), nn.SiLU(), nn.Linear(cfg.model_dim, cfg.model_dim, bias=False), nn.SiLU(), nn.Linear(cfg.model_dim, cfg.model_dim * 3, bias=False), ) self.in_proj = nn.Linear(cfg.patched_latent_dim, cfg.model_dim) self.blocks = nn.ModuleList(DiffusionBlock(cfg) for _ in range(cfg.num_layers)) self.out_norm = RMSNorm(cfg.model_dim, eps=cfg.norm_eps) self.out_proj = nn.Linear(cfg.model_dim, cfg.patched_latent_dim) # Echo/JAX training initializes decoder out projection to zero for stable early training. nn.init.zeros_(self.out_proj.weight) if self.out_proj.bias is not None: nn.init.zeros_(self.out_proj.bias) self.head_dim = cfg.model_dim // cfg.num_heads if self.head_dim % 2 != 0: raise ValueError("model head_dim must be even for RoPE") self.register_buffer( "_freqs_cis_cache", torch.empty(0, 0, dtype=torch.complex64), persistent=False ) def _rope_freqs(self, seq_len: int, device: torch.device) -> torch.Tensor: cache = self._freqs_cis_cache if cache.device != device or cache.shape[0] < seq_len: cache = precompute_freqs_cis(self.head_dim, seq_len).to(device) self._freqs_cis_cache = cache return cache[:seq_len] @staticmethod def _prepend_masked_mean_token( state: torch.Tensor, mask: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """ Prepend one global summary token computed as masked mean over time. """ mask_f = mask.unsqueeze(-1).to(dtype=state.dtype) denom = mask_f.sum(dim=1, keepdim=True).clamp_min(1.0) mean_token = (state * mask_f).sum(dim=1, keepdim=True) / denom has_any = mask.any(dim=1, keepdim=True) state = torch.cat([mean_token, state], dim=1) mask = torch.cat([has_any, mask], dim=1) return state, mask def encode_conditions( self, text_input_ids: torch.Tensor, text_mask: torch.Tensor, ref_latent: torch.Tensor | None, ref_mask: torch.Tensor | None, caption_input_ids: torch.Tensor | None = None, caption_mask: torch.Tensor | None = None, text_condition_dropout: torch.Tensor | None = None, speaker_condition_dropout: torch.Tensor | None = None, caption_condition_dropout: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, ]: if text_condition_dropout is not None: text_mask = text_mask.clone() text_mask[text_condition_dropout] = False if self.cfg.use_speaker_condition: if self.speaker_encoder is None or self.speaker_norm is None: raise RuntimeError( "Speaker conditioning is enabled but speaker modules are missing." ) if ref_latent is None or ref_mask is None: raise ValueError( "ref_latent and ref_mask are required when speaker conditioning is enabled." ) if speaker_condition_dropout is not None: ref_mask = ref_mask.clone() ref_mask[speaker_condition_dropout] = False if self.cfg.use_caption_condition: if self.caption_encoder is None or self.caption_norm is None: raise RuntimeError( "Caption conditioning is enabled but caption modules are missing." ) if caption_input_ids is None or caption_mask is None: raise ValueError( "caption_input_ids and caption_mask are required when caption conditioning is enabled." ) if caption_condition_dropout is not None: caption_mask = caption_mask.clone() caption_mask[caption_condition_dropout] = False text_state = self.text_encoder(text_input_ids, text_mask) text_state = self.text_norm(text_state) ref_state = None if self.cfg.use_speaker_condition: ref_latent, ref_mask = patch_sequence_with_mask( seq=ref_latent, mask=ref_mask, patch_size=self.cfg.speaker_patch_size, ) ref_state = self.speaker_encoder(ref_latent, ref_mask) ref_state = self.speaker_norm(ref_state) ref_state, ref_mask = self._prepend_masked_mean_token(ref_state, ref_mask) caption_state = None if self.cfg.use_caption_condition: caption_state = self.caption_encoder(caption_input_ids, caption_mask) caption_state = self.caption_norm(caption_state) return text_state, text_mask, ref_state, ref_mask, caption_state, caption_mask def forward_with_encoded_conditions( self, x_t: torch.Tensor, t: torch.Tensor, text_state: torch.Tensor, text_mask: torch.Tensor, speaker_state: torch.Tensor | None, speaker_mask: torch.Tensor | None, caption_state: torch.Tensor | None = None, caption_mask: torch.Tensor | None = None, latent_mask: torch.Tensor | None = None, context_kv_cache: list[tuple[torch.Tensor, ...]] | None = None, ) -> torch.Tensor: t_embed = get_timestep_embedding(t, self.cfg.timestep_embed_dim).to(dtype=x_t.dtype) cond_embed = self.cond_module(t_embed) cond_embed = cond_embed[:, None, :] x = self.in_proj(x_t) freqs = self._rope_freqs(x.shape[1], x.device) for i, block in enumerate(self.blocks): x = block( x=x, cond_embed=cond_embed, text_state=text_state, text_mask=text_mask, speaker_state=speaker_state, speaker_mask=speaker_mask, caption_state=caption_state, caption_mask=caption_mask, freqs_cis=freqs, self_mask=latent_mask, context_kv=context_kv_cache[i] if context_kv_cache is not None else None, ) x = self.out_norm(x) x = self.out_proj(x) return x.to(dtype=x_t.dtype) def forward( self, x_t: torch.Tensor | None, t: torch.Tensor | None, text_input_ids: torch.Tensor, text_mask: torch.Tensor, ref_latent: torch.Tensor | None, ref_mask: torch.Tensor | None, caption_input_ids: torch.Tensor | None = None, caption_mask: torch.Tensor | None = None, latent_mask: torch.Tensor | None = None, text_condition_dropout: torch.Tensor | None = None, speaker_condition_dropout: torch.Tensor | None = None, caption_condition_dropout: torch.Tensor | None = None, duration_features: torch.Tensor | None = None, duration_has_speaker: torch.Tensor | None = None, duration_only: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: if duration_features is not None: ( text_state, text_mask_full, speaker_state, speaker_mask_full, caption_state, caption_mask_full, ) = self.encode_conditions( text_input_ids=text_input_ids, text_mask=text_mask, ref_latent=ref_latent, ref_mask=ref_mask, caption_input_ids=caption_input_ids, caption_mask=caption_mask, ) if duration_only: return self.predict_duration_log_frames( text_state=text_state, text_mask=text_mask_full, speaker_state=speaker_state, speaker_mask=speaker_mask_full, duration_features=duration_features, has_speaker=duration_has_speaker, ) if x_t is None or t is None: raise ValueError("x_t and t are required unless duration_only=True.") text_mask_dit = text_mask_full speaker_mask_dit = speaker_mask_full caption_mask_dit = caption_mask_full if text_condition_dropout is not None: text_mask_dit = text_mask_dit.clone() text_mask_dit[text_condition_dropout] = False if speaker_condition_dropout is not None and speaker_mask_dit is not None: speaker_mask_dit = speaker_mask_dit.clone() speaker_mask_dit[speaker_condition_dropout] = False if caption_condition_dropout is not None and caption_mask_dit is not None: caption_mask_dit = caption_mask_dit.clone() caption_mask_dit[caption_condition_dropout] = False v_pred = self.forward_with_encoded_conditions( x_t=x_t, t=t, text_state=text_state, text_mask=text_mask_dit, speaker_state=speaker_state, speaker_mask=speaker_mask_dit, caption_state=caption_state, caption_mask=caption_mask_dit, latent_mask=latent_mask, ) duration_pred = self.predict_duration_log_frames( text_state=text_state, text_mask=text_mask_full, speaker_state=speaker_state, speaker_mask=speaker_mask_full, duration_features=duration_features, has_speaker=duration_has_speaker, ) return v_pred, duration_pred if duration_only: raise ValueError("duration_features is required when duration_only=True.") if x_t is None or t is None: raise ValueError("x_t and t are required for RF forward.") ( text_state, text_mask, speaker_state, speaker_mask, caption_state, caption_mask, ) = self.encode_conditions( text_input_ids=text_input_ids, text_mask=text_mask, ref_latent=ref_latent, ref_mask=ref_mask, caption_input_ids=caption_input_ids, caption_mask=caption_mask, text_condition_dropout=text_condition_dropout, speaker_condition_dropout=speaker_condition_dropout, caption_condition_dropout=caption_condition_dropout, ) return self.forward_with_encoded_conditions( x_t=x_t, t=t, text_state=text_state, text_mask=text_mask, speaker_state=speaker_state, speaker_mask=speaker_mask, caption_state=caption_state, caption_mask=caption_mask, latent_mask=latent_mask, ) def build_context_kv_cache( self, text_state: torch.Tensor, speaker_state: torch.Tensor | None, caption_state: torch.Tensor | None = None, ) -> list[tuple[torch.Tensor, ...]]: """ Build per-layer projected conditioning KV tensors for faster repeated sampling steps. """ return [ block.attention.project_context_kv( text_context=text_state, speaker_context=speaker_state, caption_context=caption_state, ) for block in self.blocks ] @staticmethod def masked_mean(state: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: mask_f = mask.unsqueeze(-1).to(dtype=state.dtype) denom = mask_f.sum(dim=1).clamp_min(1.0) return (state * mask_f).sum(dim=1) / denom def predict_duration_log_frames( self, *, text_state: torch.Tensor, text_mask: torch.Tensor, speaker_state: torch.Tensor | None, speaker_mask: torch.Tensor | None, duration_features: torch.Tensor, has_speaker: torch.Tensor | None, ) -> torch.Tensor: if self.duration_predictor is None: raise RuntimeError("Duration predictor is disabled for this model.") if duration_features.ndim != 2: raise ValueError( f"duration_features must have shape (B, D), got {tuple(duration_features.shape)}" ) if duration_features.shape[1] != self.cfg.duration_aux_dim: raise ValueError( "duration_features dim mismatch: " f"expected {self.cfg.duration_aux_dim}, got {duration_features.shape[1]}" ) pred = self.duration_predictor( text_state=text_state.detach(), text_mask=text_mask, aux_features=duration_features, speaker_state=None if speaker_state is None else speaker_state.detach(), speaker_mask=speaker_mask, has_speaker=has_speaker, ) return pred.float() @property def device(self) -> torch.device: return next(self.parameters()).device @property def dtype(self) -> torch.dtype: return next(self.parameters()).dtype def as_dict(self) -> dict: return asdict(self.cfg)