| """Hugging Face compatible Cola2 span-causal Full-DiT model.""" |
|
|
| from __future__ import annotations |
|
|
| import copy |
| import json |
| import math |
| from collections.abc import Mapping |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any, Optional, Union |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
| from transformers import PreTrainedModel |
| from transformers.generation.utils import GenerateDecoderOnlyOutput |
| from transformers.modeling_outputs import ModelOutput |
|
|
| from .configuration_cola2 import Cola2Config, Cola2GenerationConfig |
| from .modeling_layers import ( |
| AdaLN, |
| DiTBlock, |
| DurationHead, |
| HardDurationExpander, |
| Qwen3DecoderBlock, |
| RMSNorm, |
| TextVAEBlock, |
| TimestepEmbedding, |
| additive_attention_mask, |
| block_causal_allow, |
| causal_padding_allow, |
| span_local_allow, |
| ) |
|
|
|
|
| def _build_dit_timesteps( |
| *, |
| schedule_t: float, |
| steps: int, |
| schedule: str, |
| power: Optional[float], |
| device: torch.device, |
| ) -> torch.Tensor: |
| named_powers = { |
| "uniform": 1.0, |
| "high_noise": 0.5, |
| "low_noise": 2.0, |
| } |
| if schedule == "power": |
| exponent = float(power if power is not None else 1.0) |
| else: |
| exponent = named_powers[schedule] |
| progress = torch.linspace( |
| 1.0, |
| 0.0, |
| int(steps) + 1, |
| dtype=torch.float32, |
| device=device, |
| ) |
| return progress.pow(exponent).mul_(float(schedule_t)) |
|
|
|
|
| @dataclass |
| class Cola2EncoderOutput(ModelOutput): |
| """Dynamic compression encoder output.""" |
|
|
| latents: Optional[torch.Tensor] = None |
| latent_logvar: Optional[torch.Tensor] = None |
| latent_mask: Optional[torch.Tensor] = None |
| span_starts: Optional[torch.Tensor] = None |
| span_ends: Optional[torch.Tensor] = None |
| span_lengths: Optional[torch.Tensor] = None |
| segment_ids: Optional[torch.Tensor] = None |
| boundary_logits: Optional[torch.Tensor] = None |
| boundary_probabilities: Optional[torch.Tensor] = None |
|
|
|
|
| @dataclass |
| class Cola2ConditionalGenerationOutput(ModelOutput): |
| """Output of the same-position span-causal reconstruction forward.""" |
|
|
| loss: Optional[torch.Tensor] = None |
| logits: Optional[torch.Tensor] = None |
| latents: Optional[torch.Tensor] = None |
| latent_mask: Optional[torch.Tensor] = None |
| span_lengths: Optional[torch.Tensor] = None |
| duration_logits: Optional[torch.Tensor] = None |
| boundary_logits: Optional[torch.Tensor] = None |
|
|
|
|
| class CosineBoundaryScorer(nn.Module): |
| """Native cosine scorer layout (``segmenter.scorer.norm`` in checkpoints).""" |
|
|
| def __init__(self, config: Cola2Config) -> None: |
| super().__init__() |
| self.alpha = config.boundary_score_alpha |
| self.threshold = config.boundary_score_threshold |
| self.norm = nn.LayerNorm( |
| config.encoder_hidden_size, eps=config.layer_norm_eps |
| ) |
|
|
| def forward( |
| self, hidden_states: torch.Tensor, attention_mask: torch.Tensor |
| ) -> torch.Tensor: |
| batch, width, _ = hidden_states.shape |
| if width <= 1: |
| return hidden_states.new_zeros((batch, 0)) |
| normalized = self.norm(hidden_states) |
| cosine = F.cosine_similarity( |
| normalized[:, :-1], normalized[:, 1:], dim=-1 |
| ) |
| logits = self.alpha * (1.0 - cosine - self.threshold) |
| valid = attention_mask.bool() |
| return logits.masked_fill( |
| ~(valid[:, :-1] & valid[:, 1:]), -1e4 |
| ) |
|
|
|
|
| class DynamicSegmenter(nn.Module): |
| """Auxiliary-loss-free cosine boundary scorer with hard span constraints.""" |
|
|
| def __init__(self, config: Cola2Config) -> None: |
| super().__init__() |
| self.lmin = config.min_span_length |
| self.lmax = config.max_span_length |
| self.scorer = CosineBoundaryScorer(config) |
| rate = min(max(config.target_compression_rate, 1e-4), 1.0 - 1e-4) |
| self.register_buffer( |
| "rate_bias", |
| torch.tensor(math.log(rate / (1.0 - rate)), dtype=torch.float32), |
| ) |
|
|
| def forward( |
| self, hidden_states: torch.Tensor, attention_mask: torch.Tensor |
| ) -> dict[str, torch.Tensor]: |
| batch, width, _ = hidden_states.shape |
| valid = attention_mask.bool() |
| logits = self.scorer(hidden_states, valid) |
| if width > 1: |
| probabilities = torch.sigmoid( |
| logits + self.rate_bias.to(logits.dtype) |
| ) |
| else: |
| logits = hidden_states.new_zeros((batch, 0)) |
| probabilities = hidden_states.new_zeros((batch, 0)) |
|
|
| rows: list[list[tuple[int, int]]] = [] |
| boundaries = torch.zeros( |
| batch, width, dtype=torch.bool, device=hidden_states.device |
| ) |
| for row in range(batch): |
| length = int(valid[row].long().sum().item()) |
| if length <= 0: |
| rows.append([]) |
| continue |
| spans: list[tuple[int, int]] = [] |
| start = 0 |
| for token_index in range(length): |
| span_length = token_index - start + 1 |
| wanted = ( |
| token_index < length - 1 |
| and bool(probabilities[row, token_index] > 0.5) |
| ) |
| should_cut = ( |
| token_index == length - 1 |
| or span_length >= self.lmax |
| or (wanted and span_length >= self.lmin) |
| ) |
| if should_cut: |
| spans.append((start, token_index)) |
| boundaries[row, token_index] = True |
| start = token_index + 1 |
| rows.append(spans) |
|
|
| max_spans = max((len(row) for row in rows), default=0) |
| max_spans = max(max_spans, 1) |
| starts = torch.zeros( |
| batch, max_spans, dtype=torch.long, device=hidden_states.device |
| ) |
| ends = torch.zeros_like(starts) |
| lengths = torch.zeros_like(starts) |
| latent_mask = torch.zeros( |
| batch, max_spans, dtype=torch.bool, device=hidden_states.device |
| ) |
| segment_ids = torch.full( |
| (batch, width), -1, dtype=torch.long, device=hidden_states.device |
| ) |
| for row, spans in enumerate(rows): |
| for span_index, (start, end) in enumerate(spans): |
| starts[row, span_index] = start |
| ends[row, span_index] = end |
| lengths[row, span_index] = end - start + 1 |
| latent_mask[row, span_index] = True |
| segment_ids[row, start : end + 1] = span_index |
|
|
| return { |
| "boundary_logits": logits, |
| "boundary_probabilities": probabilities, |
| "boundary_hard": boundaries, |
| "span_starts": starts, |
| "span_ends": ends, |
| "span_lengths": lengths, |
| "latent_mask": latent_mask, |
| "segment_ids": segment_ids, |
| } |
|
|
|
|
| class SpanLocalPool(nn.Module): |
| """Span-local bidirectional contextualization followed by mean pooling.""" |
|
|
| def __init__(self, config: Cola2Config) -> None: |
| super().__init__() |
| self.blocks = nn.ModuleList( |
| TextVAEBlock( |
| config.encoder_hidden_size, |
| config.encoder_intermediate_size, |
| config.encoder_num_attention_heads, |
| norm_eps=config.layer_norm_eps, |
| rope_theta=config.encoder_rope_theta, |
| attention_dropout=config.attention_dropout, |
| hidden_dropout=config.hidden_dropout, |
| norm_style="post_legacy", |
| ) |
| for _ in range(config.pool_num_hidden_layers) |
| ) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| attention_mask: torch.Tensor, |
| segment_ids: torch.Tensor, |
| latent_mask: torch.Tensor, |
| ) -> torch.Tensor: |
| valid = attention_mask.bool() |
| positions = (torch.cumsum(valid.long(), dim=1) - 1).clamp_min(0) |
| allowed = span_local_allow(valid, segment_ids) |
| bias = additive_attention_mask(allowed, hidden_states.dtype) |
| for block in self.blocks: |
| hidden_states = block(hidden_states, bias, positions) |
|
|
| batch, _, dim = hidden_states.shape |
| max_spans = latent_mask.shape[1] |
| pooled = hidden_states.new_zeros((batch, max_spans, dim)) |
| for row in range(batch): |
| for span_index in range(max_spans): |
| if not bool(latent_mask[row, span_index]): |
| continue |
| selected = segment_ids[row] == span_index |
| pooled[row, span_index] = hidden_states[row, selected].mean( |
| dim=0 |
| ) |
| return pooled |
|
|
|
|
| class Cola2CompressionVAE(nn.Module): |
| """Dynamic text compressor and span-causal token decoder.""" |
|
|
| def __init__(self, config: Cola2Config) -> None: |
| super().__init__() |
| self.config = config |
| input_rows = max( |
| config.vocab_size, |
| int(config.pad_token_id or 0) + 1, |
| int(config.mask_token_id or 0) + 1, |
| ) |
| self.encoder = nn.ModuleDict( |
| { |
| "wte": nn.Embedding(input_rows, config.encoder_embedding_size), |
| "input_proj": ( |
| nn.Linear( |
| config.encoder_embedding_size, |
| config.encoder_hidden_size, |
| bias=False, |
| ) |
| if config.encoder_embedding_size |
| != config.encoder_hidden_size |
| else nn.Identity() |
| ), |
| "blocks": nn.ModuleList( |
| TextVAEBlock( |
| config.encoder_hidden_size, |
| config.encoder_intermediate_size, |
| config.encoder_num_attention_heads, |
| norm_eps=config.layer_norm_eps, |
| rope_theta=config.encoder_rope_theta, |
| attention_dropout=config.attention_dropout, |
| hidden_dropout=config.hidden_dropout, |
| norm_style="post_legacy", |
| ) |
| for _ in range(config.encoder_num_hidden_layers) |
| ), |
| } |
| ) |
| self.segmenter = DynamicSegmenter(config) |
| self.pool = SpanLocalPool(config) |
| self.latent_encoder_blocks = nn.ModuleList( |
| TextVAEBlock( |
| config.encoder_hidden_size, |
| config.encoder_intermediate_size, |
| config.encoder_num_attention_heads, |
| norm_eps=config.layer_norm_eps, |
| rope_theta=config.encoder_rope_theta, |
| attention_dropout=config.attention_dropout, |
| hidden_dropout=config.hidden_dropout, |
| norm_style="post_legacy", |
| ) |
| for _ in range(config.latent_encoder_num_hidden_layers) |
| ) |
| posterior_size = config.latent_dim * ( |
| 2 if config.use_variational_encoder else 1 |
| ) |
| self.posterior_proj = nn.Linear( |
| config.encoder_hidden_size, posterior_size, bias=True |
| ) |
| self.posterior_norm = nn.LayerNorm( |
| config.latent_dim, |
| eps=config.layer_norm_eps, |
| elementwise_affine=False, |
| ) |
|
|
| self.latent_in_proj = nn.Linear( |
| config.latent_dim, config.decoder_hidden_size, bias=True |
| ) |
| condition_intermediate = config.decoder_intermediate_size * 2 |
| condition_heads = config.condition_num_attention_heads |
| self.latent_decoder_blocks = nn.ModuleList( |
| TextVAEBlock( |
| config.decoder_hidden_size, |
| condition_intermediate, |
| condition_heads, |
| norm_eps=config.layer_norm_eps, |
| rope_theta=config.condition_rope_theta, |
| attention_dropout=config.attention_dropout, |
| hidden_dropout=config.hidden_dropout, |
| norm_style="reordered", |
| ) |
| for _ in range(config.latent_decoder_num_hidden_layers) |
| ) |
| self.duration_head = DurationHead( |
| config.decoder_hidden_size, |
| config.duration_hidden_size, |
| config.max_span_length, |
| ) |
| self.expander = HardDurationExpander( |
| config.decoder_hidden_size, |
| config.decoder_hidden_size, |
| config.local_position_dim, |
| ) |
| self.decoder_cond_stack = nn.ModuleList( |
| TextVAEBlock( |
| config.decoder_hidden_size, |
| condition_intermediate, |
| condition_heads, |
| norm_eps=config.layer_norm_eps, |
| rope_theta=config.condition_rope_theta, |
| attention_dropout=config.attention_dropout, |
| hidden_dropout=config.hidden_dropout, |
| norm_style="reordered", |
| ) |
| for _ in range(config.decoder_condition_num_hidden_layers) |
| ) |
|
|
| self.decoder = nn.ModuleDict( |
| { |
| "blocks": nn.ModuleList( |
| Qwen3DecoderBlock( |
| config.decoder_hidden_size, |
| config.decoder_intermediate_size, |
| config.decoder_num_attention_heads, |
| config.decoder_num_key_value_heads, |
| config.decoder_head_dim, |
| norm_eps=config.layer_norm_eps, |
| rope_theta=config.decoder_rope_theta, |
| attention_dropout=config.attention_dropout, |
| hidden_dropout=config.hidden_dropout, |
| ) |
| for _ in range(config.decoder_num_hidden_layers) |
| ), |
| "final_norm": RMSNorm( |
| config.decoder_hidden_size, eps=config.layer_norm_eps |
| ), |
| "final_layer": nn.Linear( |
| config.decoder_hidden_size, |
| config.vocab_size, |
| bias=config.decoder_lm_head_bias, |
| ), |
| } |
| ) |
| self.decoder_mask_embedding = nn.Parameter( |
| torch.zeros(config.decoder_hidden_size) |
| ) |
| self.decoder_cond_in = nn.Linear( |
| config.decoder_hidden_size, |
| config.decoder_hidden_size, |
| bias=True, |
| ) |
|
|
| def token_hidden( |
| self, input_ids: torch.Tensor, attention_mask: torch.Tensor |
| ) -> torch.Tensor: |
| hidden_states = self.encoder["input_proj"]( |
| self.encoder["wte"](input_ids) |
| ) |
| valid = attention_mask.bool() |
| positions = (torch.cumsum(valid.long(), dim=1) - 1).clamp_min(0) |
| bias = additive_attention_mask( |
| causal_padding_allow(valid), hidden_states.dtype |
| ) |
| for block in self.encoder["blocks"]: |
| hidden_states = block(hidden_states, bias, positions) |
| return hidden_states |
|
|
| def _latent_encode( |
| self, pooled: torch.Tensor, latent_mask: torch.Tensor |
| ) -> torch.Tensor: |
| if not self.latent_encoder_blocks: |
| return pooled |
| valid = latent_mask.bool() |
| positions = torch.arange( |
| pooled.shape[1], device=pooled.device, dtype=torch.long |
| ).unsqueeze(0).expand(pooled.shape[0], -1) |
| if self.config.latent_block_causal: |
| block_ids = positions // self.config.latent_block_size |
| allowed = block_causal_allow(valid, block_ids) |
| else: |
| allowed = causal_padding_allow(valid) |
| bias = additive_attention_mask(allowed, pooled.dtype) |
| for block in self.latent_encoder_blocks: |
| pooled = block(pooled, bias, positions) |
| return pooled |
|
|
| def _posterior( |
| self, hidden_states: torch.Tensor |
| ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: |
| parameters = self.posterior_proj(hidden_states) |
| if self.config.use_variational_encoder: |
| mean, logvar = parameters.chunk(2, dim=-1) |
| return self.posterior_norm(mean), logvar.clamp(-30.0, 20.0) |
| return self.posterior_norm(parameters), None |
|
|
| def encode_dynamic( |
| self, input_ids: torch.Tensor, attention_mask: torch.Tensor |
| ) -> Cola2EncoderOutput: |
| hidden_states = self.token_hidden(input_ids, attention_mask) |
| segmentation = self.segmenter(hidden_states, attention_mask) |
| pooled = self.pool( |
| hidden_states, |
| attention_mask, |
| segmentation["segment_ids"], |
| segmentation["latent_mask"], |
| ) |
| pooled = self._latent_encode(pooled, segmentation["latent_mask"]) |
| latents, logvar = self._posterior(pooled) |
| latents = latents * segmentation["latent_mask"].unsqueeze(-1) |
| return Cola2EncoderOutput( |
| latents=latents, |
| latent_logvar=logvar, |
| latent_mask=segmentation["latent_mask"], |
| span_starts=segmentation["span_starts"], |
| span_ends=segmentation["span_ends"], |
| span_lengths=segmentation["span_lengths"], |
| segment_ids=segmentation["segment_ids"], |
| boundary_logits=segmentation["boundary_logits"], |
| boundary_probabilities=segmentation["boundary_probabilities"], |
| ) |
|
|
| def encode_given_spans( |
| self, input_ids: torch.Tensor, span_lengths: torch.Tensor |
| ) -> torch.Tensor: |
| """Re-encode committed tokens under decoder-predicted span geometry.""" |
|
|
| if input_ids.ndim != 1 or span_lengths.ndim != 1: |
| raise ValueError("encode_given_spans expects unpadded 1-D rows.") |
| if int(span_lengths.sum().item()) != int(input_ids.numel()): |
| raise ValueError( |
| "Feedback spans must be positive and cover every committed token." |
| ) |
| if bool((span_lengths <= 0).any()): |
| raise ValueError("Feedback span lengths must all be positive.") |
| token_batch = input_ids.unsqueeze(0) |
| valid = torch.ones_like(token_batch, dtype=torch.bool) |
| hidden_states = self.token_hidden(token_batch, valid) |
| segment_ids = torch.empty_like(token_batch) |
| offset = 0 |
| for span_index, length_tensor in enumerate(span_lengths): |
| length = int(length_tensor.item()) |
| segment_ids[:, offset : offset + length] = span_index |
| offset += length |
| latent_mask = torch.ones( |
| 1, |
| span_lengths.numel(), |
| dtype=torch.bool, |
| device=input_ids.device, |
| ) |
| pooled = self.pool( |
| hidden_states, valid, segment_ids, latent_mask |
| ) |
| pooled = self._latent_encode(pooled, latent_mask) |
| latents, _ = self._posterior(pooled) |
| return latents[0] |
|
|
| def latent_decode( |
| self, latents: torch.Tensor |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| if latents.ndim != 2: |
| raise ValueError("latent_decode expects [K, latent_dim].") |
| hidden_states = self.latent_in_proj(latents).unsqueeze(0) |
| length = hidden_states.shape[1] |
| valid = torch.ones( |
| 1, length, dtype=torch.bool, device=latents.device |
| ) |
| positions = torch.arange( |
| length, device=latents.device, dtype=torch.long |
| ).unsqueeze(0) |
| if self.config.latent_block_causal: |
| block_ids = positions // self.config.latent_block_size |
| allowed = block_causal_allow(valid, block_ids) |
| else: |
| allowed = causal_padding_allow(valid) |
| bias = additive_attention_mask(allowed, hidden_states.dtype) |
| for block in self.latent_decoder_blocks: |
| hidden_states = block(hidden_states, bias, positions) |
| duration_logits = self.duration_head(hidden_states) |
| return hidden_states[0], duration_logits[0] |
|
|
| def expand_and_refine( |
| self, latents: torch.Tensor, span_lengths: torch.Tensor |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Create token-aligned latent conditions for one complete sequence.""" |
|
|
| latent_hidden, duration_logits = self.latent_decode(latents) |
| condition = self.expander(latent_hidden, span_lengths) |
| if condition.shape[0] == 0: |
| return condition, duration_logits |
| latent_blocks = ( |
| torch.arange( |
| span_lengths.numel(), |
| device=latents.device, |
| dtype=torch.long, |
| ) |
| // self.config.dit_block_size |
| ) |
| token_blocks = torch.repeat_interleave( |
| latent_blocks, span_lengths.long(), dim=0 |
| ).unsqueeze(0) |
| valid = torch.ones( |
| 1, condition.shape[0], dtype=torch.bool, device=latents.device |
| ) |
| positions = torch.arange( |
| condition.shape[0], device=latents.device, dtype=torch.long |
| ).unsqueeze(0) |
| bias = additive_attention_mask( |
| block_causal_allow(valid, token_blocks), condition.dtype |
| ) |
| refined = condition.unsqueeze(0) |
| for block in self.decoder_cond_stack: |
| refined = block(refined, bias, positions) |
| return refined[0], duration_logits |
|
|
| def span_query_logits( |
| self, |
| clean_ids: torch.Tensor, |
| query_condition: torch.Tensor, |
| *, |
| conditional: bool, |
| ) -> torch.Tensor: |
| """Predict one MASK-query chunk against the clean token prefix.""" |
|
|
| if clean_ids.ndim != 1 or query_condition.ndim != 2: |
| raise ValueError("span_query_logits expects one unpadded row.") |
| query_width = query_condition.shape[0] |
| if query_width == 0: |
| return query_condition.new_zeros((0, self.config.vocab_size)) |
| output_weight = self.decoder["final_layer"].weight |
| if clean_ids.numel(): |
| clean_embedding = F.embedding( |
| clean_ids.clamp(0, self.config.vocab_size - 1), |
| output_weight, |
| ) |
| else: |
| clean_embedding = query_condition.new_zeros( |
| (0, self.config.decoder_hidden_size) |
| ) |
| query_embedding = self.decoder_mask_embedding.unsqueeze(0).expand( |
| query_width, -1 |
| ) |
| if conditional: |
| query_embedding = query_embedding + self.decoder_cond_in( |
| query_condition |
| ) |
| hidden_states = torch.cat( |
| (clean_embedding, query_embedding), dim=0 |
| ).unsqueeze(0) |
| total = hidden_states.shape[1] |
| valid = torch.ones( |
| 1, total, dtype=torch.bool, device=hidden_states.device |
| ) |
| positions = torch.arange( |
| total, device=hidden_states.device, dtype=torch.long |
| ).unsqueeze(0) |
| bias = additive_attention_mask( |
| causal_padding_allow(valid), hidden_states.dtype |
| ) |
| for block in self.decoder["blocks"]: |
| hidden_states = block(hidden_states, bias, positions) |
| query_states = hidden_states[:, -query_width:] |
| query_states = self.decoder["final_norm"](query_states) |
| return self.decoder["final_layer"](query_states)[0] |
|
|
| def teacher_forced_logits( |
| self, |
| input_ids: torch.Tensor, |
| latents: torch.Tensor, |
| span_lengths: torch.Tensor, |
| *, |
| chunk_spans: int, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Same-position dual-stream logits used by :meth:`forward`.""" |
|
|
| condition, duration_logits = self.expand_and_refine( |
| latents, span_lengths |
| ) |
| logits = condition.new_zeros( |
| input_ids.numel(), self.config.vocab_size |
| ) |
| bounds = torch.cat( |
| ( |
| span_lengths.new_zeros(1), |
| torch.cumsum(span_lengths, dim=0), |
| ), |
| dim=0, |
| ) |
| for span_lo in range(0, span_lengths.numel(), chunk_spans): |
| span_hi = min(span_lo + chunk_spans, span_lengths.numel()) |
| token_lo = int(bounds[span_lo].item()) |
| token_hi = int(bounds[span_hi].item()) |
| logits[token_lo:token_hi] = self.span_query_logits( |
| input_ids[:token_lo], |
| condition[token_lo:token_hi], |
| conditional=True, |
| ) |
| return logits, duration_logits |
|
|
|
|
| class _PatchProjection(nn.Module): |
| def __init__(self, input_size: int, output_size: int) -> None: |
| super().__init__() |
| self.proj = nn.Linear(input_size, output_size) |
|
|
| def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| return self.proj(hidden_states) |
|
|
|
|
| class Cola2DiT(nn.Module): |
| """Block-causal latent flow transformer.""" |
|
|
| def __init__(self, config: Cola2Config) -> None: |
| super().__init__() |
| self.block_size = config.dit_block_size |
| self.txt_in = _PatchProjection( |
| config.latent_dim, config.dit_hidden_size |
| ) |
| self.emb_in = TimestepEmbedding( |
| config.dit_hidden_size, config.dit_hidden_size |
| ) |
| self.blocks = nn.ModuleList( |
| DiTBlock( |
| config.dit_hidden_size, |
| config.dit_intermediate_size, |
| config.dit_num_attention_heads, |
| config.dit_head_dim, |
| config.dit_rope_dim, |
| norm_eps=config.dit_layer_norm_eps, |
| rope_theta=config.dit_rope_theta, |
| attention_dropout=config.attention_dropout, |
| ) |
| for _ in range(config.dit_num_hidden_layers) |
| ) |
| self.txt_out_norm = nn.LayerNorm( |
| config.dit_hidden_size, |
| eps=config.dit_layer_norm_eps, |
| elementwise_affine=True, |
| ) |
| self.txt_out_ada = AdaLN(config.dit_hidden_size, ("out",)) |
| self.txt_out = _PatchProjection( |
| config.dit_hidden_size, config.latent_dim |
| ) |
|
|
| def forward( |
| self, latent_states: torch.Tensor, timesteps: torch.Tensor |
| ) -> torch.Tensor: |
| if latent_states.ndim != 2: |
| raise ValueError("DiT expects one latent row [L,D].") |
| length = latent_states.shape[0] |
| if timesteps.ndim == 0: |
| timesteps = timesteps.expand(length) |
| if timesteps.numel() == 1: |
| timesteps = timesteps.expand(length) |
| if timesteps.numel() != length: |
| raise ValueError("DiT needs one timestep per latent position.") |
| hidden_states = self.txt_in(latent_states).unsqueeze(0) |
| embedding = self.emb_in(timesteps).unsqueeze(0) |
| positions = torch.arange( |
| length, device=latent_states.device, dtype=torch.long |
| ).unsqueeze(0) |
| valid = torch.ones( |
| 1, length, dtype=torch.bool, device=latent_states.device |
| ) |
| block_ids = positions // self.block_size |
| bias = additive_attention_mask( |
| block_causal_allow(valid, block_ids), hidden_states.dtype |
| ) |
| for block in self.blocks: |
| hidden_states = block( |
| hidden_states, embedding, bias, positions |
| ) |
| hidden_states = self.txt_out_ada.modulate( |
| hidden_states, |
| embedding, |
| "out", |
| self.txt_out_norm, |
| ) |
| return self.txt_out(hidden_states)[0] |
|
|
|
|
| class Cola2ForConditionalGeneration(PreTrainedModel): |
| """Runnable span-causal Cola2 Full-DiT generation wrapper. |
| |
| Generation follows the production sequence: |
| |
| ``prompt encode -> DiT denoise -> dynamic duration -> span-causal decode |
| -> feedback re-encode``. |
| """ |
|
|
| config_class = Cola2Config |
| generation_config_class = Cola2GenerationConfig |
| base_model_prefix = "cola2" |
| main_input_name = "input_ids" |
| supports_gradient_checkpointing = False |
| _no_split_modules = ["Qwen3DecoderBlock", "DiTBlock", "TextVAEBlock"] |
|
|
| def __init__(self, config: Cola2Config) -> None: |
| super().__init__(config) |
| self.vae = Cola2CompressionVAE(config) |
| self.dit = Cola2DiT(config) |
| self.last_generation_stats: dict[str, int] = {} |
| self.post_init() |
| self.generation_config = Cola2GenerationConfig( |
| pad_token_id=config.pad_token_id, |
| bos_token_id=config.bos_token_id, |
| eos_token_id=config.eos_token_id, |
| ) |
|
|
| def _init_weights(self, module: nn.Module) -> None: |
| if isinstance(module, nn.Linear): |
| nn.init.normal_( |
| module.weight, mean=0.0, std=self.config.initializer_range |
| ) |
| if module.bias is not None: |
| nn.init.zeros_(module.bias) |
| elif isinstance(module, nn.Embedding): |
| nn.init.normal_( |
| module.weight, mean=0.0, std=self.config.initializer_range |
| ) |
| elif isinstance(module, nn.LayerNorm): |
| if module.weight is not None: |
| nn.init.ones_(module.weight) |
| if module.bias is not None: |
| nn.init.zeros_(module.bias) |
| elif isinstance(module, RMSNorm): |
| nn.init.ones_(module.weight) |
|
|
| @classmethod |
| def from_pretrained(cls, *args: Any, **kwargs: Any) -> Any: |
| """Preserve Cola2-specific generation fields after HF loading.""" |
|
|
| result = super().from_pretrained(*args, **kwargs) |
| model = result[0] if isinstance(result, tuple) else result |
| if not isinstance(model.generation_config, Cola2GenerationConfig): |
| model.generation_config = Cola2GenerationConfig.from_dict( |
| model.generation_config.to_dict() |
| ) |
| return result |
|
|
| def get_input_embeddings(self) -> nn.Module: |
| return self.vae.encoder["wte"] |
|
|
| def set_input_embeddings(self, value: nn.Module) -> None: |
| self.vae.encoder["wte"] = value |
|
|
| def get_output_embeddings(self) -> nn.Module: |
| return self.vae.decoder["final_layer"] |
|
|
| def set_output_embeddings(self, value: nn.Module) -> None: |
| self.vae.decoder["final_layer"] = value |
|
|
| def encode( |
| self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None |
| ) -> Cola2EncoderOutput: |
| if attention_mask is None: |
| attention_mask = input_ids.ne(self.config.pad_token_id) |
| return self.vae.encode_dynamic(input_ids, attention_mask) |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| attention_mask: Optional[torch.Tensor] = None, |
| labels: Optional[torch.Tensor] = None, |
| return_dict: Optional[bool] = None, |
| **_: Any, |
| ) -> Union[Cola2ConditionalGenerationOutput, tuple[torch.Tensor, ...]]: |
| """Encode and run the real span-causal reconstruction graph. |
| |
| Unlike an autoregressive LM, Cola2 predicts each token at the same |
| position from its latent condition and the clean prefix available at the |
| start of that complete-span chunk. |
| """ |
|
|
| return_dict = ( |
| self.config.use_return_dict |
| if return_dict is None |
| else bool(return_dict) |
| ) |
| if input_ids.ndim != 2: |
| raise ValueError("input_ids must have shape [batch, sequence].") |
| if attention_mask is None: |
| attention_mask = input_ids.ne(self.config.pad_token_id) |
| attention_mask = attention_mask.bool() |
| encoder_output = self.vae.encode_dynamic(input_ids, attention_mask) |
| batch, width = input_ids.shape |
| logits = next(self.parameters()).new_zeros( |
| batch, width, self.config.vocab_size |
| ) |
| duration_logits = next(self.parameters()).new_zeros( |
| batch, |
| encoder_output.latents.shape[1], |
| self.config.max_span_length, |
| ) |
| chunk_spans = int( |
| getattr(self.generation_config, "span_causal_chunk_size", 1) |
| ) |
| for row in range(batch): |
| token_count = int(attention_mask[row].long().sum().item()) |
| latent_count = int( |
| encoder_output.latent_mask[row].long().sum().item() |
| ) |
| if token_count <= 0 or latent_count <= 0: |
| continue |
| row_logits, row_duration = self.vae.teacher_forced_logits( |
| input_ids[row, attention_mask[row]], |
| encoder_output.latents[row, :latent_count], |
| encoder_output.span_lengths[row, :latent_count], |
| chunk_spans=chunk_spans, |
| ) |
| logits[row, attention_mask[row]] = row_logits |
| duration_logits[row, :latent_count] = row_duration |
|
|
| loss = None |
| if labels is not None: |
| if labels.shape != input_ids.shape: |
| raise ValueError("labels must have the same shape as input_ids.") |
| supervised = attention_mask & labels.ge(0) |
| if bool(supervised.any()): |
| loss = F.cross_entropy( |
| logits[supervised].float(), |
| labels[supervised].clamp( |
| min=0, max=self.config.vocab_size - 1 |
| ), |
| ) |
| else: |
| loss = logits.sum() * 0.0 |
|
|
| output = Cola2ConditionalGenerationOutput( |
| loss=loss, |
| logits=logits, |
| latents=encoder_output.latents, |
| latent_mask=encoder_output.latent_mask, |
| span_lengths=encoder_output.span_lengths, |
| duration_logits=duration_logits, |
| boundary_logits=encoder_output.boundary_logits, |
| ) |
| if return_dict: |
| return output |
| values: tuple[torch.Tensor, ...] = ( |
| logits, |
| encoder_output.latents, |
| encoder_output.span_lengths, |
| duration_logits, |
| ) |
| return ((loss,) + values) if loss is not None else values |
|
|
| def prepare_inputs_for_generation( |
| self, |
| input_ids: torch.Tensor, |
| attention_mask: Optional[torch.Tensor] = None, |
| **kwargs: Any, |
| ) -> dict[str, Any]: |
| return { |
| "input_ids": input_ids, |
| "attention_mask": attention_mask, |
| **kwargs, |
| } |
|
|
| def _scale_latents(self, latents: torch.Tensor) -> torch.Tensor: |
| return ( |
| latents - self.config.shifting_factor |
| ) * self.config.scaling_factor |
|
|
| def _unscale_latents(self, latents: torch.Tensor) -> torch.Tensor: |
| return ( |
| latents / self.config.scaling_factor |
| ) + self.config.shifting_factor |
|
|
| def _dit_drift( |
| self, |
| prefix: torch.Tensor, |
| block: torch.Tensor, |
| block_timesteps: torch.Tensor, |
| *, |
| conditional: bool, |
| ) -> torch.Tensor: |
| if conditional and prefix.numel(): |
| states = torch.cat((prefix, block), dim=0) |
| timesteps = torch.cat( |
| ( |
| block_timesteps.new_zeros(prefix.shape[0]), |
| block_timesteps, |
| ), |
| dim=0, |
| ) |
| return self.dit(states, timesteps)[-block.shape[0] :] |
| return self.dit(block, block_timesteps) |
|
|
| def _denoise_block( |
| self, |
| prefix: torch.Tensor, |
| block_init: torch.Tensor, |
| clean_prefix_length: int, |
| generation_config: Cola2GenerationConfig, |
| ) -> torch.Tensor: |
| state = block_init.float() |
| clean = block_init[:clean_prefix_length].float().clone() |
| timesteps = _build_dit_timesteps( |
| schedule_t=self.config.diffusion_time, |
| steps=generation_config.dit_num_steps, |
| schedule=generation_config.dit_timestep_schedule, |
| power=generation_config.dit_timestep_power, |
| device=state.device, |
| ) |
| model_dtype = self.dit.txt_in.proj.weight.dtype |
| for current, following in zip( |
| timesteps[:-1], timesteps[1:], strict=True |
| ): |
| step = torch.full( |
| (state.shape[0],), |
| float(current.item()), |
| dtype=torch.float32, |
| device=state.device, |
| ) |
| if clean_prefix_length: |
| step[:clean_prefix_length] = 0.0 |
| state_model = state.to(model_dtype) |
| prefix_model = prefix.to(model_dtype) |
| step_model = step.to(model_dtype) |
| conditional = self._dit_drift( |
| prefix_model, |
| state_model, |
| step_model, |
| conditional=True, |
| ).float() |
| apply_cfg = ( |
| abs(generation_config.dit_guidance_scale - 1.0) > 1e-8 |
| and not ( |
| generation_config.dit_cfg_disable_empty_prefix |
| and prefix.numel() == 0 |
| ) |
| ) |
| if apply_cfg: |
| unconditional = self._dit_drift( |
| prefix_model, |
| state_model, |
| step_model, |
| conditional=False, |
| ).float() |
| drift = ( |
| generation_config.dit_guidance_scale |
| * (conditional - unconditional) |
| + unconditional |
| ) |
| else: |
| drift = conditional |
| dt = float((current - following).item()) / self.config.diffusion_time |
| state = state - drift * dt |
| if clean_prefix_length: |
| state[:clean_prefix_length] = clean |
| self.last_generation_stats["dit_denoise_steps"] += 1 |
| return state |
|
|
| @staticmethod |
| def _filter_logits( |
| logits: torch.Tensor, top_k: int, top_p: float |
| ) -> torch.Tensor: |
| filtered = logits.clone() |
| if top_k > 0: |
| k = min(int(top_k), filtered.shape[-1]) |
| threshold = torch.topk(filtered, k, dim=-1).values[:, -1:] |
| filtered.masked_fill_(filtered < threshold, float("-inf")) |
| if 0.0 < top_p < 1.0: |
| sorted_logits, sorted_indices = filtered.sort( |
| dim=-1, descending=True |
| ) |
| cumulative = sorted_logits.softmax(-1).cumsum(-1) |
| remove = cumulative > top_p |
| remove[:, 1:] = remove[:, :-1].clone() |
| remove[:, 0] = False |
| remove = torch.zeros_like(remove).scatter( |
| 1, sorted_indices, remove |
| ) |
| filtered.masked_fill_(remove, float("-inf")) |
| return filtered |
|
|
| def _sample_chunk( |
| self, |
| logits: torch.Tensor, |
| clean_ids: torch.Tensor, |
| generation_config: Cola2GenerationConfig, |
| generator: Optional[torch.Generator], |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| scores = logits.float().clone() |
| penalty = float(generation_config.repetition_penalty) |
| if penalty != 1.0 and clean_ids.numel(): |
| token_ids = clean_ids.unique().clamp( |
| 0, self.config.vocab_size - 1 |
| ) |
| selected = scores[:, token_ids] |
| selected = torch.where( |
| selected < 0, selected * penalty, selected / penalty |
| ) |
| scores[:, token_ids] = selected |
| if not generation_config.do_sample: |
| return scores.argmax(dim=-1), scores |
| temperature = float(generation_config.temperature) |
| if temperature <= 0: |
| raise ValueError("temperature must be positive when do_sample=True.") |
| scores = scores / temperature |
| scores = self._filter_logits( |
| scores, |
| int(generation_config.top_k), |
| float(generation_config.top_p), |
| ) |
| probabilities = scores.softmax(dim=-1) |
| sampled = torch.multinomial( |
| probabilities, num_samples=1, generator=generator |
| ).squeeze(-1) |
| return sampled, scores |
|
|
| def _span_causal_sample( |
| self, |
| clean_ids: torch.Tensor, |
| query_condition: torch.Tensor, |
| span_lengths: torch.Tensor, |
| generation_config: Cola2GenerationConfig, |
| generator: Optional[torch.Generator], |
| ) -> tuple[torch.Tensor, list[torch.Tensor]]: |
| committed = clean_ids |
| sampled_pieces: list[torch.Tensor] = [] |
| score_rows: list[torch.Tensor] = [] |
| bounds = torch.cat( |
| ( |
| span_lengths.new_zeros(1), |
| torch.cumsum(span_lengths, dim=0), |
| ), |
| dim=0, |
| ) |
| chunk_spans = generation_config.span_causal_chunk_size |
| for span_lo in range(0, span_lengths.numel(), chunk_spans): |
| span_hi = min(span_lo + chunk_spans, span_lengths.numel()) |
| token_lo = int(bounds[span_lo].item()) |
| token_hi = int(bounds[span_hi].item()) |
| condition = query_condition[token_lo:token_hi] |
| logits_conditional = self.vae.span_query_logits( |
| committed, condition, conditional=True |
| ).float() |
| weight = generation_config.decoder_guidance_scale |
| if abs(weight) > 1e-12: |
| logits_unconditional = self.vae.span_query_logits( |
| committed, condition, conditional=False |
| ).float() |
| logits = ( |
| (1.0 + weight) * logits_conditional |
| - weight * logits_unconditional |
| ) |
| else: |
| logits = logits_conditional |
| sampled, scores = self._sample_chunk( |
| logits, committed, generation_config, generator |
| ) |
| committed = torch.cat((committed, sampled), dim=0) |
| sampled_pieces.append(sampled) |
| score_rows.extend(scores.unbind(dim=0)) |
| self.last_generation_stats["decoder_chunks"] += 1 |
| if not sampled_pieces: |
| return clean_ids.new_zeros((0,)), score_rows |
| return torch.cat(sampled_pieces, dim=0), score_rows |
|
|
| def _generate_one( |
| self, |
| prompt_ids: torch.Tensor, |
| max_new_tokens: int, |
| generation_config: Cola2GenerationConfig, |
| generator: Optional[torch.Generator], |
| ) -> tuple[torch.Tensor, list[torch.Tensor]]: |
| if prompt_ids.numel() == 0: |
| raise ValueError("Cola2 generation requires a non-empty prompt.") |
| attention_mask = torch.ones_like(prompt_ids, dtype=torch.bool) |
| encoded = self.vae.encode_dynamic( |
| prompt_ids.unsqueeze(0), attention_mask.unsqueeze(0) |
| ) |
| latent_count = int(encoded.latent_mask[0].long().sum().item()) |
| current_unscaled = encoded.latents[0, :latent_count].float() |
| current_lengths = encoded.span_lengths[0, :latent_count].long() |
| committed_ids = prompt_ids.clone() |
| visible_tokens: list[int] = [] |
| visible_scores: list[torch.Tensor] = [] |
| eos_token_ids = generation_config.eos_token_id |
| if eos_token_ids is None: |
| eos_set: set[int] = set() |
| elif isinstance(eos_token_ids, int): |
| eos_set = {int(eos_token_ids)} |
| else: |
| eos_set = {int(value) for value in eos_token_ids} |
|
|
| while len(visible_tokens) < max_new_tokens: |
| scaled = self._scale_latents(current_unscaled) |
| block_size = self.config.dit_block_size |
| remainder = int(scaled.shape[0]) % block_size |
| if remainder: |
| prefix_length = int(scaled.shape[0]) - remainder |
| clean_block = scaled[-remainder:] |
| new_latent_count = block_size - remainder |
| else: |
| prefix_length = int(scaled.shape[0]) |
| clean_block = scaled.new_zeros((0, scaled.shape[-1])) |
| new_latent_count = block_size |
| prefix = scaled[:prefix_length] |
| noise = torch.randn( |
| new_latent_count, |
| self.config.latent_dim, |
| dtype=torch.float32, |
| device=scaled.device, |
| generator=generator, |
| ) |
| block_init = torch.cat((clean_block.float(), noise), dim=0) |
| denoised = self._denoise_block( |
| prefix, |
| block_init, |
| remainder, |
| generation_config, |
| ) |
| block_unscaled = self._unscale_latents(denoised) |
| candidate = torch.cat( |
| (current_unscaled[:prefix_length], block_unscaled), dim=0 |
| ) |
|
|
| _, all_duration_logits = self.vae.latent_decode( |
| candidate.to(self.vae.latent_in_proj.weight.dtype) |
| ) |
| self.last_generation_stats["duration_predictions"] += 1 |
| if generation_config.dynamic_duration: |
| predicted = all_duration_logits.argmax(dim=-1).long() + 1 |
| else: |
| fixed = max( |
| 1, |
| int( |
| round( |
| 1.0 / self.config.target_compression_rate |
| ) |
| ), |
| ) |
| predicted = torch.full( |
| (candidate.shape[0],), |
| min(fixed, self.config.max_span_length), |
| dtype=torch.long, |
| device=candidate.device, |
| ) |
| if remainder: |
| predicted[prefix_length : prefix_length + remainder] = ( |
| current_lengths[-remainder:] |
| ) |
| new_lengths = predicted[-new_latent_count:] |
| candidate_lengths = torch.cat( |
| (current_lengths, new_lengths), dim=0 |
| ) |
| condition, _ = self.vae.expand_and_refine( |
| candidate.to(self.vae.latent_in_proj.weight.dtype), |
| candidate_lengths, |
| ) |
| new_token_offset = int(current_lengths.sum().item()) |
| new_condition = condition[new_token_offset:] |
| sampled, sampled_scores = self._span_causal_sample( |
| committed_ids, |
| new_condition, |
| new_lengths, |
| generation_config, |
| generator, |
| ) |
| committed_ids = torch.cat((committed_ids, sampled), dim=0) |
| current_lengths = candidate_lengths |
|
|
| stop = False |
| for token, score in zip( |
| sampled.detach().tolist(), sampled_scores, strict=True |
| ): |
| if len(visible_tokens) >= max_new_tokens: |
| break |
| visible_tokens.append(int(token)) |
| visible_scores.append(score) |
| if int(token) in eos_set: |
| stop = True |
| break |
|
|
| if generation_config.latent_feedback: |
| current_unscaled = self.vae.encode_given_spans( |
| committed_ids, current_lengths |
| ).float() |
| self.last_generation_stats["feedback_rounds"] += 1 |
| else: |
| current_unscaled = candidate.float() |
| self.last_generation_stats["generation_rounds"] += 1 |
| if stop: |
| break |
|
|
| return torch.tensor( |
| visible_tokens, dtype=torch.long, device=prompt_ids.device |
| ), visible_scores |
|
|
| def _resolve_generation_config( |
| self, |
| generation_config: Optional[Union[Cola2GenerationConfig, Mapping[str, Any]]], |
| kwargs: dict[str, Any], |
| ) -> Cola2GenerationConfig: |
| if generation_config is None: |
| config = copy.deepcopy(self.generation_config) |
| elif isinstance(generation_config, Mapping): |
| config = Cola2GenerationConfig.from_dict(dict(generation_config)) |
| elif isinstance(generation_config, Cola2GenerationConfig): |
| config = copy.deepcopy(generation_config) |
| else: |
| config = Cola2GenerationConfig.from_dict( |
| generation_config.to_dict() |
| ) |
| unused = config.update(**kwargs) |
| if unused: |
| names = ", ".join(sorted(unused)) |
| raise ValueError(f"Unused generation arguments: {names}.") |
| config.validate() |
| if int(config.num_return_sequences) != 1: |
| raise ValueError("Cola2 currently supports num_return_sequences=1.") |
| if int(config.num_beams) != 1: |
| raise ValueError( |
| "Beam search is not defined for block-parallel span-causal decoding." |
| ) |
| return config |
|
|
| @torch.no_grad() |
| def generate( |
| self, |
| input_ids: Optional[torch.Tensor] = None, |
| generation_config: Optional[ |
| Union[Cola2GenerationConfig, Mapping[str, Any]] |
| ] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| generator: Optional[torch.Generator] = None, |
| streamer: Any = None, |
| **kwargs: Any, |
| ) -> Union[torch.Tensor, GenerateDecoderOnlyOutput]: |
| """Generate continuations with Full-DiT and span-causal decoding.""" |
|
|
| if input_ids is None: |
| bos = self.config.bos_token_id |
| if bos is None: |
| raise ValueError("input_ids or bos_token_id is required.") |
| input_ids = torch.tensor( |
| [[bos]], dtype=torch.long, device=self.device |
| ) |
| if input_ids.ndim != 2: |
| raise ValueError("input_ids must have shape [batch, sequence].") |
| if attention_mask is None: |
| attention_mask = input_ids.ne(self.config.pad_token_id) |
| if attention_mask.shape != input_ids.shape: |
| raise ValueError("attention_mask must match input_ids.") |
|
|
| config = self._resolve_generation_config(generation_config, kwargs) |
| self.last_generation_stats = { |
| "prompt_encodes": int(input_ids.shape[0]), |
| "dit_denoise_steps": 0, |
| "duration_predictions": 0, |
| "decoder_chunks": 0, |
| "feedback_rounds": 0, |
| "generation_rounds": 0, |
| } |
| generated_rows: list[torch.Tensor] = [] |
| score_rows: list[list[torch.Tensor]] = [] |
| prompt_lengths: list[int] = [] |
| for row in range(input_ids.shape[0]): |
| valid_ids = input_ids[row, attention_mask[row].bool()] |
| prompt_lengths.append(int(valid_ids.numel())) |
| if config.max_new_tokens is not None: |
| budget = int(config.max_new_tokens) |
| else: |
| budget = max(0, int(config.max_length) - int(valid_ids.numel())) |
| generated, scores = self._generate_one( |
| valid_ids, budget, config, generator |
| ) |
| generated_rows.append(generated) |
| score_rows.append(scores) |
|
|
| max_generated = max((row.numel() for row in generated_rows), default=0) |
| pad_id = config.pad_token_id |
| if pad_id is None: |
| pad_id = ( |
| config.eos_token_id |
| if isinstance(config.eos_token_id, int) |
| else self.config.pad_token_id |
| ) |
| if pad_id is None: |
| pad_id = 0 |
| sequences = input_ids.new_full( |
| (input_ids.shape[0], input_ids.shape[1] + max_generated), |
| int(pad_id), |
| ) |
| sequences[:, : input_ids.shape[1]] = input_ids |
| for row, generated in enumerate(generated_rows): |
| sequences[ |
| row, |
| input_ids.shape[1] : input_ids.shape[1] + generated.numel(), |
| ] = generated |
|
|
| if streamer is not None: |
| if input_ids.shape[0] != 1: |
| raise ValueError("streamer is supported only for batch size 1.") |
| for token in generated_rows[0]: |
| streamer.put(token.reshape(1, 1).cpu()) |
| streamer.end() |
|
|
| if not config.return_dict_in_generate: |
| return sequences |
| scores_tuple = None |
| if config.output_scores: |
| score_steps: list[torch.Tensor] = [] |
| for step in range(max_generated): |
| rows = [] |
| for row_scores in score_rows: |
| if step < len(row_scores): |
| rows.append(row_scores[step]) |
| else: |
| filler = next(self.parameters()).new_full( |
| (self.config.vocab_size,), float("-inf") |
| ) |
| if 0 <= int(pad_id) < self.config.vocab_size: |
| filler[int(pad_id)] = 0.0 |
| rows.append(filler) |
| score_steps.append(torch.stack(rows, dim=0)) |
| scores_tuple = tuple(score_steps) |
| return GenerateDecoderOnlyOutput( |
| sequences=sequences, |
| scores=scores_tuple, |
| ) |
|
|
| @staticmethod |
| def _read_checkpoint_state( |
| source: Union[str, Path, Mapping[str, torch.Tensor]], |
| map_location: Union[str, torch.device] = "cpu", |
| ) -> dict[str, torch.Tensor]: |
| if isinstance(source, Mapping): |
| raw: Any = dict(source) |
| else: |
| path = Path(source) |
| if path.is_dir(): |
| safe_index = path / "model.safetensors.index.json" |
| bin_index = path / "pytorch_model.bin.index.json" |
| if safe_index.exists() or bin_index.exists(): |
| index_path = ( |
| safe_index if safe_index.exists() else bin_index |
| ) |
| with index_path.open("r", encoding="utf-8") as handle: |
| shard_names = sorted( |
| set(json.load(handle)["weight_map"].values()) |
| ) |
| state: dict[str, torch.Tensor] = {} |
| for shard_name in shard_names: |
| shard = Cola2ForConditionalGeneration._read_checkpoint_state( |
| path / shard_name, map_location=map_location |
| ) |
| overlap = state.keys() & shard.keys() |
| if overlap: |
| raise RuntimeError( |
| f"Duplicate checkpoint keys: {sorted(overlap)[:5]}" |
| ) |
| state.update(shard) |
| return state |
| candidates = ( |
| path / "model.safetensors", |
| path / "pytorch_model.bin", |
| ) |
| path = next( |
| (candidate for candidate in candidates if candidate.exists()), |
| path, |
| ) |
| if path.is_dir(): |
| raise FileNotFoundError( |
| f"No model.safetensors or pytorch_model.bin in {path}." |
| ) |
| if path.suffix == ".safetensors": |
| from safetensors.torch import load_file |
|
|
| raw = load_file(str(path), device=str(map_location)) |
| else: |
| raw = torch.load( |
| path, map_location=map_location, weights_only=True |
| ) |
| if isinstance(raw, Mapping): |
| for key in ("state_dict", "model"): |
| nested = raw.get(key) |
| if isinstance(nested, Mapping) and all( |
| torch.is_tensor(value) for value in nested.values() |
| ): |
| raw = nested |
| break |
| if not isinstance(raw, Mapping) or not all( |
| isinstance(key, str) and torch.is_tensor(value) |
| for key, value in raw.items() |
| ): |
| raise TypeError("Checkpoint must resolve to a tensor state dictionary.") |
| return dict(raw) |
|
|
| def load_strict_state_dict( |
| self, |
| source: Union[str, Path, Mapping[str, torch.Tensor]], |
| *, |
| map_location: Union[str, torch.device] = "cpu", |
| prefix: Optional[str] = None, |
| assign: bool = False, |
| ) -> "Cola2ForConditionalGeneration": |
| """Load package weights with exact key and shape validation. |
| |
| ``module.`` and ``_orig_mod.`` wrappers are stripped only when they wrap |
| every key. Other conversions must pass an explicit ``prefix`` so a |
| partial or accidentally wrong native checkpoint cannot load silently. |
| """ |
|
|
| state = self._read_checkpoint_state(source, map_location) |
| if prefix is not None: |
| if not all(key.startswith(prefix) for key in state): |
| raise RuntimeError( |
| f"Not every checkpoint key starts with prefix {prefix!r}." |
| ) |
| state = {key[len(prefix) :]: value for key, value in state.items()} |
| for wrapper in ("module.", "_orig_mod."): |
| if state and all(key.startswith(wrapper) for key in state): |
| state = { |
| key[len(wrapper) :]: value for key, value in state.items() |
| } |
|
|
| expected = self.state_dict() |
| missing = sorted(expected.keys() - state.keys()) |
| unexpected = sorted(state.keys() - expected.keys()) |
| mismatched = sorted( |
| ( |
| key, |
| tuple(state[key].shape), |
| tuple(expected[key].shape), |
| ) |
| for key in expected.keys() & state.keys() |
| if tuple(state[key].shape) != tuple(expected[key].shape) |
| ) |
| if missing or unexpected or mismatched: |
| parts = [] |
| if missing: |
| parts.append(f"missing={missing[:20]}") |
| if unexpected: |
| parts.append(f"unexpected={unexpected[:20]}") |
| if mismatched: |
| parts.append(f"shape_mismatches={mismatched[:20]}") |
| raise RuntimeError( |
| "Strict Cola2 state_dict validation failed: " + "; ".join(parts) |
| ) |
| self.load_state_dict(state, strict=True, assign=assign) |
| return self |
|
|
|
|
| def build_export_components( |
| exported_config: Mapping[str, Any], |
| ) -> dict[str, nn.Module]: |
| """Build strict-validation components from a legacy export config. |
| |
| This compatibility factory accepts either a public Cola2 config or an |
| earlier nested, path-free config and constructs the same DiT and VAE |
| modules. |
| """ |
|
|
| raw = dict(exported_config) |
| if "dit_config" not in raw or "vae_config" not in raw: |
| config = Cola2Config.from_dict(raw) |
| return {"dit": Cola2DiT(config), "vae": Cola2CompressionVAE(config)} |
|
|
| dit = dict(raw["dit_config"]) |
| vae = dict(raw["vae_config"]) |
| backbone = dict(vae.get("qwen3_backbone_config") or {}) |
| backbone = dict(backbone.get("text_config", backbone)) |
| model_config = Cola2Config( |
| vocab_size=int(vae["vocab_size"]), |
| encoder_embedding_size=( |
| int(backbone.get("hidden_size", vae["dim"])) |
| if bool(vae.get("qwen_encoder_token_weights", False)) |
| else int(vae["dim"]) |
| ), |
| encoder_hidden_size=int(vae["dim"]), |
| encoder_intermediate_size=int(vae["ffn_dim"]), |
| encoder_num_attention_heads=int(vae["num_heads"]), |
| encoder_num_hidden_layers=int(vae["token_encoder_num_blocks"]), |
| pool_num_hidden_layers=int(vae["pool_num_blocks"]), |
| latent_encoder_num_hidden_layers=int( |
| vae["latent_encoder_num_blocks"] |
| ), |
| latent_decoder_num_hidden_layers=int( |
| vae["latent_decoder_num_blocks"] |
| ), |
| latent_dim=int(vae["latent_dim"]), |
| use_variational_encoder=bool(vae["use_variation"]), |
| boundary_score_alpha=float(vae["scorer_score_alpha"]), |
| boundary_score_threshold=float(vae["scorer_score_tau0"]), |
| target_compression_rate=float(vae["rho_target"]), |
| min_span_length=int(vae["lmin"]), |
| max_span_length=int(vae["lmax"]), |
| latent_block_causal=bool(vae["latent_block_causal"]), |
| latent_block_size=int(vae["latent_block_size"]), |
| dit_hidden_size=int(dit["txt_dim"]), |
| dit_intermediate_size=int(dit["txt_dim"]) |
| * int(dit["expand_ratio"]), |
| dit_num_hidden_layers=int(dit["num_layers"]), |
| dit_num_attention_heads=int(dit["heads"]), |
| dit_head_dim=int(dit["head_dim"]), |
| dit_rope_dim=int( |
| dit["rope_dim"][0] |
| if isinstance(dit["rope_dim"], list) |
| else dit["rope_dim"] |
| ), |
| dit_block_size=int(dit["block_size"]), |
| decoder_hidden_size=int( |
| backbone.get("hidden_size", vae["decoder_dim"]) |
| ), |
| decoder_intermediate_size=int( |
| backbone.get( |
| "intermediate_size", |
| int(vae["decoder_ffn_dim"]) // 2, |
| ) |
| ), |
| decoder_num_hidden_layers=int(vae["token_decoder_num_blocks"]), |
| decoder_num_attention_heads=int( |
| backbone.get( |
| "num_attention_heads", vae["decoder_num_heads"] |
| ) |
| ), |
| decoder_num_key_value_heads=int( |
| backbone.get( |
| "num_key_value_heads", vae["decoder_num_heads"] |
| ) |
| ), |
| decoder_head_dim=int( |
| backbone.get( |
| "head_dim", |
| int(vae["decoder_dim"]) |
| // int(vae["decoder_num_heads"]), |
| ) |
| ), |
| condition_num_attention_heads=int(vae["decoder_num_heads"]), |
| decoder_condition_num_hidden_layers=int(vae["decoder_cond_blocks"]), |
| decoder_lm_head_bias=bool(vae.get("bias", True)), |
| duration_hidden_size=( |
| None |
| if vae.get("duration_hidden_dim") is None |
| else int(vae["duration_hidden_dim"]) |
| ), |
| local_position_dim=int(vae["loc_dim"]), |
| hidden_dropout=float(vae["dropout"]), |
| attention_dropout=float(vae["attn_dropout"]), |
| layer_norm_eps=float(vae["layer_norm_eps"]), |
| dit_layer_norm_eps=float(dit["norm_eps"]), |
| encoder_rope_theta=float(vae["rope_theta"]), |
| condition_rope_theta=float(vae["rope_theta"]), |
| decoder_rope_theta=float( |
| backbone.get("rope_theta", 1_000_000.0) |
| ), |
| dit_rope_theta=10_000.0, |
| rope_theta=float(backbone.get("rope_theta", 1_000_000.0)), |
| initializer_range=float(vae["init_std"]), |
| scaling_factor=float(raw.get("latent_scaling_factor", 1.0)), |
| shifting_factor=float(raw.get("latent_shifting_factor", 0.0)), |
| pad_token_id=int(raw["pad_token_id"]), |
| mask_token_id=int(raw["mask_token_id"]), |
| eos_token_id=int(raw["eos_token_id"]), |
| bos_token_id=int(raw["eos_token_id"]), |
| ) |
| return { |
| "dit": Cola2DiT(model_config), |
| "vae": Cola2CompressionVAE(model_config), |
| } |
|
|
