# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 """Sparse Autoencoder for image/video tokenization. This module provides: - DiagonalGaussianDistribution: VAE latent distribution - SparseTransformerBase: Base transformer class for encoder/decoder - Encoder: Vision transformer encoder with pretrained weight loading - Decoder: Transformer decoder with multiscale support - AutoencoderKLConfig: Configuration dataclass for the autoencoder - AutoencoderKL: Base VAE model with KL loss for encoding/decoding """ import os from dataclasses import dataclass from typing import Any, Literal import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.loaders.single_file_model import FromOriginalModelMixin from diffusers.models.autoencoders.vae import DecoderOutput from diffusers.models.modeling_utils import ModelMixin from diffusers.utils.accelerate_utils import apply_forward_hook from diffusers.utils.torch_utils import randn_tensor from einops import rearrange from loguru import logger as logging from transformers import Siglip2Model from cosmos_framework.model.tokenizer.models.dense_backends import ( resolve_dense_backend, run_batched_block_stack, run_varlen_block_stack, ) from cosmos_framework.model.tokenizer.models.modules import ( AbsolutePositionEmbedder, LayerNorm32, LearnedPositionEmbedder, LearnedPositionEmbedder4D, RMSNorm32, SparseLinear, SparseMultiheadAttentionPoolingHead, SparseTensor, SparseTransformerBlock, ) from cosmos_framework.model.tokenizer.models.modules.quantizers import ( FSQ, LFQ, RQBottleneck, levels_from_codebook_size, ) from cosmos_framework.model.tokenizer.models.utils import ( batch_tensor_to_sparse, reconstruct_from_temporal_slices, sparse_to_batched_tensor, sparse_to_img_list, ) from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs, stack_with_bounded_inputs # ============================================================================= # Helper Functions # ============================================================================= def multiply_all_factors(config: dict[int, dict[str, Any]]) -> list[int]: """Multiply all factors from a multiscale config. Args: config: Dict mapping layer indices to config dicts with 'factor' keys. Returns: List of multiplied factors. """ config_values = list(config.values()) if not config_values: return [] first_factor = config_values[0]["factor"] result = [1] * len(first_factor) for value_dict in config_values: factor_list = value_dict["factor"] for i in range(len(result)): result[i] *= factor_list[i] return result def expand_multiscale_cfg(config: dict[int, dict[str, Any]], num_layer: int) -> dict[int, dict[str, Any]]: """Expand multiscale config to cover layers that have corresponding configs. Args: config: Dict with layer numbers as keys and config dicts as values. num_layer: Total number of layers (0 to num_layer-1). Returns: Dict with config for layers that have a corresponding config key. """ expanded_config = {} config_keys = sorted(config.keys()) for layer in range(num_layer): best_key = None for key in config_keys: if key <= layer: best_key = key else: break if best_key is not None: if layer == best_key: expanded_config[layer] = config[best_key].copy() else: layer_config = {} for key, value in config[best_key].items(): if key not in ["factor", "channel_duplication"]: layer_config[key] = value expanded_config[layer] = layer_config return expanded_config def _sparse_tensor_to_dense_batch_tokens(sparse_tensor: SparseTensor) -> tuple[torch.Tensor, bool] | None: """Convert a uniform sparse batch into dense `[B, S, D]` token features. Returns: Tuple of (`dense_feats`, `used_reshape_fast_path`) when every batch has the same token count, otherwise `None`. """ q_seqlen = sparse_tensor.get_batch_seq_lens() if not q_seqlen: return None seq_len = q_seqlen[0] if any(current_seq_len != seq_len for current_seq_len in q_seqlen): return None batch_size = sparse_tensor.shape[0] feature_shape = sparse_tensor.feats.shape[1:] contiguous_layout = all( batch_slice.start == batch_idx * seq_len and batch_slice.stop == (batch_idx + 1) * seq_len for batch_idx, batch_slice in enumerate(sparse_tensor.layout) ) if contiguous_layout: dense_feats = sparse_tensor.feats.reshape(batch_size, seq_len, *feature_shape) return dense_feats, True dense_feats = stack_with_bounded_inputs( [sparse_tensor.feats[batch_slice] for batch_slice in sparse_tensor.layout], dim=0 ) # [B,S,*F] return dense_feats, False def _dense_batch_tokens_to_flat_features(dense_feats: torch.Tensor, used_reshape_fast_path: bool) -> torch.Tensor: """Restore dense `[B, S, D]` features to the sparse flat token order.""" if used_reshape_fast_path: return dense_feats.reshape(-1, *dense_feats.shape[2:]) return cat_with_bounded_inputs(dense_feats.unbind(dim=0), dim=0) def _crop_temporal_slices_to_ownership( slices: list[SparseTensor], *, frame_batch_size: int, frame_batch_strides: int, ) -> list[SparseTensor]: """Crop overlapping temporal slices to non-overlapping ownership regions. Adjacent slices split their actual overlap at the midpoint, with odd-length overlaps assigning the extra timestep to the later slice. This keeps one owner per reconstructed timestep even when the final slice is truncated. """ if frame_batch_size <= frame_batch_strides or len(slices) <= 1: return slices slice_bounds: list[tuple[int, int, int | None] | None] = [] owned_starts: list[int] = [] owned_ends: list[int] = [] for slice_tensor in slices: min_frame, max_frame = slice_tensor.get_temporal_range() temporal_offset = slice_tensor.get_spatial_cache("temporal_offset") if min_frame is None or max_frame is None: slice_bounds.append(None) owned_starts.append(0) owned_ends.append(0) continue global_start = min_frame + (0 if temporal_offset is None else temporal_offset) global_end = max_frame + 1 + (0 if temporal_offset is None else temporal_offset) slice_bounds.append((global_start, global_end, temporal_offset)) owned_starts.append(global_start) owned_ends.append(global_end) for slice_index in range(len(slices) - 1): current_bounds = slice_bounds[slice_index] next_bounds = slice_bounds[slice_index + 1] if current_bounds is None or next_bounds is None: continue current_start, current_end, _ = current_bounds next_start, next_end, _ = next_bounds overlap_start = max(current_start, next_start) overlap_end = min(current_end, next_end) if overlap_start >= overlap_end: continue boundary = overlap_start + (overlap_end - overlap_start) // 2 owned_ends[slice_index] = min(owned_ends[slice_index], boundary) owned_starts[slice_index + 1] = max(owned_starts[slice_index + 1], boundary) cropped_slices: list[SparseTensor] = [] for slice_tensor, bounds, owned_start, owned_end in zip(slices, slice_bounds, owned_starts, owned_ends): if bounds is None: cropped_slices.append(slice_tensor) continue _, _, temporal_offset = bounds offset_value = 0 if temporal_offset is None else temporal_offset local_start = owned_start - offset_value local_end = owned_end - offset_value cropped_slice = slice_tensor.slice_temporal_range(local_start, local_end, adjust_temporal=False) if temporal_offset is not None: cropped_slice.register_spatial_cache("temporal_offset", temporal_offset) cropped_slices.append(cropped_slice) return cropped_slices def _validate_no_legacy_attention_block_cfg( cfg: dict[int, dict[str, Any]], num_blocks: int, ) -> None: """Reject removed sparse-attention overrides inside per-block configs.""" for block_index in range(num_blocks): block_cfg = cfg.get(block_index, {}) block_attn_mode = block_cfg.get("attn_mode") if block_attn_mode is not None and block_attn_mode != "full": raise ValueError( "Tokenizer sparse attention only supports full attention. " f"Got attn_mode={block_attn_mode!r} for block {block_index}." ) if block_cfg.get("window_size") is not None: raise ValueError( "Tokenizer sparse attention no longer supports windowed attention. " f"Got window_size={block_cfg['window_size']!r} for block {block_index}." ) for legacy_key in ("shift_sequence", "shift_window", "serialize_mode"): if block_cfg.get(legacy_key) is not None: raise ValueError( "Tokenizer sparse attention only supports full attention. " f"Got legacy override {legacy_key}={block_cfg[legacy_key]!r} for block {block_index}." ) # ============================================================================= # Distribution Classes # ============================================================================= class DiagonalGaussianDistribution: """Diagonal Gaussian distribution for VAE latent space. Parameterized by mean and log-variance, supports sampling and KL divergence. """ def __init__( self, parameters: torch.Tensor | SparseTensor, deterministic: bool = False, ): """Initialize distribution from parameters. Args: parameters: Tensor containing concatenated [mean, logvar] along dim 1. deterministic: If True, sampling returns the mean. """ if isinstance(parameters, SparseTensor): self.raw_parameters = parameters self.parameters = parameters.feats else: self.raw_parameters = None self.parameters = parameters self.original_dtype = self.parameters.dtype self.parameters = self.parameters.to(torch.float32) self.mean, self.logvar = torch.chunk(self.parameters, 2, dim=1) self.deterministic = deterministic self.std = torch.exp(0.5 * self.logvar) self.var = torch.exp(self.logvar) if self.deterministic: self.var = self.std = torch.zeros_like( self.mean, device=self.parameters.device, dtype=self.parameters.dtype ) def sample(self, generator: torch.Generator | None = None) -> torch.Tensor: """Sample from the distribution. Args: generator: Optional random generator for reproducibility. Returns: Sampled tensor. """ sample = randn_tensor( self.mean.shape, generator=generator, device=self.parameters.device, dtype=self.parameters.dtype, ) x = self.mean + self.std * sample return x.to(self.original_dtype) def kl(self, other: "DiagonalGaussianDistribution | None" = None) -> torch.Tensor: """Compute KL divergence. Args: other: Other distribution to compare against (defaults to N(0,1)). Returns: KL divergence tensor. """ reduce_dims = tuple(range(1, self.mean.ndim)) if self.deterministic: return torch.zeros( self.mean.shape[0], device=self.parameters.device, dtype=self.parameters.dtype, ) else: if other is None: return 0.5 * torch.sum( torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, dim=reduce_dims, ) else: return 0.5 * torch.sum( torch.pow(self.mean - other.mean, 2) / other.var + self.var / other.var - 1.0 - self.logvar + other.logvar, dim=reduce_dims, ) def nll(self, sample: torch.Tensor, dims: tuple[int, ...] = (1, 2, 3)) -> torch.Tensor: """Compute negative log-likelihood. Args: sample: Sample to evaluate. dims: Dimensions to sum over. Returns: NLL tensor. """ if self.deterministic: return torch.zeros( self.mean.shape[0], device=self.parameters.device, dtype=self.parameters.dtype, ) logtwopi = np.log(2.0 * np.pi) reduce_dims = dims if sample.ndim > max(dims, default=0) else tuple(range(1, sample.ndim)) return 0.5 * torch.sum( logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, dim=reduce_dims, ) def mode(self) -> torch.Tensor: """Return the mode (mean) of the distribution.""" return self.mean.to(self.original_dtype) # ============================================================================= # Transformer Base Classes # ============================================================================= class SparseTransformerBase(nn.Module): """Base class for sparse transformer encoder and decoder.""" def __init__( self, in_channels: int, model_channels: int, num_blocks: int, num_heads: int | None = None, num_head_channels: int | None = 64, mlp_channels: int = 2048, pe_mode: Literal["ape", "rope", "learned", "learned4d", "joint"] = "rope", use_checkpoint: bool = False, qk_rms_norm: bool = False, use_bias: bool = False, use_rms_norm: bool = True, ln_affine: bool = True, multiscale: dict[int, dict[str, Any]] | None = None, dense_train_backend: Literal["disabled", "varlen", "batched", "auto"] = "disabled", ): """Initialize SparseTransformerBase. Args: in_channels: Number of input channels. model_channels: Hidden dimension size. num_blocks: Number of transformer blocks. num_heads: Number of attention heads. num_head_channels: Channels per head (used if num_heads is None). mlp_channels: MLP hidden size. pe_mode: Position embedding mode. use_checkpoint: Whether to use gradient checkpointing. qk_rms_norm: Whether to apply RMS norm to Q/K. use_bias: Whether to use bias in linear layers. use_rms_norm: Whether to use RMSNorm (vs LayerNorm). ln_affine: Whether to use affine parameters in LayerNorm. multiscale: Multiscale configuration. dense_train_backend: Optional dense tensor backend to use for the supported fixed-shape train-time subset. """ super().__init__() self.in_channels = in_channels self.model_channels = model_channels self.num_blocks = num_blocks self.num_heads = num_heads or model_channels // num_head_channels self.mlp_channels = mlp_channels self.pe_mode = pe_mode self.use_checkpoint = use_checkpoint self.qk_rms_norm = qk_rms_norm self.use_bias = use_bias self.use_rms_norm = use_rms_norm self.dense_train_backend = dense_train_backend if pe_mode == "ape": self.pos_embedder = AbsolutePositionEmbedder(model_channels) elif pe_mode == "learned4d": self.pos_embedder = LearnedPositionEmbedder4D(model_channels) elif pe_mode in ["joint", "learned"]: self.pos_embedder = LearnedPositionEmbedder(model_channels) self.input_layer = SparseLinear(in_channels, model_channels, bias=use_bias) if multiscale is not None: cfg = expand_multiscale_cfg(multiscale, num_blocks) else: cfg = {} _validate_no_legacy_attention_block_cfg(cfg=cfg, num_blocks=num_blocks) self.blocks = nn.ModuleList( [ SparseTransformerBlock( cfg[i]["model_channels"] if i in cfg else model_channels, num_heads=cfg[i]["num_heads"] if i in cfg else self.num_heads, mlp_channels=cfg[i]["mlp_channels"] if i in cfg else mlp_channels, use_checkpoint=self.use_checkpoint, use_rope=(pe_mode in ["rope", "joint"]), qk_rms_norm=self.qk_rms_norm, use_bias=self.use_bias, use_rms_norm=self.use_rms_norm, multiscale=cfg.get(i, None), ln_affine=ln_affine, layer_idx=i, ) for i in range(num_blocks) ] ) @property def device(self) -> torch.device: """Return the device of the model.""" return next(self.parameters()).device def _resolve_dense_train_backend(self) -> Literal["varlen", "batched"] | None: """Resolve the optional dense-train backend for the current execution mode.""" if self.dense_train_backend == "disabled": return None return resolve_dense_backend( self.dense_train_backend, use_compile=torch.compiler.is_compiling(), ) def forward( self, x: SparseTensor, kv_cache: dict[str, SparseTensor] | None = None, kv_cache_size: int | None = None, kv_cache_detach: bool = True, temporal_causal_mask: bool = False, collect_hidden_states: bool = True, ) -> tuple[SparseTensor, list[SparseTensor] | None, dict[str, SparseTensor]]: """Forward pass through transformer blocks. Args: x: Input SparseTensor. kv_cache: Optional KV cache for autoregressive generation. kv_cache_size: Maximum KV cache size. kv_cache_detach: Whether tensors stored in the KV cache should be detached. temporal_causal_mask: Whether to apply temporal-causal, same-timestep bidirectional self-attention in decoder blocks. collect_hidden_states: Whether to retain per-block hidden states for downstream consumers such as concat_latent or multiscale outputs. Returns: Tuple of (output, hidden_states, updated_kv_cache). """ hs: list[SparseTensor] | None = [] if collect_hidden_states else None input_dtype = next(self.input_layer.parameters()).dtype if x.dtype != input_dtype: x = x.to(input_dtype) h = self.input_layer(x) if self.pe_mode == "ape": h = h + self.pos_embedder(x.coords[:, 1:]).to(h.dtype) elif self.pe_mode in ["learned", "learned4d", "joint"]: h = h + self.pos_embedder(x) block_dtype = next(self.blocks.parameters()).dtype if h.dtype != block_dtype: h = h.to(block_dtype) if hs is not None: hs.append(h) dense_train_backend = self._resolve_dense_train_backend() if self.training else None can_use_dense_train_path = dense_train_backend is not None and all( ( hs is None, kv_cache_size is None, not temporal_causal_mask, not h.has_special_tokens(), all(getattr(block, "multiscale", None) is None for block in self.blocks), all(getattr(block.attn, "_type", None) == "self" for block in self.blocks), ) ) if can_use_dense_train_path: q_freqs_cis: torch.Tensor | None = None blocks_with_rope = [block for block in self.blocks if getattr(block.attn, "use_rope", False)] if blocks_with_rope: rope_configs = { ( block.attn.rope.head_dim, block.attn.rope.pos_cls_token, ) for block in blocks_with_rope } if len(rope_configs) != 1: can_use_dense_train_path = False else: q_freqs_cis = blocks_with_rope[0].attn.rope.compute_freqs_cis( h.coords[:, 1:], has_special_tokens=False, ) dense_batch_tokens = _sparse_tensor_to_dense_batch_tokens(h) if can_use_dense_train_path else None if dense_batch_tokens is None: can_use_dense_train_path = False if can_use_dense_train_path and dense_batch_tokens is not None: dense_feats, used_reshape_fast_path = dense_batch_tokens q_seqlen = h.get_batch_seq_lens() max_q_seqlen = q_seqlen[0] if q_seqlen else 0 cu_seqlens_q = h.get_cu_seqlens(device=h.device) if dense_train_backend == "varlen": dense_feats = run_varlen_block_stack( self.blocks, dense_feats, q_seqlen=q_seqlen, cu_seqlens_q=cu_seqlens_q, max_q_seqlen=max_q_seqlen, q_freqs_cis=q_freqs_cis, ) else: dense_feats = run_batched_block_stack( self.blocks, dense_feats, q_freqs_cis=q_freqs_cis, ) h = h.replace(_dense_batch_tokens_to_flat_features(dense_feats, used_reshape_fast_path)) empty_kv_cache: dict[str, SparseTensor] if kv_cache is None: empty_kv_cache = {} else: kv_cache.clear() empty_kv_cache = kv_cache return h, hs, empty_kv_cache can_use_tensor_no_cache_fast_path = ( not self.training and hs is None and kv_cache_size is None and not temporal_causal_mask and all(callable(getattr(block, "forward_tensor_no_cache", None)) for block in self.blocks) and all(getattr(block, "multiscale", None) is None for block in self.blocks) and all(getattr(block.attn, "_type", None) == "self" for block in self.blocks) ) if can_use_tensor_no_cache_fast_path: q_freqs_cis: torch.Tensor | None = None if ( any(getattr(block.attn, "_debug_capture_enabled", False) for block in self.blocks) and h.has_special_tokens() ): raise AssertionError( "Debug attention capture for the encoder tensor fast path requires image-like tokens without " "special sparse tokens." ) blocks_with_rope = [block for block in self.blocks if getattr(block.attn, "use_rope", False)] if blocks_with_rope: rope_configs = { ( block.attn.rope.head_dim, block.attn.rope.pos_cls_token, ) for block in blocks_with_rope } if len(rope_configs) != 1: can_use_tensor_no_cache_fast_path = False else: q_freqs_cis = blocks_with_rope[0].attn.rope.compute_freqs_cis( h.coords[:, 1:], has_special_tokens=h.has_special_tokens(), ) if can_use_tensor_no_cache_fast_path: q_seqlen = h.get_batch_seq_lens() max_q_seqlen = max(q_seqlen) if q_seqlen else 0 cu_seqlens_q = h.get_cu_seqlens(device=h.device) feats = h.feats for block in self.blocks: feats = block.forward_tensor_no_cache( feats, q_seqlen=q_seqlen, cu_seqlens_q=cu_seqlens_q, max_q_seqlen=max_q_seqlen, q_freqs_cis=q_freqs_cis, ) h = h.replace(feats) empty_kv_cache: dict[str, SparseTensor] if kv_cache is None: empty_kv_cache = {} else: kv_cache.clear() empty_kv_cache = kv_cache return h, hs, empty_kv_cache can_use_tensor_flat_kv_fast_path = ( not self.training and hs is None and kv_cache_size is not None and kv_cache_size > 0 and not temporal_causal_mask and h.shape[0] == 1 and not h.has_special_tokens() and all(callable(getattr(block, "forward_tensor_flat_kv", None)) for block in self.blocks) and all(getattr(block, "multiscale", None) is None for block in self.blocks) and all(getattr(block.attn, "_type", None) == "self" for block in self.blocks) ) if can_use_tensor_flat_kv_fast_path: q_seqlen = h.get_batch_seq_lens() max_q_seqlen = max(q_seqlen) if q_seqlen else 0 cu_seqlens_q = h.get_cu_seqlens(device=h.device) current_times = h.coords[:, 1].contiguous() q_freqs_cis: torch.Tensor | None = None blocks_with_rope = [block for block in self.blocks if getattr(block.attn, "use_rope", False)] if blocks_with_rope: rope_configs = { ( block.attn.rope.head_dim, block.attn.rope.pos_cls_token, ) for block in blocks_with_rope } if len(rope_configs) != 1: can_use_tensor_flat_kv_fast_path = False else: q_freqs_cis = blocks_with_rope[0].attn.rope.compute_freqs_cis( h.coords[:, 1:], has_special_tokens=h.has_special_tokens(), ) if kv_cache is not None: for block_idx, block in enumerate(self.blocks): block_cache = kv_cache.get(f"block_{block_idx}", {}) if not block.attn._has_compatible_flat_kv_cache_state(block_cache): can_use_tensor_flat_kv_fast_path = False break if can_use_tensor_flat_kv_fast_path: if kv_cache is None: kv_cache = {} feats = h.feats for block_idx, block in enumerate(self.blocks): block_cache_key = f"block_{block_idx}" if block_cache_key not in kv_cache: kv_cache[block_cache_key] = {} block_cache = kv_cache[block_cache_key] feats, updated_block_cache = block.forward_tensor_flat_kv( feats, current_times=current_times, q_seqlen=q_seqlen, cu_seqlens_q=cu_seqlens_q, max_q_seqlen=max_q_seqlen, kv_cache=block_cache, kv_cache_size=kv_cache_size, kv_cache_detach=kv_cache_detach, q_freqs_cis=q_freqs_cis, ) kv_cache[block_cache_key] = updated_block_cache h = h.replace(feats) return h, hs, kv_cache if kv_cache_size is None and not temporal_causal_mask: for block in self.blocks: forward_no_cache = getattr(block, "forward_no_cache", None) if callable(forward_no_cache): h = forward_no_cache(h) else: h, _ = block( h, kv_cache=None, kv_cache_size=None, kv_cache_detach=kv_cache_detach, temporal_causal_mask=False, ) if hs is not None: hs.append(h) empty_kv_cache: dict[str, SparseTensor] if kv_cache is None: empty_kv_cache = {} else: kv_cache.clear() empty_kv_cache = kv_cache return h, hs, empty_kv_cache if kv_cache is None: kv_cache = {} for block_idx, block in enumerate(self.blocks): block_cache_key = f"block_{block_idx}" if block_cache_key not in kv_cache: kv_cache[block_cache_key] = {} block_cache = kv_cache[block_cache_key] h, updated_block_cache = block( h, kv_cache=block_cache, kv_cache_size=kv_cache_size, kv_cache_detach=kv_cache_detach, temporal_causal_mask=temporal_causal_mask, ) kv_cache[block_cache_key] = updated_block_cache if hs is not None: hs.append(h) return h, hs, kv_cache class Encoder(SparseTransformerBase): """Sparse Transformer Encoder with optional pretrained weight loading.""" def __init__( self, in_channels: int, model_channels: int, num_blocks: int, num_heads: int | None = None, num_head_channels: int = 64, mlp_channels: int = 2048, pe_mode: Literal["ape", "rope", "learned", "learned4d", "joint"] = "ape", use_checkpoint: bool = False, qk_rms_norm: bool = False, use_bias: bool = False, use_rms_norm: bool = True, pretrained_model: nn.Module | None = None, concat_latent: list[int] | None = None, use_head: bool = False, dense_train_backend: Literal["disabled", "varlen", "batched", "auto"] = "disabled", ): """Initialize Encoder. Args: in_channels: Number of input channels. model_channels: Hidden dimension size. num_blocks: Number of transformer blocks. num_heads: Number of attention heads. num_head_channels: Channels per head. mlp_channels: MLP hidden size. pe_mode: Position embedding mode. use_checkpoint: Whether to use gradient checkpointing. qk_rms_norm: Whether to apply RMS norm to Q/K. use_bias: Whether to use bias in linear layers. use_rms_norm: Whether to use RMSNorm. pretrained_model: Optional pretrained vision model for weight init. concat_latent: List of layer indices to concatenate for output. use_head: Whether to apply pooling head for image features. dense_train_backend: Optional dense tensor backend for the supported fixed-shape train-time subset. """ super().__init__( in_channels, model_channels, num_blocks, num_heads, num_head_channels, mlp_channels, pe_mode, use_checkpoint, qk_rms_norm, use_bias=use_bias, use_rms_norm=use_rms_norm, dense_train_backend=dense_train_backend, ) self.concat_latent = concat_latent self.use_head = use_head if use_rms_norm: self.post_layernorm = RMSNorm32(model_channels, eps=1e-6) else: self.post_layernorm = LayerNorm32(model_channels, elementwise_affine=True, eps=1e-6) self.head = SparseMultiheadAttentionPoolingHead( hidden_size=model_channels, num_attention_heads=self.num_heads, intermediate_size=mlp_channels, use_bias=use_bias, use_rms_norm=use_rms_norm, qk_rms_norm=qk_rms_norm, ) if pretrained_model is not None: self._initialize_weights(pretrained_model=pretrained_model) def _initialize_weights(self, pretrained_model: nn.Module) -> None: """Initialize weights from pretrained vision model. Args: pretrained_model: Pretrained vision model (e.g., SigLIP2). """ pretrained_weight = pretrained_model.state_dict() # Copy patch embedding weights if "embeddings.patch_embedding.weight" in pretrained_weight: s1 = pretrained_weight["embeddings.patch_embedding.weight"].shape[1] self.input_layer.weight.data[:, :s1].copy_(pretrained_weight["embeddings.patch_embedding.weight"]) if hasattr(self.input_layer, "bias") and self.input_layer.bias is not None: self.input_layer.bias.data.copy_(pretrained_weight["embeddings.patch_embedding.bias"]) # Copy position embedding if ( "embeddings.position_embedding.weight" in pretrained_weight and getattr(self, "pos_embedder", None) is not None ): self.pos_embedder.position_embedding.weight.data.copy_( pretrained_weight["embeddings.position_embedding.weight"] ) # Copy transformer block weights for block_idx, block in enumerate(self.blocks): prefix = f"encoder.layers.{block_idx}" # Layer norms if f"{prefix}.layer_norm1.weight" in pretrained_weight: block.norm1.weight.data.copy_(pretrained_weight[f"{prefix}.layer_norm1.weight"]) if f"{prefix}.layer_norm1.bias" in pretrained_weight: block.norm1.bias.data.copy_(pretrained_weight[f"{prefix}.layer_norm1.bias"]) if f"{prefix}.layer_norm2.weight" in pretrained_weight: block.norm2.weight.data.copy_(pretrained_weight[f"{prefix}.layer_norm2.weight"]) if f"{prefix}.layer_norm2.bias" in pretrained_weight: block.norm2.bias.data.copy_(pretrained_weight[f"{prefix}.layer_norm2.bias"]) # Attention QKV if hasattr(block.attn, "to_qkv"): q_weight = pretrained_weight.get(f"{prefix}.self_attn.q_proj.weight") k_weight = pretrained_weight.get(f"{prefix}.self_attn.k_proj.weight") v_weight = pretrained_weight.get(f"{prefix}.self_attn.v_proj.weight") if q_weight is not None and k_weight is not None and v_weight is not None: qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) block.attn.to_qkv.weight.data.copy_(qkv_weight) q_bias = pretrained_weight.get(f"{prefix}.self_attn.q_proj.bias") k_bias = pretrained_weight.get(f"{prefix}.self_attn.k_proj.bias") v_bias = pretrained_weight.get(f"{prefix}.self_attn.v_proj.bias") if q_bias is not None and k_bias is not None and v_bias is not None: qkv_bias = torch.cat([q_bias, k_bias, v_bias], dim=0) block.attn.to_qkv.bias.data.copy_(qkv_bias) # Attention output projection if hasattr(block.attn, "to_out"): if f"{prefix}.self_attn.out_proj.weight" in pretrained_weight: block.attn.to_out.weight.data.copy_(pretrained_weight[f"{prefix}.self_attn.out_proj.weight"]) if f"{prefix}.self_attn.out_proj.bias" in pretrained_weight: block.attn.to_out.bias.data.copy_(pretrained_weight[f"{prefix}.self_attn.out_proj.bias"]) # MLP weights if hasattr(block.mlp, "mlp") and len(block.mlp.mlp) > 2: if f"{prefix}.mlp.fc1.weight" in pretrained_weight: block.mlp.mlp[0].weight.data.copy_(pretrained_weight[f"{prefix}.mlp.fc1.weight"]) if f"{prefix}.mlp.fc1.bias" in pretrained_weight: block.mlp.mlp[0].bias.data.copy_(pretrained_weight[f"{prefix}.mlp.fc1.bias"]) if f"{prefix}.mlp.fc2.weight" in pretrained_weight: block.mlp.mlp[2].weight.data.copy_(pretrained_weight[f"{prefix}.mlp.fc2.weight"]) if f"{prefix}.mlp.fc2.bias" in pretrained_weight: block.mlp.mlp[2].bias.data.copy_(pretrained_weight[f"{prefix}.mlp.fc2.bias"]) # Post layer norm if hasattr(self, "post_layernorm"): if "post_layernorm.weight" in pretrained_weight: self.post_layernorm.weight.data.copy_(pretrained_weight["post_layernorm.weight"]) if "post_layernorm.bias" in pretrained_weight: self.post_layernorm.bias.data.copy_(pretrained_weight["post_layernorm.bias"]) # Head weights if hasattr(self, "head"): self._initialize_head_weights(pretrained_weight) logging.info("Successfully initialized encoder weights from pretrained model") def _initialize_head_weights(self, pretrained_weight: dict[str, torch.Tensor]) -> None: """Initialize pooling head weights from pretrained model. Args: pretrained_weight: State dict from pretrained model. """ head_prefix = "head" if f"{head_prefix}.probe" in pretrained_weight: probe_weight = pretrained_weight[f"{head_prefix}.probe"].squeeze() self.head.probe.data.copy_(probe_weight) # Cross-attention weights if hasattr(self.head.attention, "to_q") and hasattr(self.head.attention, "to_kv"): if f"{head_prefix}.attention.in_proj_weight" in pretrained_weight: in_proj_weight = pretrained_weight[f"{head_prefix}.attention.in_proj_weight"] q_weight, k_weight, v_weight = in_proj_weight.chunk(3, dim=0) self.head.attention.to_q.weight.data.copy_(q_weight) kv_weight = torch.cat([k_weight, v_weight], dim=0) self.head.attention.to_kv.weight.data.copy_(kv_weight) if f"{head_prefix}.attention.in_proj_bias" in pretrained_weight: in_proj_bias = pretrained_weight[f"{head_prefix}.attention.in_proj_bias"] q_bias, k_bias, v_bias = in_proj_bias.chunk(3, dim=0) self.head.attention.to_q.bias.data.copy_(q_bias) kv_bias = torch.cat([k_bias, v_bias], dim=0) self.head.attention.to_kv.bias.data.copy_(kv_bias) # Attention output projection if f"{head_prefix}.attention.out_proj.weight" in pretrained_weight: self.head.attention.to_out.weight.data.copy_(pretrained_weight[f"{head_prefix}.attention.out_proj.weight"]) if f"{head_prefix}.attention.out_proj.bias" in pretrained_weight: self.head.attention.to_out.bias.data.copy_(pretrained_weight[f"{head_prefix}.attention.out_proj.bias"]) # Head layer norm if f"{head_prefix}.layernorm.weight" in pretrained_weight: self.head.layernorm.weight.data.copy_(pretrained_weight[f"{head_prefix}.layernorm.weight"]) if f"{head_prefix}.layernorm.bias" in pretrained_weight: self.head.layernorm.bias.data.copy_(pretrained_weight[f"{head_prefix}.layernorm.bias"]) # Head MLP if hasattr(self.head.mlp, "mlp") and len(self.head.mlp.mlp) > 2: if f"{head_prefix}.mlp.fc1.weight" in pretrained_weight: self.head.mlp.mlp[0].weight.data.copy_(pretrained_weight[f"{head_prefix}.mlp.fc1.weight"]) if f"{head_prefix}.mlp.fc1.bias" in pretrained_weight: self.head.mlp.mlp[0].bias.data.copy_(pretrained_weight[f"{head_prefix}.mlp.fc1.bias"]) if f"{head_prefix}.mlp.fc2.weight" in pretrained_weight: self.head.mlp.mlp[2].weight.data.copy_(pretrained_weight[f"{head_prefix}.mlp.fc2.weight"]) if f"{head_prefix}.mlp.fc2.bias" in pretrained_weight: self.head.mlp.mlp[2].bias.data.copy_(pretrained_weight[f"{head_prefix}.mlp.fc2.bias"]) def forward( self, x: SparseTensor, kv_cache: dict[str, SparseTensor] | None = None, ) -> tuple[SparseTensor, dict[str, SparseTensor]] | tuple[SparseTensor, torch.Tensor, dict[str, SparseTensor]]: """Forward pass through encoder. Args: x: Input SparseTensor. kv_cache: Optional KV cache. Returns: If use_head: (encoded, pooled_features, kv_cache). Otherwise: (encoded, kv_cache). """ collect_hidden_states = self.concat_latent is not None h, hs, kv_cache = super().forward( x, kv_cache=kv_cache, collect_hidden_states=collect_hidden_states, ) h = h.replace(self.post_layernorm(h.feats)) if self.concat_latent is not None: assert hs is not None, "concat_latent requires hidden states to be collected" assert all(0 <= idx < len(hs) for idx in self.concat_latent), ( f"concat_latent indices {self.concat_latent} must be within [0, {len(hs) - 1}]" ) hs_concat = [hs[i].feats for i in self.concat_latent] hs_concat = torch.cat(hs_concat, dim=1) h = h.replace(hs_concat) if self.use_head: pooler_output = self.head(h) return h, pooler_output.feats, kv_cache else: return h, kv_cache class Decoder(SparseTransformerBase): """Sparse Transformer Decoder with multiscale output support.""" def __init__( self, in_channels: int, out_channels: int, model_channels: int, num_blocks: int, num_heads: int | None = None, num_head_channels: int = 64, mlp_channels: int = 2048, pe_mode: Literal["ape", "rope", "learned", "learned4d", "joint"] = "ape", use_checkpoint: bool = False, qk_rms_norm: bool = False, use_bias: bool = False, use_rms_norm: bool = True, multiscale: dict[int, dict[str, Any]] | None = None, multiscale_outputs: dict[int, dict[str, Any]] | None = None, dense_train_backend: Literal["disabled", "varlen", "batched", "auto"] = "disabled", ): """Initialize Decoder. Args: in_channels: Number of input channels (latent dimension). out_channels: Number of output channels. model_channels: Hidden dimension size. num_blocks: Number of transformer blocks. num_heads: Number of attention heads. num_head_channels: Channels per head. mlp_channels: MLP hidden size. pe_mode: Position embedding mode. use_checkpoint: Whether to use gradient checkpointing. qk_rms_norm: Whether to apply RMS norm to Q/K. use_bias: Whether to use bias in linear layers. use_rms_norm: Whether to use RMSNorm. multiscale: Multiscale expansion configuration. multiscale_outputs: Multiscale output layer configuration. dense_train_backend: Optional dense tensor backend for the supported fixed-shape train-time subset. """ super().__init__( in_channels, model_channels, num_blocks, num_heads, num_head_channels, mlp_channels, pe_mode, use_checkpoint, qk_rms_norm, use_bias=use_bias, use_rms_norm=use_rms_norm, multiscale=multiscale, dense_train_backend=dense_train_backend, ) self.multiscale = multiscale self.multiscale_outputs = multiscale_outputs # Select last multiscale channel configuration if multiscale is not None and multiscale_outputs is None: last_layer = max(multiscale.keys()) model_channels = multiscale[last_layer]["model_channels"] if use_rms_norm: self.out_norm = RMSNorm32(model_channels, eps=1e-6) else: self.out_norm = LayerNorm32(model_channels, elementwise_affine=False, eps=1e-6) self.out_layer = SparseLinear( model_channels, out_channels, bias=False, ) if multiscale is not None: self.recover_factor = multiply_all_factors(multiscale) if multiscale_outputs is not None: self.multiscale_out_layers = nn.ModuleList() self.multiscale_out_norms = nn.ModuleList() self.recover_factors = [] for cfg in multiscale_outputs[1:]: if use_rms_norm: out_norm = RMSNorm32(cfg["model_channels"], eps=1e-6) else: out_norm = LayerNorm32(cfg["model_channels"], elementwise_affine=False, eps=1e-6) out_layer = SparseLinear( cfg["model_channels"], cfg["out_channels"], bias=False, ) scale = {k: v for k, v in multiscale.items() if k < cfg["layer_id"]} recover_factor = multiply_all_factors(scale) self.multiscale_out_norms.append(out_norm) self.multiscale_out_layers.append(out_layer) self.recover_factors.append(recover_factor) def forward( self, x: SparseTensor, kv_cache: dict[str, SparseTensor] | None = None, kv_cache_size: int | None = None, kv_cache_detach: bool = True, temporal_causal_mask: bool = False, ) -> tuple[SparseTensor, dict[str, SparseTensor]]: """Forward pass through decoder. Args: x: Input SparseTensor (latent representation). kv_cache: Optional KV cache. kv_cache_size: Maximum KV cache size. kv_cache_detach: Whether tensors stored in the KV cache should be detached. temporal_causal_mask: Whether to apply temporal-causal, same-timestep bidirectional self-attention in decoder blocks. Returns: Tuple of (decoded output, updated kv_cache). """ collect_hidden_states = self.multiscale_outputs is not None h, hs, kv_cache = super().forward( x, kv_cache=kv_cache, kv_cache_size=kv_cache_size, kv_cache_detach=kv_cache_detach, temporal_causal_mask=temporal_causal_mask, collect_hidden_states=collect_hidden_states, ) if self.multiscale_outputs is not None: assert hs is not None, "multiscale_outputs requires hidden states to be collected" requested_layer_ids = [cfg["layer_id"] for cfg in self.multiscale_outputs] assert all(0 <= layer_id < len(hs) for layer_id in requested_layer_ids), ( f"multiscale layer_ids {requested_layer_ids} must be within [0, {len(hs) - 1}]" ) h = hs[self.multiscale_outputs[0]["layer_id"]] h = h.replace(self.out_norm(h.feats)) h = self.out_layer(h) for i, cfg in enumerate(self.multiscale_outputs[1:]): h_i = hs[cfg["layer_id"]] h_i = h_i.replace(self.multiscale_out_norms[i](h_i.feats)) h_i = self.multiscale_out_layers[i](h_i) if self.multiscale is not None: h_i = h_i.shrink_by_factors(self.recover_factors[i]) h = h + h_i h = h.replace(h.feats / len(self.multiscale_outputs)) else: h = h.replace(self.out_norm(h.feats)) h = self.out_layer(h) if self.multiscale is not None: h = h.shrink_by_factors(self.recover_factor) return h, kv_cache @dataclass class AutoencoderKLConfig: """Configuration for AutoencoderKL model.""" patch_size: tuple[int, int, int] = (1, 8, 8) in_channels: int = 192 out_channels: int = 192 latent_channels: int = 8 encoder_model_channels: int = 768 encoder_num_blocks: int = 12 encoder_num_heads: int | None = None encoder_mlp_channels: float = 2048 encoder_pe_mode: str = "rope" encoder_qk_rms_norm: bool = True encoder_use_bias: bool = True encoder_use_rms_norm: bool = False decoder_model_channels: int = 768 decoder_num_blocks: int = 12 decoder_num_heads: int | None = None decoder_mlp_channels: float = 2048 decoder_pe_mode: str = "rope" decoder_qk_rms_norm: bool = True decoder_use_bias: bool = True decoder_use_rms_norm: bool = False decoder_multiscale: dict[int, dict[str, Any]] | None = None decoder_multiscale_outputs: dict[int, dict[str, Any]] | None = None use_decoder: bool = True use_quantizer: bool = False quantizer_type: Literal["fsq", "lfq", "rq"] = "rq" quantizer_codebook_size: int = 65536 quantizer_num_codebooks: int = 1 quantizer_feature_dim: int = 48 quantizer_chunk_size: int = 1 use_text_alignment: bool = False use_post_text_alignment: bool = False use_text_decoder: bool = False use_post_text_decoder: bool = False spatial_pool_size: int = 1 text_decoder_model_name: str | None = None # e.g., "Qwen/Qwen3-0.6B" or local path text_decoder_family: str = "qwen3" text_decoder_gradient_checkpointing: bool = True text_decoder_packed_attention_backend: Literal["sdpa", "natten"] = "sdpa" encoder_use_checkpoint: bool | None = None decoder_use_checkpoint: bool | None = None encoder_dense_train_backend: Literal["disabled", "varlen", "batched", "auto"] = "disabled" decoder_dense_train_backend: Literal["disabled", "varlen", "batched", "auto"] = "disabled" decoder_temporal_mode: Literal["bidirectional", "causal_mask", "causal"] = "bidirectional" decoder_temporal_query_latent_steps: int = 1 decoder_temporal_cache_latent_steps: int | None = None decoder_temporal_detach_cache: bool = True inference_num_sample_frames_batch_size: int = 16 inference_num_sample_frames_stride: int = 16 inference_kv_cache_size: int = 0 task_decode_runtime_settings: dict[str, dict[str, int]] | None = None use_vf_loss: bool = False freeze_encoder: bool = False pretrained_model_name: str | None = None concat_latent: list | None = None random_num_sample_frames_batch_sizes: list[int] | None = None task_random_num_sample_frames_batch_sizes: dict[str, list[int]] | None = None use_dual_latent: bool = False class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin): """Sparse Autoencoder with KL loss for image/video encoding. This model encodes images/videos into latent representations using a sparse transformer architecture with SigLIP2 pretrained features. Args: patch_size: Patch size for tokenization (T, H, W). in_channels: Number of input channels per patch. out_channels: Number of output channels per patch. latent_channels: Number of latent channels. encoder_*: Encoder configuration parameters. decoder_*: Decoder configuration parameters. use_quantizer: Whether to use quantization. quantizer_type: Type of quantizer ("fsq", "lfq", "rq"). use_text_alignment: Whether to enable text alignment. pretrained_model_name: HuggingFace model name for SigLIP2. """ _supports_gradient_checkpointing = True _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"] @register_to_config def __init__( self, patch_size: tuple[int, int, int] = (1, 8, 8), in_channels: int = 192, out_channels: int = 192, latent_channels: int = 8, encoder_model_channels: int = 768, encoder_num_blocks: int = 12, encoder_num_heads: int | None = None, encoder_mlp_channels: float = 2048, encoder_pe_mode: str = "rope", encoder_qk_rms_norm: bool = True, encoder_use_bias: bool = False, encoder_use_rms_norm: bool = True, decoder_model_channels: int = 768, decoder_num_blocks: int = 12, decoder_num_heads: int | None = None, decoder_mlp_channels: float = 2048, decoder_pe_mode: str = "rope", decoder_qk_rms_norm: bool = True, decoder_use_bias: bool = False, decoder_use_rms_norm: bool = True, decoder_multiscale: dict[int, dict[str, Any]] | None = None, decoder_multiscale_outputs: dict[int, dict[str, Any]] | None = None, use_decoder: bool = True, use_quantizer: bool = False, quantizer_type: Literal["fsq", "lfq", "rq"] = "rq", quantizer_codebook_size: int = 16384, quantizer_num_codebooks: int = 1, quantizer_feature_dim: int = 48, quantizer_chunk_size: int = 1, use_text_alignment: bool = False, use_post_text_alignment: bool = False, use_text_decoder: bool = False, use_post_text_decoder: bool = False, spatial_pool_size: int = 1, text_decoder_model_name: str | None = None, text_decoder_family: str = "qwen3", text_decoder_gradient_checkpointing: bool = True, text_decoder_packed_attention_backend: Literal["sdpa", "natten"] = "sdpa", encoder_use_checkpoint: bool | None = None, decoder_use_checkpoint: bool | None = None, encoder_dense_train_backend: Literal["disabled", "varlen", "batched", "auto"] = "disabled", decoder_dense_train_backend: Literal["disabled", "varlen", "batched", "auto"] = "disabled", decoder_temporal_mode: Literal["bidirectional", "causal_mask", "causal"] = "bidirectional", decoder_temporal_query_latent_steps: int = 1, decoder_temporal_cache_latent_steps: int | None = None, decoder_temporal_detach_cache: bool = True, inference_num_sample_frames_batch_size: int = 16, inference_num_sample_frames_stride: int = 16, inference_kv_cache_size: int = 0, task_decode_runtime_settings: dict[str, dict[str, int]] | None = None, use_vf_loss: bool = False, freeze_encoder: bool = False, pretrained_model_name: str | None = None, concat_latent: list | None = None, random_num_sample_frames_batch_sizes: list[int] | None = None, task_random_num_sample_frames_batch_sizes: dict[str, list[int]] | None = None, use_dual_latent: bool = False, use_checkpoint: bool = True, ) -> None: super().__init__() self.use_checkpoint = use_checkpoint self.encoder_use_checkpoint = use_checkpoint if encoder_use_checkpoint is None else encoder_use_checkpoint self.decoder_use_checkpoint = use_checkpoint if decoder_use_checkpoint is None else decoder_use_checkpoint self.encoder_dense_train_backend = encoder_dense_train_backend self.decoder_dense_train_backend = decoder_dense_train_backend self.patch_size = patch_size self.use_text_alignment = use_text_alignment self.use_post_text_alignment = use_post_text_alignment self.use_text_decoder = use_text_decoder self.use_post_text_decoder = use_post_text_decoder self.spatial_pool_size = spatial_pool_size self.text_decoder_family = text_decoder_family self.text_decoder_packed_attention_backend: Literal["sdpa", "natten"] = text_decoder_packed_attention_backend self.decoder_temporal_mode = decoder_temporal_mode self.decoder_temporal_query_latent_steps = decoder_temporal_query_latent_steps self.decoder_temporal_cache_latent_steps = decoder_temporal_cache_latent_steps self.decoder_temporal_detach_cache = decoder_temporal_detach_cache self.inference_num_sample_frames_batch_size = inference_num_sample_frames_batch_size self.inference_num_sample_frames_stride = inference_num_sample_frames_stride self.inference_kv_cache_size = inference_kv_cache_size self.task_decode_runtime_settings = task_decode_runtime_settings self.use_quantizer = use_quantizer self.quantizer_type = quantizer_type self.quantizer_codebook_size = quantizer_codebook_size self.quantizer_num_codebooks = quantizer_num_codebooks self.quantizer_feature_dim = quantizer_feature_dim self.use_vf_loss = use_vf_loss self.freeze_encoder = freeze_encoder self.quantizer_chunk_size = quantizer_chunk_size self.use_dual_latent = use_dual_latent self.latent_channels = latent_channels self.in_channels = in_channels self.out_channels = out_channels self.random_num_sample_frames_batch_sizes = random_num_sample_frames_batch_sizes self.task_random_num_sample_frames_batch_sizes = task_random_num_sample_frames_batch_sizes self.num_sample_frames_batch_size = 16 self.num_sample_frames_stride = 12 self.kv_cache_size = 4 self._logged_decoder_temporal_plan = False # Load SigLIP2 pretrained model (text encoder always needed for text alignment) # Prefer HF_HUB_CACHE (set by configure_hf_cache_env to HF_HOME/hub), then fall back to # HF_HOME/hub (where transformers stores models), then HF_HOME itself. This handles jobs # that set HF_HOME but not HF_HUB_CACHE (e.g. VFM training via submit_helper). _hf_home = os.environ.get("HF_HOME") hf_cache_dir = ( os.environ.get("HF_HUB_CACHE") or (os.path.join(_hf_home, "hub") if _hf_home else None) or _hf_home ) local_files_only = hf_cache_dir is not None pretrained_model = None pretrained_vision_model = None need_pretrained_for_text = use_text_alignment or use_post_text_alignment if pretrained_model_name is not None or need_pretrained_for_text: model_name = pretrained_model_name or "google/siglip2-so400m-patch16-naflex" pretrained_model = Siglip2Model.from_pretrained( model_name, cache_dir=hf_cache_dir, local_files_only=local_files_only, ) if pretrained_model_name is not None: pretrained_vision_model = pretrained_model.vision_model # Initialize encoder self.encoder = Encoder( in_channels=in_channels, model_channels=encoder_model_channels, num_blocks=encoder_num_blocks, num_heads=encoder_num_heads, mlp_channels=encoder_mlp_channels, pe_mode=encoder_pe_mode, qk_rms_norm=encoder_qk_rms_norm, use_bias=encoder_use_bias, use_rms_norm=encoder_use_rms_norm, pretrained_model=pretrained_vision_model, concat_latent=concat_latent, use_checkpoint=self.encoder_use_checkpoint, dense_train_backend=self.encoder_dense_train_backend, ) # Initialize teacher encoder (frozen) — only needed for ITD loss # which requires use_text_alignment=True. Skip when only using text decoder (ITG) # to save ~3GB GPU memory. if use_text_alignment: self.teacher_encoder = Encoder( in_channels=in_channels, model_channels=encoder_model_channels, num_blocks=encoder_num_blocks, num_heads=encoder_num_heads, mlp_channels=encoder_mlp_channels, pe_mode="learned", qk_rms_norm=False, use_bias=True, use_rms_norm=False, pretrained_model=pretrained_vision_model, use_head=True, dense_train_backend="disabled", ) self.teacher_encoder.requires_grad_(False) # Projection layer if concat_latent is not None: proj_in_channels = encoder_model_channels * len(concat_latent) else: proj_in_channels = encoder_model_channels self.proj = SparseLinear(proj_in_channels, 2 * latent_channels) # Initialize quantizer if use_quantizer: if self.quantizer_type == "lfq": self.quantizer = LFQ( dim=self.quantizer_feature_dim // self.quantizer_chunk_size, codebook_size=self.quantizer_codebook_size // self.quantizer_chunk_size, num_codebooks=self.quantizer_num_codebooks, sample_minimization_weight=1.0, batch_maximization_weight=1.0, token_factorization=False, factorized_bits=[9, 9], ) elif self.quantizer_type == "rq": self.quantizer = RQBottleneck( latent_shape=(16, 16, latent_channels), code_shape=(16, 16, self.quantizer_num_codebooks), n_embed=self.quantizer_codebook_size, decay=0.99, shared_codebook=True, restart_unused_codes=True, ) else: levels, _ = levels_from_codebook_size(self.quantizer_codebook_size) self.quantizer = FSQ( levels=levels, dim=self.quantizer_feature_dim // self.quantizer_chunk_size, num_codebooks=self.quantizer_num_codebooks, ) # Initialize decoder if use_decoder: self.decoder = Decoder( in_channels=latent_channels, out_channels=out_channels, model_channels=decoder_model_channels, num_blocks=decoder_num_blocks, num_heads=decoder_num_heads, mlp_channels=decoder_mlp_channels, qk_rms_norm=decoder_qk_rms_norm, use_bias=decoder_use_bias, use_rms_norm=decoder_use_rms_norm, pe_mode=decoder_pe_mode, use_checkpoint=self.decoder_use_checkpoint, multiscale=decoder_multiscale, multiscale_outputs=decoder_multiscale_outputs, dense_train_backend=self.decoder_dense_train_backend, ) if use_decoder and use_dual_latent: self.discrete_decoder = Decoder( in_channels=latent_channels, out_channels=out_channels, model_channels=decoder_model_channels, num_blocks=decoder_num_blocks, num_heads=decoder_num_heads, mlp_channels=decoder_mlp_channels, qk_rms_norm=decoder_qk_rms_norm, use_bias=decoder_use_bias, use_rms_norm=decoder_use_rms_norm, pe_mode=decoder_pe_mode, use_checkpoint=self.decoder_use_checkpoint, multiscale=decoder_multiscale, multiscale_outputs=decoder_multiscale_outputs, dense_train_backend=self.decoder_dense_train_backend, ) self.use_slicing = False self.logit_bias = None self.logit_scale = None # Text alignment if self.use_text_alignment: self.text_encoder = pretrained_model.text_model self.text_encoder.requires_grad_(False) self.logit_scale = nn.Parameter(pretrained_model.logit_scale.data.clone()) self.logit_bias = nn.Parameter(pretrained_model.logit_bias.data.clone()) self.teacher_logit_scale = nn.Parameter(pretrained_model.logit_scale.data.clone()) self.teacher_logit_bias = nn.Parameter(pretrained_model.logit_bias.data.clone()) self.teacher_logit_scale.requires_grad = False self.teacher_logit_bias.requires_grad = False else: self.teacher_logit_bias = None self.teacher_logit_scale = None if self.use_post_text_alignment: self.post_logit_scale = nn.Parameter(pretrained_model.logit_scale.data.clone()) self.post_logit_bias = nn.Parameter(pretrained_model.logit_bias.data.clone()) self.post_alignment_proj = SparseLinear(latent_channels, proj_in_channels) self.post_alignment_head = SparseMultiheadAttentionPoolingHead( hidden_size=encoder_model_channels, num_attention_heads=encoder_num_heads, intermediate_size=encoder_mlp_channels, use_bias=encoder_use_bias, use_rms_norm=encoder_use_rms_norm, qk_rms_norm=False, ) else: self.post_logit_scale = None self.post_logit_bias = None # Text decoder (configured causal LM for image-to-text generation) if self.use_text_decoder and text_decoder_model_name is not None: from cosmos_framework.model.tokenizer.models.text_decoder import ( TextDecoderWrapper, get_text_decoder_family_spec, ) self.text_decoder_wrapper = TextDecoderWrapper( model_name=text_decoder_model_name, image_hidden_size=encoder_model_channels, spatial_pool_size=spatial_pool_size, gradient_checkpointing=text_decoder_gradient_checkpointing, packed_attention_backend=text_decoder_packed_attention_backend, family_spec=get_text_decoder_family_spec( family=text_decoder_family, model_name=text_decoder_model_name, ), ) else: self.text_decoder_wrapper = None def decode_text( self, input_ids: torch.Tensor, image_feats: "SparseTensor", image_patch_indices: torch.Tensor, segment_ids: torch.Tensor | None = None, return_hidden_states: bool = False, ) -> tuple[Any, int]: """Decode text from image features using the configured text decoder. Passes encoder output (x_no_proj) through spatial merger + causal LM. Args: input_ids: [B, S] text token IDs with vision placeholders. image_feats: SparseTensor from encoder output (x_no_proj). image_patch_indices: [N_pooled] flat indices into [B*S]. segment_ids: [B, S] segment IDs for packed sequences. Enables segment-isolated attention with per-segment position reset. return_hidden_states: Return decoder hidden states instead of full vocabulary logits for chunked training loss projection. Returns: Tuple of decoder output [B,S,Vocab] or [B,S,D] and pooled-token count. """ if self.text_decoder_wrapper is None: raise RuntimeError("Text decoder not initialized. Set use_text_decoder=True and text_decoder_model_name.") return self.text_decoder_wrapper( input_ids=input_ids, image_feats_tensor=image_feats.feats, image_coords=image_feats.coords, image_patch_indices=image_patch_indices, image_layout=image_feats.layout if hasattr(image_feats, "layout") else None, segment_ids=segment_ids, return_hidden_states=return_hidden_states, ) def encode_text(self, text: torch.Tensor, normalize: bool = False) -> torch.Tensor: """Encode text using the text encoder. Args: text: Input text token IDs. normalize: Whether to L2 normalize the output. Returns: Text embeddings. """ text_outputs = self.text_encoder(text) pooled_output = text_outputs.pooler_output return F.normalize(pooled_output, dim=-1) if normalize else pooled_output def _set_gradient_checkpointing(self, module, value=False): if isinstance(module, (Encoder, Decoder)): module.gradient_checkpointing = value def enable_slicing(self): """Enable sliced VAE decoding for memory efficiency.""" self.use_slicing = True def disable_slicing(self): """Disable sliced VAE decoding.""" self.use_slicing = False def _frame_count_to_latent_steps(self, frame_count: int, name: str, *, allow_zero: bool = False) -> int: """Convert a raw frame count to latent temporal steps with strict divisibility checks.""" frame_count = int(frame_count) temporal_patch_size = int(self.patch_size[0]) if temporal_patch_size <= 0: raise ValueError(f"patch_size[0] must be positive, got {temporal_patch_size}.") if frame_count == 0 and allow_zero: return 0 if frame_count < 0 and allow_zero: raise ValueError(f"{name} must be non-negative, got {frame_count}.") if frame_count <= 0: raise ValueError(f"{name} must be positive, got {frame_count}.") if frame_count % temporal_patch_size != 0: raise ValueError(f"{name} must be divisible by patch_size[0]={temporal_patch_size}, got {frame_count}.") return frame_count // temporal_patch_size def _encode( self, x: SparseTensor, normalize: bool = False, compute_image_feat: bool = True, ) -> tuple[SparseTensor, torch.Tensor | None, SparseTensor]: """Internal encode method with temporal batching. Args: x: Input SparseTensor. normalize: Whether to normalize image features. compute_image_feat: Whether to compute image features via encoder.head. Set to False for reconstruction-only tasks to save computation. Returns: Tuple of (projected latent, image features or None, unprojected encoder output). """ if self.training and self.random_num_sample_frames_batch_sizes is not None: num_sample_frames_batch_size = np.random.choice(self.random_num_sample_frames_batch_sizes) else: num_sample_frames_batch_size = self.num_sample_frames_batch_size frame_batch_size = self._frame_count_to_latent_steps( int(num_sample_frames_batch_size), "num_sample_frames_batch_size", ) temporal_slices = x.split_by_temporal_batches(frame_batch_size, adjust_temporal=True) processed_slices = [] for x_slice in temporal_slices: if x_slice.coords.shape[0] > 0: if self.freeze_encoder: with torch.no_grad(): enc_slice, _ = self.encoder(x_slice) else: enc_slice, _ = self.encoder(x_slice) processed_slices.append(enc_slice) else: processed_slices.append(x_slice) enc_full = reconstruct_from_temporal_slices(processed_slices, target_coords=x.coords, use_cached_offsets=True) # Only compute image features if needed (e.g., for image-text alignment tasks) # Skip for reconstruction-only tasks to save computation image_feat = None if compute_image_feat: image_feat = self.encoder.head(enc_full).feats if normalize: image_feat = F.normalize(image_feat, dim=-1) enc_proj = self.proj(enc_full) return enc_proj, image_feat, enc_full @apply_forward_hook def encode( self, x: SparseTensor | torch.Tensor, normalize: bool = False, return_dict: bool = False, compute_image_feat: bool = True, ) -> tuple[SparseTensor, torch.Tensor | None, SparseTensor]: """Encode input into latent representation. Args: x: Input tensor or SparseTensor. normalize: Whether to normalize image features. return_dict: Unused, kept for compatibility. compute_image_feat: Whether to compute image features via encoder.head. Set to False for reconstruction-only tasks to save computation. Returns: Tuple of (latent, image_features or None, encoder_output). """ del return_dict if self.use_slicing and isinstance(x, torch.Tensor) and x.shape[0] > 1: raise ValueError("Legacy tensor slicing not implemented yet") if isinstance(x, torch.Tensor): x = batch_tensor_to_sparse(x, self.patch_size) x, image_feat, x_no_proj = self._encode(x, normalize=normalize, compute_image_feat=compute_image_feat) return x, image_feat, x_no_proj def _decode( self, z: SparseTensor, return_dict: bool = True, training: bool = True, discrete_decoder: bool = False, ) -> DecoderOutput | SparseTensor: """Internal decode method with temporal batching or causal-mask decoding. Args: z: Latent SparseTensor. return_dict: Whether to return DecoderOutput. training: Whether in training mode. discrete_decoder: Whether to use discrete decoder. Returns: Decoded output. """ decoder_temporal_mode = "causal_mask" if self.decoder_temporal_mode == "causal" else self.decoder_temporal_mode if decoder_temporal_mode == "causal_mask" and training: if not self._logged_decoder_temporal_plan: logging.info( "Decoder temporal plan: mode=causal_mask, full-clip masked training, inference uses KV cache." ) self._logged_decoder_temporal_plan = True decoder_module = self.discrete_decoder if discrete_decoder else self.decoder dec, _ = decoder_module( z, kv_cache=None, kv_cache_size=None, kv_cache_detach=True, temporal_causal_mask=True, ) if not return_dict: return (dec,) return DecoderOutput(sample=dec) frame_batch_size, frame_batch_strides, kv_cache_size, kv_cache_detach = self._get_decode_temporal_plan( z=z, training=training, ) decoder_kv_cache_size: int | None = None if kv_cache_size == 0 else kv_cache_size kv_cache = None temporal_slices = z.split_by_temporal_batches( frame_batch_size, frame_batch_strides, adjust_temporal=True, offset=kv_cache_size, ) processed_slices = [] for z_slice in temporal_slices: if z_slice.coords.shape[0] > 0: if discrete_decoder: dec_slice, updated_kv_cache = self.discrete_decoder( z_slice, kv_cache if decoder_kv_cache_size is not None else None, decoder_kv_cache_size, kv_cache_detach=kv_cache_detach, temporal_causal_mask=False, ) else: dec_slice, updated_kv_cache = self.decoder( z_slice, kv_cache if decoder_kv_cache_size is not None else None, decoder_kv_cache_size, kv_cache_detach=kv_cache_detach, temporal_causal_mask=False, ) if decoder_kv_cache_size is not None: kv_cache = updated_kv_cache processed_slices.append(dec_slice) else: processed_slices.append(z_slice) if not training and frame_batch_size > frame_batch_strides: processed_slices = _crop_temporal_slices_to_ownership( processed_slices, frame_batch_size=frame_batch_size, frame_batch_strides=frame_batch_strides, ) dec = reconstruct_from_temporal_slices(processed_slices, target_coords=z.coords, use_cached_offsets=True) if not return_dict: return (dec,) return DecoderOutput(sample=dec) def _get_decode_temporal_plan( self, z: SparseTensor, training: bool, ) -> tuple[int, int, int, bool]: """Resolve temporal decode scheduling. Returns: Tuple of (query_latent_steps, stride_latent_steps, cache_latent_steps, detach_cache). """ decoder_temporal_mode = "causal_mask" if self.decoder_temporal_mode == "causal" else self.decoder_temporal_mode if decoder_temporal_mode == "causal_mask": query_latent_steps = self.decoder_temporal_query_latent_steps if query_latent_steps != 1: raise ValueError( "Causal-mask decoder currently requires decoder_temporal_query_latent_steps=1; " f"got {query_latent_steps}. Multi-step causal windows are not implemented yet." ) auto_cache_resolution = self.decoder_temporal_cache_latent_steps is None if self.decoder_temporal_cache_latent_steps is None: min_t, max_t = z.get_temporal_range() cache_latent_steps = 0 if min_t is None or max_t is None else max_t - min_t + 1 else: cache_latent_steps = self.decoder_temporal_cache_latent_steps if cache_latent_steps < 0: raise ValueError( f"decoder_temporal_cache_latent_steps must be non-negative or None, got {cache_latent_steps}" ) if not self._logged_decoder_temporal_plan: logging.info( "Decoder temporal plan: mode=causal_mask, " f"query_latent_steps={query_latent_steps}, " f"cache_latent_steps={cache_latent_steps}" f"{' (auto)' if auto_cache_resolution else ''}, " f"detach_cache={self.decoder_temporal_detach_cache}" ) self._logged_decoder_temporal_plan = True return ( query_latent_steps, query_latent_steps, cache_latent_steps, self.decoder_temporal_detach_cache, ) if training and self.random_num_sample_frames_batch_sizes is not None: num_sample_frames_batch_size = np.random.choice(self.random_num_sample_frames_batch_sizes) else: num_sample_frames_batch_size = self.num_sample_frames_batch_size if training: frame_batch_size = self._frame_count_to_latent_steps( int(num_sample_frames_batch_size), "num_sample_frames_batch_size", ) frame_batch_strides = frame_batch_size kv_cache_size = 0 else: frame_batch_size = self._frame_count_to_latent_steps( self.inference_num_sample_frames_batch_size, "inference_num_sample_frames_batch_size", ) frame_batch_strides = self._frame_count_to_latent_steps( self.inference_num_sample_frames_stride, "inference_num_sample_frames_stride", ) kv_cache_size = self._frame_count_to_latent_steps( self.inference_kv_cache_size, "inference_kv_cache_size", allow_zero=True, ) if frame_batch_size < frame_batch_strides: raise ValueError( "Non-causal inference requires inference_num_sample_frames_batch_size >= " "inference_num_sample_frames_stride; " f"got batch={self.inference_num_sample_frames_batch_size}, " f"stride={self.inference_num_sample_frames_stride}." ) return frame_batch_size, frame_batch_strides, kv_cache_size, True @apply_forward_hook def decode( self, z: SparseTensor | torch.Tensor, return_dict: bool = True, return_batched_tensor: bool = False, training: bool = True, discrete_decoder: bool = False, ) -> DecoderOutput | SparseTensor: """Decode latent representation. Args: z: Latent tensor or SparseTensor. return_dict: Whether to return DecoderOutput. return_batched_tensor: Whether to return as batched tensor. training: Whether in training mode. discrete_decoder: Whether to use discrete decoder. Returns: Decoded output. """ if self.use_slicing and isinstance(z, torch.Tensor) and z.shape[0] > 1: raise ValueError("Legacy tensor slicing not implemented yet") else: decoded = self._decode(z, training=training, discrete_decoder=discrete_decoder).sample if return_batched_tensor and isinstance(decoded, SparseTensor): patch_volume = int(np.prod(self.patch_size)) channels = self.out_channels // patch_volume if self.out_channels % patch_volume == 0 else 3 decoded_batched = sparse_to_batched_tensor(decoded, self.patch_size, channels=channels) if decoded_batched is not None: decoded = rearrange(decoded_batched, "b t c h w -> b t h w c") if decoded.shape[1] == 1: decoded = decoded.squeeze(1) else: decoded_list = sparse_to_img_list(decoded, self.patch_size) if len(set(x.shape for x in decoded_list)) > 1: logging.warning(f"Decoded shapes are not the same: {[x.shape for x in decoded_list]}") decoded = decoded_list else: decoded = stack_with_bounded_inputs(decoded_list, dim=0) # [B,T,C,H,W] decoded = rearrange(decoded, "b t c h w -> b t h w c") if decoded.shape[1] == 1: decoded = decoded.squeeze(1) if not return_dict: return (decoded,) return DecoderOutput(sample=decoded)