| """Minimal inference layers used by Cola2. |
| |
| The module intentionally depends only on PyTorch. It contains the dense |
| Qwen3-shaped token decoder, the TextVAE-style encoder/condition block and the |
| latent DiT block needed by the Full-DiT generation path. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from typing import Optional |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
|
|
| def additive_attention_mask( |
| allowed: torch.Tensor, dtype: torch.dtype |
| ) -> torch.Tensor: |
| """Convert a boolean ``[B,Q,K]`` allow tensor to an SDPA bias.""" |
|
|
| if allowed.dtype is not torch.bool: |
| allowed = allowed.bool() |
| mask = torch.zeros( |
| allowed.shape[0], |
| 1, |
| allowed.shape[1], |
| allowed.shape[2], |
| dtype=dtype, |
| device=allowed.device, |
| ) |
| return mask.masked_fill(~allowed.unsqueeze(1), torch.finfo(dtype).min) |
|
|
|
|
| def causal_padding_allow(valid: torch.Tensor) -> torch.Tensor: |
| """Causal attention with finite diagonal entries for padded queries.""" |
|
|
| valid = valid.bool() |
| batch, length = valid.shape |
| index = torch.arange(length, device=valid.device) |
| allowed = valid[:, None, :] & (index[None, None, :] <= index[None, :, None]) |
| eye = torch.eye(length, dtype=torch.bool, device=valid.device) |
| return allowed | eye.unsqueeze(0).expand(batch, -1, -1) |
|
|
|
|
| def block_causal_allow( |
| valid: torch.Tensor, block_ids: torch.Tensor |
| ) -> torch.Tensor: |
| """Allow all valid keys in the query's block or an earlier block.""" |
|
|
| valid = valid.bool() |
| allowed = valid[:, None, :] & ( |
| block_ids[:, None, :] <= block_ids[:, :, None] |
| ) |
| length = valid.shape[1] |
| eye = torch.eye(length, dtype=torch.bool, device=valid.device) |
| return allowed | eye.unsqueeze(0) |
|
|
|
|
| def span_local_allow(valid: torch.Tensor, segment_ids: torch.Tensor) -> torch.Tensor: |
| """Bidirectional attention inside each dynamic span.""" |
|
|
| valid = valid.bool() |
| allowed = ( |
| valid[:, :, None] |
| & valid[:, None, :] |
| & (segment_ids[:, :, None] == segment_ids[:, None, :]) |
| ) |
| length = valid.shape[1] |
| eye = torch.eye(length, dtype=torch.bool, device=valid.device) |
| return allowed | eye.unsqueeze(0) |
|
|
|
|
| class RMSNorm(nn.Module): |
| """Qwen/T5 style RMSNorm with a float32 variance reduction.""" |
|
|
| def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(hidden_size)) |
| self.variance_epsilon = float(eps) |
|
|
| def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| dtype = hidden_states.dtype |
| states = hidden_states.float() |
| variance = states.square().mean(dim=-1, keepdim=True) |
| states = states * torch.rsqrt(variance + self.variance_epsilon) |
| return states.to(dtype) * self.weight |
|
|
|
|
| def _rotate_half(x: torch.Tensor) -> torch.Tensor: |
| x1, x2 = x.chunk(2, dim=-1) |
| return torch.cat((-x2, x1), dim=-1) |
|
|
|
|
| def _rotate_every_two(x: torch.Tensor) -> torch.Tensor: |
| first = x[..., 0::2] |
| second = x[..., 1::2] |
| return torch.stack((-second, first), dim=-1).flatten(-2) |
|
|
|
|
| def _rope_cos_sin( |
| position_ids: torch.Tensor, |
| rotary_dim: int, |
| theta: float, |
| dtype: torch.dtype, |
| *, |
| interleaved: bool, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| inv_freq = 1.0 / ( |
| theta |
| ** ( |
| torch.arange( |
| 0, |
| rotary_dim, |
| 2, |
| dtype=torch.float32, |
| device=position_ids.device, |
| ) |
| / rotary_dim |
| ) |
| ) |
| freqs = position_ids.float().unsqueeze(-1) * inv_freq |
| angles = ( |
| freqs.repeat_interleave(2, dim=-1) |
| if interleaved |
| else torch.cat((freqs, freqs), dim=-1) |
| ).unsqueeze(1) |
| return angles.cos().to(dtype), angles.sin().to(dtype) |
|
|
|
|
| def apply_rotary( |
| query: torch.Tensor, |
| key: torch.Tensor, |
| query_positions: torch.Tensor, |
| key_positions: Optional[torch.Tensor] = None, |
| *, |
| rotary_dim: Optional[int] = None, |
| theta: float = 1_000_000.0, |
| interleaved: bool = False, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Apply non-interleaved RoPE to ``[B,H,L,C]`` query and key.""" |
|
|
| key_positions = query_positions if key_positions is None else key_positions |
| rotary_dim = int(rotary_dim or query.shape[-1]) |
| q_rot, q_pass = query[..., :rotary_dim], query[..., rotary_dim:] |
| k_rot, k_pass = key[..., :rotary_dim], key[..., rotary_dim:] |
| q_cos, q_sin = _rope_cos_sin( |
| query_positions, |
| rotary_dim, |
| theta, |
| q_rot.dtype, |
| interleaved=interleaved, |
| ) |
| k_cos, k_sin = _rope_cos_sin( |
| key_positions, |
| rotary_dim, |
| theta, |
| k_rot.dtype, |
| interleaved=interleaved, |
| ) |
| rotate = _rotate_every_two if interleaved else _rotate_half |
| q_rot = q_rot * q_cos + rotate(q_rot) * q_sin |
| k_rot = k_rot * k_cos + rotate(k_rot) * k_sin |
| return torch.cat((q_rot, q_pass), dim=-1), torch.cat( |
| (k_rot, k_pass), dim=-1 |
| ) |
|
|
|
|
| def _repeat_kv(states: torch.Tensor, repeats: int) -> torch.Tensor: |
| if repeats == 1: |
| return states |
| batch, heads, length, dim = states.shape |
| return ( |
| states[:, :, None, :, :] |
| .expand(batch, heads, repeats, length, dim) |
| .reshape(batch, heads * repeats, length, dim) |
| ) |
|
|
|
|
| class Qwen3Attention(nn.Module): |
| """Dense grouped-query attention with Qwen3-compatible parameter names.""" |
|
|
| def __init__( |
| self, |
| hidden_size: int, |
| num_heads: int, |
| num_key_value_heads: int, |
| head_dim: int, |
| *, |
| norm_eps: float, |
| rope_theta: float, |
| attention_dropout: float, |
| ) -> None: |
| super().__init__() |
| if num_heads % num_key_value_heads: |
| raise ValueError("num_heads must be divisible by num_key_value_heads.") |
| self.num_heads = int(num_heads) |
| self.num_key_value_heads = int(num_key_value_heads) |
| self.num_key_value_groups = self.num_heads // self.num_key_value_heads |
| self.head_dim = int(head_dim) |
| self.rope_theta = float(rope_theta) |
| self.attention_dropout = float(attention_dropout) |
| self.q_proj = nn.Linear( |
| hidden_size, self.num_heads * self.head_dim, bias=False |
| ) |
| self.k_proj = nn.Linear( |
| hidden_size, self.num_key_value_heads * self.head_dim, bias=False |
| ) |
| self.v_proj = nn.Linear( |
| hidden_size, self.num_key_value_heads * self.head_dim, bias=False |
| ) |
| self.o_proj = nn.Linear( |
| self.num_heads * self.head_dim, hidden_size, bias=False |
| ) |
| self.q_norm = RMSNorm(self.head_dim, eps=norm_eps) |
| self.k_norm = RMSNorm(self.head_dim, eps=norm_eps) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| attention_mask: torch.Tensor, |
| position_ids: torch.Tensor, |
| ) -> torch.Tensor: |
| batch, query_length, _ = hidden_states.shape |
| query = self.q_proj(hidden_states).view( |
| batch, query_length, self.num_heads, self.head_dim |
| ) |
| key = self.k_proj(hidden_states).view( |
| batch, query_length, self.num_key_value_heads, self.head_dim |
| ) |
| value = self.v_proj(hidden_states).view( |
| batch, query_length, self.num_key_value_heads, self.head_dim |
| ) |
| query = self.q_norm(query).transpose(1, 2) |
| key = self.k_norm(key).transpose(1, 2) |
| value = value.transpose(1, 2) |
| query, key = apply_rotary( |
| query, |
| key, |
| position_ids, |
| rotary_dim=self.head_dim, |
| theta=self.rope_theta, |
| ) |
| key = _repeat_kv(key, self.num_key_value_groups) |
| value = _repeat_kv(value, self.num_key_value_groups) |
| output = F.scaled_dot_product_attention( |
| query, |
| key, |
| value, |
| attn_mask=attention_mask.to(query.dtype), |
| dropout_p=self.attention_dropout if self.training else 0.0, |
| ) |
| output = output.transpose(1, 2).reshape( |
| batch, query_length, self.num_heads * self.head_dim |
| ) |
| return self.o_proj(output) |
|
|
|
|
| class Qwen3MLP(nn.Module): |
| def __init__(self, hidden_size: int, intermediate_size: int) -> None: |
| super().__init__() |
| self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) |
| self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) |
| self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) |
|
|
| def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| return self.down_proj( |
| F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states) |
| ) |
|
|
|
|
| class Qwen3DecoderBlock(nn.Module): |
| """One pre-norm dense Qwen3 decoder block.""" |
|
|
| def __init__( |
| self, |
| hidden_size: int, |
| intermediate_size: int, |
| num_heads: int, |
| num_key_value_heads: int, |
| head_dim: int, |
| *, |
| norm_eps: float, |
| rope_theta: float, |
| attention_dropout: float = 0.0, |
| hidden_dropout: float = 0.0, |
| ) -> None: |
| super().__init__() |
| self.self_attn = Qwen3Attention( |
| hidden_size, |
| num_heads, |
| num_key_value_heads, |
| head_dim, |
| norm_eps=norm_eps, |
| rope_theta=rope_theta, |
| attention_dropout=attention_dropout, |
| ) |
| self.mlp = Qwen3MLP(hidden_size, intermediate_size) |
| self.input_layernorm = RMSNorm(hidden_size, eps=norm_eps) |
| self.post_attention_layernorm = RMSNorm(hidden_size, eps=norm_eps) |
| self.dropout = nn.Dropout(hidden_dropout) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| attention_mask: torch.Tensor, |
| position_ids: torch.Tensor, |
| ) -> torch.Tensor: |
| residual = hidden_states |
| hidden_states = self.self_attn( |
| self.input_layernorm(hidden_states), |
| attention_mask, |
| position_ids, |
| ) |
| hidden_states = residual + self.dropout(hidden_states) |
| residual = hidden_states |
| hidden_states = self.mlp( |
| self.post_attention_layernorm(hidden_states) |
| ) |
| return residual + self.dropout(hidden_states) |
|
|
|
|
| class SwiGLU(nn.Module): |
| def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| values, gates = hidden_states.chunk(2, dim=-1) |
| return values * F.silu(gates) |
|
|
|
|
| class TextVAEBlock(nn.Module): |
| """TextVAE transformer block used by encoder, pool and condition stacks.""" |
|
|
| def __init__( |
| self, |
| hidden_size: int, |
| intermediate_size: int, |
| num_heads: int, |
| *, |
| norm_eps: float, |
| rope_theta: float, |
| attention_dropout: float = 0.0, |
| hidden_dropout: float = 0.0, |
| norm_style: str = "reordered", |
| ) -> None: |
| super().__init__() |
| if hidden_size % num_heads: |
| raise ValueError("hidden_size must be divisible by num_heads.") |
| if intermediate_size % 2: |
| raise ValueError("TextVAE SwiGLU intermediate size must be even.") |
| self.num_heads = int(num_heads) |
| self.head_dim = hidden_size // num_heads |
| self.rope_theta = float(rope_theta) |
| self.attention_dropout = float(attention_dropout) |
| self.norm_style = str(norm_style) |
| self.norm_attn = nn.LayerNorm(hidden_size, eps=norm_eps) |
| self.qkv_proj = nn.Linear(hidden_size, hidden_size * 3, bias=False) |
| |
| |
| self.q_norm = nn.LayerNorm(hidden_size, eps=norm_eps) |
| self.k_norm = nn.LayerNorm(hidden_size, eps=norm_eps) |
| self.attn_out_proj = nn.Linear(hidden_size, hidden_size, bias=True) |
| self.dropout_layer = nn.Dropout(hidden_dropout) |
| self.norm_ffn = nn.LayerNorm(hidden_size, eps=norm_eps) |
| self.ffn_proj = nn.Linear(hidden_size, intermediate_size, bias=True) |
| self.ffn_act = SwiGLU() |
| self.ffn_out = nn.Linear(intermediate_size // 2, hidden_size, bias=True) |
|
|
| def _attention( |
| self, |
| hidden_states: torch.Tensor, |
| attention_mask: torch.Tensor, |
| position_ids: torch.Tensor, |
| ) -> torch.Tensor: |
| batch, length, hidden_size = hidden_states.shape |
| query, key, value = self.qkv_proj(hidden_states).chunk(3, dim=-1) |
| query = self.q_norm(query).view( |
| batch, length, self.num_heads, self.head_dim |
| ) |
| key = self.k_norm(key).view( |
| batch, length, self.num_heads, self.head_dim |
| ) |
| value = value.view( |
| batch, length, self.num_heads, self.head_dim |
| ).transpose(1, 2) |
| query = query.transpose(1, 2) |
| key = key.transpose(1, 2) |
| query, key = apply_rotary( |
| query, |
| key, |
| position_ids, |
| rotary_dim=self.head_dim, |
| theta=self.rope_theta, |
| ) |
| output = F.scaled_dot_product_attention( |
| query, |
| key, |
| value, |
| attn_mask=attention_mask.to(query.dtype), |
| dropout_p=self.attention_dropout if self.training else 0.0, |
| ) |
| output = output.transpose(1, 2).reshape(batch, length, hidden_size) |
| return self.attn_out_proj(output) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| attention_mask: torch.Tensor, |
| position_ids: torch.Tensor, |
| ) -> torch.Tensor: |
| if self.norm_style == "pre": |
| attention_input = self.norm_attn(hidden_states) |
| hidden_states = hidden_states + self.dropout_layer( |
| self._attention( |
| attention_input, attention_mask, position_ids |
| ) |
| ) |
| residual = hidden_states |
| hidden_states = self.norm_ffn(hidden_states) |
| hidden_states = self.ffn_out( |
| self.ffn_act(self.ffn_proj(hidden_states)) |
| ) |
| return residual + self.dropout_layer(hidden_states) |
|
|
| attention = self._attention( |
| hidden_states, attention_mask, position_ids |
| ) |
| if self.norm_style == "reordered": |
| attention = self.norm_attn(attention) |
| elif self.norm_style == "post_legacy": |
| hidden_states = self.norm_attn(hidden_states) |
| else: |
| raise ValueError(f"Unknown norm_style={self.norm_style!r}.") |
| hidden_states = hidden_states + self.dropout_layer(attention) |
| residual = hidden_states |
| feed_forward = self.ffn_out( |
| self.ffn_act(self.ffn_proj(hidden_states)) |
| ) |
| feed_forward = self.norm_ffn(feed_forward) |
| return residual + self.dropout_layer(feed_forward) |
|
|
|
|
| def sinusoidal_timestep_embedding( |
| timesteps: torch.Tensor, embedding_dim: int |
| ) -> torch.Tensor: |
| """Diffusers-compatible sinusoidal embedding used by the native DiT.""" |
|
|
| half = embedding_dim // 2 |
| exponent = -math.log(10_000.0) * torch.arange( |
| half, dtype=torch.float32, device=timesteps.device |
| ) |
| exponent = exponent / max(half, 1) |
| frequencies = torch.exp(exponent) |
| angles = timesteps.float().unsqueeze(-1) * frequencies.unsqueeze(0) |
| embedding = torch.cat((angles.sin(), angles.cos()), dim=-1) |
| if embedding_dim % 2: |
| embedding = F.pad(embedding, (0, 1)) |
| return embedding |
|
|
|
|
| class TimestepEmbedding(nn.Module): |
| def __init__(self, hidden_size: int, output_size: int) -> None: |
| super().__init__() |
| self.sinusoidal_dim = 256 |
| self.proj_in = nn.Linear(self.sinusoidal_dim, hidden_size) |
| self.proj_hid = nn.Linear(hidden_size, hidden_size) |
| self.proj_out = nn.Linear(hidden_size, output_size) |
|
|
| def forward(self, timesteps: torch.Tensor) -> torch.Tensor: |
| embedding = sinusoidal_timestep_embedding( |
| timesteps.reshape(-1), self.sinusoidal_dim |
| ).to(self.proj_in.weight.dtype) |
| embedding = F.silu(self.proj_in(embedding)) |
| embedding = F.silu(self.proj_hid(embedding)) |
| return self.proj_out(embedding) |
|
|
|
|
| class AdaLN(nn.Module): |
| """Adaptive LayerNorm with native Cola DiT parameter naming.""" |
|
|
| def __init__(self, hidden_size: int, layers: tuple[str, ...]) -> None: |
| super().__init__() |
| for layer in layers: |
| setattr( |
| self, |
| f"{layer}_in", |
| nn.Sequential( |
| nn.SiLU(), |
| nn.Linear(hidden_size, hidden_size * 2, bias=True), |
| ), |
| ) |
| if layer != "out": |
| setattr( |
| self, |
| f"{layer}_out", |
| nn.Sequential( |
| nn.SiLU(), |
| nn.Linear(hidden_size, hidden_size, bias=True), |
| ), |
| ) |
|
|
| def modulate( |
| self, |
| hidden_states: torch.Tensor, |
| embedding: torch.Tensor, |
| layer: str, |
| norm: nn.Module, |
| ) -> torch.Tensor: |
| shift, scale = getattr(self, f"{layer}_in")(embedding).chunk(2, -1) |
| return norm(hidden_states) * (1.0 + scale) + shift |
|
|
| def residual( |
| self, |
| update: torch.Tensor, |
| residual: torch.Tensor, |
| embedding: torch.Tensor, |
| layer: str, |
| ) -> torch.Tensor: |
| gate = getattr(self, f"{layer}_out")(embedding) |
| return residual + update * gate |
|
|
|
|
| class DiTAttention(nn.Module): |
| def __init__( |
| self, |
| hidden_size: int, |
| num_heads: int, |
| head_dim: int, |
| rope_dim: int, |
| *, |
| norm_eps: float, |
| rope_theta: float, |
| attention_dropout: float, |
| ) -> None: |
| super().__init__() |
| self.num_heads = int(num_heads) |
| self.head_dim = int(head_dim) |
| self.rope_dim = int(rope_dim) |
| self.rope_theta = float(rope_theta) |
| self.attention_dropout = float(attention_dropout) |
| self.proj_qkv = nn.Linear( |
| hidden_size, 3 * num_heads * head_dim, bias=False |
| ) |
| self.proj_out = nn.Linear( |
| num_heads * head_dim, hidden_size, bias=False |
| ) |
| self.norm_q = nn.LayerNorm(head_dim, eps=norm_eps) |
| self.norm_k = nn.LayerNorm(head_dim, eps=norm_eps) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| attention_mask: torch.Tensor, |
| position_ids: torch.Tensor, |
| ) -> torch.Tensor: |
| batch, length, _ = hidden_states.shape |
| qkv = self.proj_qkv(hidden_states).view( |
| batch, length, 3, self.num_heads, self.head_dim |
| ) |
| query, key, value = qkv.unbind(dim=2) |
| query = self.norm_q(query).transpose(1, 2) |
| key = self.norm_k(key).transpose(1, 2) |
| value = value.transpose(1, 2) |
| query, key = apply_rotary( |
| query, |
| key, |
| position_ids, |
| rotary_dim=self.rope_dim, |
| theta=self.rope_theta, |
| interleaved=True, |
| ) |
| output = F.scaled_dot_product_attention( |
| query, |
| key, |
| value, |
| attn_mask=attention_mask.to(query.dtype), |
| dropout_p=self.attention_dropout if self.training else 0.0, |
| ) |
| output = output.transpose(1, 2).reshape( |
| batch, length, self.num_heads * self.head_dim |
| ) |
| return self.proj_out(output) |
|
|
|
|
| class DiTMLP(nn.Module): |
| def __init__(self, hidden_size: int, intermediate_size: int) -> None: |
| super().__init__() |
| self.proj_in = nn.Linear(hidden_size, intermediate_size, bias=True) |
| self.act = nn.GELU(approximate="tanh") |
| self.proj_out = nn.Linear(intermediate_size, hidden_size, bias=True) |
|
|
| def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| return self.proj_out(self.act(self.proj_in(hidden_states))) |
|
|
|
|
| class DiTBlock(nn.Module): |
| def __init__( |
| self, |
| hidden_size: int, |
| intermediate_size: int, |
| num_heads: int, |
| head_dim: int, |
| rope_dim: int, |
| *, |
| norm_eps: float, |
| rope_theta: float, |
| attention_dropout: float, |
| ) -> None: |
| super().__init__() |
| self.msa_norm = nn.LayerNorm( |
| hidden_size, eps=norm_eps, elementwise_affine=False |
| ) |
| self.msa = DiTAttention( |
| hidden_size, |
| num_heads, |
| head_dim, |
| rope_dim, |
| norm_eps=norm_eps, |
| rope_theta=rope_theta, |
| attention_dropout=attention_dropout, |
| ) |
| self.mlp_norm = nn.LayerNorm( |
| hidden_size, eps=norm_eps, elementwise_affine=False |
| ) |
| self.mlp = DiTMLP(hidden_size, intermediate_size) |
| self.ada = AdaLN(hidden_size, ("msa", "mlp")) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| embedding: torch.Tensor, |
| attention_mask: torch.Tensor, |
| position_ids: torch.Tensor, |
| ) -> torch.Tensor: |
| attention_input = self.ada.modulate( |
| hidden_states, embedding, "msa", self.msa_norm |
| ) |
| attention_output = self.msa( |
| attention_input, attention_mask, position_ids |
| ) |
| hidden_states = self.ada.residual( |
| attention_output, hidden_states, embedding, "msa" |
| ) |
| mlp_input = self.ada.modulate( |
| hidden_states, embedding, "mlp", self.mlp_norm |
| ) |
| mlp_output = self.mlp(mlp_input) |
| return self.ada.residual( |
| mlp_output, hidden_states, embedding, "mlp" |
| ) |
|
|
|
|
| class DurationHead(nn.Module): |
| def __init__(self, hidden_size: int, intermediate_size: int, lmax: int): |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.Linear(hidden_size, intermediate_size), |
| nn.SiLU(), |
| nn.Linear(intermediate_size, lmax), |
| ) |
|
|
| def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| return self.net(hidden_states) |
|
|
|
|
| class HardDurationExpander(nn.Module): |
| """Repeat latent conditions by duration and add a local-position feature.""" |
|
|
| def __init__( |
| self, hidden_size: int, output_size: int, local_position_dim: int |
| ) -> None: |
| super().__init__() |
| self.output_size = int(output_size) |
| self.loc_mlp = nn.Sequential( |
| nn.Linear(3, local_position_dim), |
| nn.SiLU(), |
| nn.Linear(local_position_dim, local_position_dim), |
| ) |
| self.cond_proj = nn.Linear( |
| hidden_size + local_position_dim, output_size |
| ) |
|
|
| def forward( |
| self, hidden_states: torch.Tensor, lengths: torch.Tensor |
| ) -> torch.Tensor: |
| """Expand one unpadded latent row into token-aligned conditions.""" |
|
|
| if hidden_states.ndim != 2 or lengths.ndim != 1: |
| raise ValueError("Expander expects hidden [K,D] and lengths [K].") |
| if hidden_states.shape[0] != lengths.shape[0]: |
| raise ValueError("Each latent needs exactly one duration.") |
| pieces = [] |
| for latent, length_tensor in zip(hidden_states, lengths, strict=True): |
| length = int(length_tensor.item()) |
| if length <= 0: |
| continue |
| positions = torch.arange( |
| 1, |
| length + 1, |
| dtype=hidden_states.dtype, |
| device=hidden_states.device, |
| ) |
| length_f = positions.new_full((length,), float(length)) |
| local_features = torch.stack( |
| ( |
| positions / length_f, |
| (length_f - positions) / length_f, |
| length_f.log(), |
| ), |
| dim=-1, |
| ) |
| local = self.loc_mlp(local_features) |
| repeated = latent.unsqueeze(0).expand(length, -1) |
| pieces.append(self.cond_proj(torch.cat((repeated, local), dim=-1))) |
| if not pieces: |
| return hidden_states.new_zeros((0, self.output_size)) |
| return torch.cat(pieces, dim=0) |
|
|
