Spaces:
Running on Zero
Running on Zero
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: OpenMDW-1.1 | |
| """Dense wrapper runtime for frozen tokenizer encode/decode.""" | |
| from __future__ import annotations | |
| import math | |
| from collections import OrderedDict | |
| from dataclasses import dataclass | |
| from typing import Literal | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from einops import rearrange | |
| from cosmos_framework.model.tokenizer.models.dense_backends import ( | |
| DenseResolvedBackend, | |
| DenseRuntimeBackend, | |
| resolve_dense_backend, | |
| run_batched_block_stack, | |
| run_varlen_block_stack, | |
| ) | |
| from cosmos_framework.model.tokenizer.models.modules.transformer.blocks import LearnedPositionEmbedder | |
| from cosmos_framework.model.tokenizer.models.sparse_autoencoder import AutoencoderKL, SparseTransformerBase | |
| from cosmos_framework.model.tokenizer.utils.tensors import cat_with_bounded_inputs | |
| class DenseTemporalChunkSpec: | |
| """Temporal chunk configuration for the dense runtime.""" | |
| raw_frames: int | |
| patch_frames: int | |
| class DenseGridMetadata: | |
| """Precomputed dense-grid metadata shared across chunk executions.""" | |
| batch_size: int | |
| temporal_patches: int | |
| height_patches: int | |
| width_patches: int | |
| learned_pe: torch.Tensor | None | |
| rope_freqs_cis: torch.Tensor | None | |
| cu_seqlens: torch.Tensor | |
| q_seqlen: list[int] | |
| max_seq_len: int | |
| DenseGridMetadataKey = tuple[str, int, int, int, int, str, str, int] | |
| DenseImageTemporalPadding = Literal["repeat", "zero"] | |
| class DenseAutoencoderRuntime(nn.Module): | |
| """Dense frozen-runtime wrapper around an existing sparse autoencoder. | |
| The wrapper intentionally holds the original ``AutoencoderKL`` as a single | |
| registered submodule and only exposes compile-friendly dense orchestration. | |
| Backend math is added incrementally in follow-up changes. | |
| """ | |
| autoencoder: AutoencoderKL | |
| backend: DenseRuntimeBackend | |
| image_temporal_padding: DenseImageTemporalPadding | |
| metadata_cache_max_entries: int | |
| _metadata_cache: OrderedDict[DenseGridMetadataKey, DenseGridMetadata] | |
| def __init__( | |
| self, | |
| autoencoder: AutoencoderKL, | |
| backend: DenseRuntimeBackend = "auto", | |
| pad_frames: int = 0, | |
| pixel_trim: bool = True, | |
| chunk_size: int = 16, | |
| image_temporal_padding: DenseImageTemporalPadding = "zero", | |
| metadata_cache_max_entries: int = 32, | |
| ) -> None: | |
| """Initialize the dense runtime wrapper. | |
| Args: | |
| autoencoder: The sparse autoencoder to wrap. | |
| backend: Backend selection for block-stack execution. | |
| pad_frames: Number of boundary frames to replicate at each end of | |
| every temporal chunk before encoding. Must be divisible by | |
| ``patch_size[0]``. Set ``0`` to disable boundary padding; | |
| set ``>0`` (typically one temporal patch, e.g. ``4``) to give | |
| the non-causal encoder additional context across chunk edges, | |
| eliminating the per-chunk-boundary PSNR dip. | |
| pixel_trim: When ``True`` and ``pad_frames > 0``, boundary latents | |
| are kept in the encoded output and trimmed in pixel space after | |
| decoding. When ``False``, boundary latents are trimmed | |
| immediately after encoding. ``True`` should always be used | |
| for the best reconstruction quality. | |
| chunk_size: Number of *raw* frames consumed by the encoder per | |
| temporal chunk. Forwarded to | |
| ``autoencoder.num_sample_frames_batch_size`` and used to | |
| slice the input video into encode batches. Must satisfy | |
| ``2 * pad_frames < chunk_size``. Default ``16``. | |
| image_temporal_padding: How a one-frame image fills a temporal | |
| patch when ``patch_size[0] > 1``. ``"zero"`` matches the | |
| canonical sparse path and keeps the first decoded frame. | |
| ``"repeat"`` preserves the legacy deployed dense-runtime | |
| contract and keeps the last decoded frame. | |
| metadata_cache_max_entries: Maximum number of device-resident | |
| dense-grid metadata entries retained by the runtime. Set to | |
| ``0`` to disable caching. | |
| """ | |
| super().__init__() | |
| self._validate_autoencoder(autoencoder) | |
| self.autoencoder = autoencoder | |
| self.backend = backend | |
| autoencoder.num_sample_frames_batch_size = chunk_size | |
| if pad_frames < 0: | |
| raise ValueError(f"pad_frames must be non-negative, got {pad_frames}.") | |
| if 2 * pad_frames >= chunk_size: | |
| raise ValueError(f"pad_frames must be less than chunk_size / 2, got {pad_frames=}, {chunk_size=}.") | |
| if image_temporal_padding not in {"repeat", "zero"}: | |
| raise ValueError( | |
| f"image_temporal_padding must be either 'repeat' or 'zero', got {image_temporal_padding!r}." | |
| ) | |
| if metadata_cache_max_entries < 0: | |
| raise ValueError(f"metadata_cache_max_entries must be non-negative, got {metadata_cache_max_entries}.") | |
| self.pad_frames = pad_frames | |
| self.pixel_trim = pixel_trim | |
| self.image_temporal_padding = image_temporal_padding | |
| self.metadata_cache_max_entries = metadata_cache_max_entries | |
| self._metadata_cache: OrderedDict[DenseGridMetadataKey, DenseGridMetadata] = OrderedDict() | |
| self.cg_compiled = False | |
| def from_autoencoder( | |
| cls, | |
| autoencoder: AutoencoderKL, | |
| backend: DenseRuntimeBackend = "auto", | |
| pad_frames: int = 0, | |
| pixel_trim: bool = True, | |
| chunk_size: int = 16, | |
| image_temporal_padding: DenseImageTemporalPadding = "zero", | |
| metadata_cache_max_entries: int = 32, | |
| ) -> "DenseAutoencoderRuntime": | |
| """Build a dense runtime from a supported sparse autoencoder.""" | |
| return cls( | |
| autoencoder=autoencoder, | |
| backend=backend, | |
| pad_frames=pad_frames, | |
| pixel_trim=pixel_trim, | |
| chunk_size=chunk_size, | |
| image_temporal_padding=image_temporal_padding, | |
| metadata_cache_max_entries=metadata_cache_max_entries, | |
| ) | |
| def _validate_autoencoder(autoencoder: AutoencoderKL) -> None: | |
| """Validate that the sparse autoencoder fits the dense-runtime V1 scope.""" | |
| if not hasattr(autoencoder, "decoder"): | |
| raise ValueError("Dense runtime V1 requires use_decoder=True.") | |
| encoder = autoencoder.encoder | |
| decoder = autoencoder.decoder | |
| if encoder.concat_latent is not None: | |
| raise ValueError("Dense runtime V1 does not support concat_latent.") | |
| if autoencoder.use_dual_latent: | |
| raise ValueError("Dense runtime V1 does not support dual latent.") | |
| if autoencoder.use_quantizer: | |
| raise ValueError("Dense runtime V1 does not support quantized latent paths.") | |
| if autoencoder.decoder_temporal_mode != "bidirectional": | |
| raise ValueError( | |
| "Dense runtime V1 only supports decoder_temporal_mode='bidirectional', " | |
| f"got {autoencoder.decoder_temporal_mode!r}." | |
| ) | |
| if int(autoencoder.inference_kv_cache_size) != 0: | |
| raise ValueError( | |
| "Dense runtime V1 does not support decoder KV cache; " | |
| f"got inference_kv_cache_size={autoencoder.inference_kv_cache_size}." | |
| ) | |
| if decoder.multiscale is not None or decoder.multiscale_outputs is not None: | |
| raise ValueError("Dense runtime V1 does not support decoder multiscale outputs.") | |
| if any(getattr(block, "multiscale", None) is not None for block in encoder.blocks): | |
| raise ValueError("Dense runtime V1 does not support encoder multiscale blocks.") | |
| if any(getattr(block, "multiscale", None) is not None for block in decoder.blocks): | |
| raise ValueError("Dense runtime V1 does not support decoder multiscale blocks.") | |
| if encoder.pe_mode not in {"joint", "learned"}: | |
| raise ValueError(f"Dense runtime V1 currently requires encoder learned/joint PE, got {encoder.pe_mode}.") | |
| if decoder.pe_mode not in {"joint", "learned"}: | |
| raise ValueError(f"Dense runtime V1 currently requires decoder learned/joint PE, got {decoder.pe_mode}.") | |
| def patch_size(self) -> tuple[int, int, int]: | |
| """Return the tokenizer patch size.""" | |
| patch_size = self.autoencoder.patch_size | |
| return int(patch_size[0]), int(patch_size[1]), int(patch_size[2]) | |
| def patch_volume(self) -> int: | |
| """Return the tokenizer patch volume.""" | |
| return math.prod(self.patch_size) | |
| def encoder_chunk_spec(self) -> DenseTemporalChunkSpec: | |
| """Return the fixed encoder chunk configuration used in eval mode.""" | |
| raw_frames = int(self.autoencoder.num_sample_frames_batch_size) | |
| patch_frames = self._raw_frames_to_patch_frames(raw_frames) | |
| return DenseTemporalChunkSpec(raw_frames=raw_frames, patch_frames=patch_frames) | |
| def decoder_window_spec(self) -> DenseTemporalChunkSpec: | |
| """Return the decoder inference window configuration.""" | |
| raw_frames = int(self.autoencoder.inference_num_sample_frames_batch_size) | |
| patch_frames = self._raw_frames_to_patch_frames(raw_frames) | |
| return DenseTemporalChunkSpec(raw_frames=raw_frames, patch_frames=patch_frames) | |
| def decoder_stride_spec(self) -> DenseTemporalChunkSpec: | |
| """Return the decoder inference stride configuration.""" | |
| raw_frames = int(self.autoencoder.inference_num_sample_frames_stride) | |
| patch_frames = self._raw_frames_to_patch_frames(raw_frames) | |
| return DenseTemporalChunkSpec(raw_frames=raw_frames, patch_frames=patch_frames) | |
| def decoder_cache_spec(self) -> DenseTemporalChunkSpec: | |
| """Return the decoder inference cache configuration.""" | |
| raw_frames = int(self.autoencoder.inference_kv_cache_size) | |
| patch_frames = self._raw_frames_to_patch_frames(raw_frames) | |
| return DenseTemporalChunkSpec(raw_frames=raw_frames, patch_frames=patch_frames) | |
| def resolve_backend(self, use_compile: bool = False) -> DenseResolvedBackend: | |
| """Resolve the backend for the current execution mode.""" | |
| return resolve_dense_backend(self.backend, use_compile=use_compile) | |
| def clear_metadata_cache(self) -> None: | |
| """Drop cached dense-grid metadata after an eval-time state mutation.""" | |
| self._metadata_cache.clear() | |
| def encode( | |
| self, | |
| dense_video: torch.Tensor, | |
| sample_posterior: bool = False, | |
| pad_to: int | None = None, | |
| chunk_raw_frames: int | None = None, | |
| encode_chunk_batch_size: int = 1, | |
| ) -> torch.Tensor: | |
| """Encode a dense video tensor into latent moments or posterior samples.""" | |
| moments = self.encode_moments( | |
| dense_video, | |
| chunk_raw_frames=chunk_raw_frames, | |
| pad_to=pad_to, | |
| encode_chunk_batch_size=encode_chunk_batch_size, | |
| ) | |
| if not sample_posterior: | |
| return moments | |
| return self._sample_dense_posterior(moments) | |
| def encode_moments( | |
| self, | |
| video: torch.Tensor, | |
| chunk_raw_frames: int | None = None, | |
| pad_to: int | None = None, | |
| encode_chunk_batch_size: int = 1, | |
| ) -> torch.Tensor: | |
| """Encode a dense video tensor into `[B, T_p, H_p, W_p, 2C]` latent moments. | |
| Args: | |
| video: Dense channels-last video tensor ``[B, T, H, W, 3]``. | |
| chunk_raw_frames: Number of raw frames per encoder chunk. Defaults | |
| to ``self.encoder_chunk_spec.raw_frames``. | |
| pad_to: Sequence-length padding target for the ``batched_with_padding`` | |
| backend (reduces CUDA graph recapture). | |
| encode_chunk_batch_size: Number of full temporal chunks to encode | |
| together. Only supported for the ``batched`` backend; defaults | |
| to ``1`` (sequential encoding). | |
| Shapes (example): | |
| Config: ``patch_size = (1, 16, 16)``, ``chunk_size = 16``, | |
| ``pad_frames = 1`` (1 raw frame replicated on each chunk edge). | |
| Whole-video input: ``[B=1, T=28, H=480, W=832, 3]``. | |
| Per-chunk pipeline (loop slices the 28 frames into 2 chunks of | |
| ``chunk_raw_frames = 16 - 2*1 = 14``): | |
| :: | |
| step shape notes | |
| --------------------------------------------------------------------------------------- | |
| 1. raw chunk [1, 14, 480, 832, 3] 1 of 2 chunks | |
| 2. after input padding [1, 16, 480, 832, 3] 1 pre + 14 raw + 1 post | |
| 3. after encoding (latent) [1, 16, 30, 52, 2C] T_p=16/1, H_p=480/16, W_p=832/16 | |
| 4. after decoding [1, 16, 480, 832, 3] | |
| 5. after pixel trim [1, 14, 480, 832, 3] drops pad_frames=1 pixel frame | |
| on each end | |
| Across both chunks the concatenated pixel-space output is | |
| ``[1, 28, 480, 832, 3]``; the latent fed to a downstream DiT is | |
| ``[1, 32, 30, 52, 2C]``. | |
| For images (``T = 1``) the input is padded to one temporal patch | |
| (``T = patch_time``) according to ``image_temporal_padding`` and | |
| ``latents_per_boundary = 0``, so the DiT-facing shape is | |
| ``[B, 1, H_p, W_p, 2C]``. | |
| """ | |
| if video.ndim != 5: | |
| raise ValueError(f"Dense runtime expects 5D video tensor, got {video.ndim}D") | |
| if video.shape[4] != 3: | |
| raise ValueError(f"Dense runtime expects video tensor with 3 channels, got {video.shape[4]}") | |
| batch_size, raw_frames, height, width, _ = video.shape | |
| patch_time, patch_height, patch_width = self.patch_size | |
| assert batch_size == 1 or encode_chunk_batch_size == 1, ( | |
| "Dense runtime with batching currently only supports batch size 1" | |
| ) | |
| if chunk_raw_frames is None: | |
| chunk_raw_frames = self.encoder_chunk_spec.raw_frames | |
| chunk_raw_frames = chunk_raw_frames - 2 * self.pad_frames | |
| assert chunk_raw_frames > 0, ( | |
| f"Padding frames must be less than chunk_raw_frames, got {chunk_raw_frames=}, {self.pad_frames=}." | |
| ) | |
| if chunk_raw_frames <= 0: | |
| raise ValueError(f"chunk_raw_frames must be positive, got {chunk_raw_frames}.") | |
| if encode_chunk_batch_size < 1: | |
| raise ValueError(f"encode_chunk_batch_size must be positive, got {encode_chunk_batch_size}.") | |
| if encode_chunk_batch_size > 1 and self.backend != "batched": | |
| raise ValueError( | |
| f"encode_chunk_batch_size > 1 is only supported for the batched backend, got backend={self.backend!r}." | |
| ) | |
| # if input is an image, we pad to form single temporal patch | |
| if raw_frames == 1: | |
| is_image = True | |
| if self.image_temporal_padding == "repeat": | |
| video = video.repeat(1, patch_time, 1, 1, 1) # [B,Pt,H,W,3] | |
| else: | |
| temporal_padding = video.new_zeros((video.shape[0], patch_time - 1, *video.shape[2:])) # [B,Pt-1,H,W,3] | |
| video = torch.cat((video, temporal_padding), dim=1) # [B,Pt,H,W,3] | |
| raw_frames = patch_time | |
| else: | |
| is_image = False | |
| if (chunk_raw_frames + 2 * self.pad_frames) % patch_time != 0: | |
| raise ValueError( | |
| f"chunk_raw_frames + 2 * pad_frames must be divisible by patch_size[0]={patch_time}, got {chunk_raw_frames=}, {self.pad_frames=}." | |
| ) | |
| if not is_image: | |
| # Noncausal scheme: first frame is its own chunk; remaining must fill complete regular chunks. | |
| remaining_frames = raw_frames - 1 | |
| remainder = remaining_frames % chunk_raw_frames | |
| if remainder != 0 and (remainder + 2 * self.pad_frames) % patch_time != 0: | |
| raise ValueError( | |
| f"Dense runtime requires (frame_count - 1) equal to " | |
| f"chunk_raw_frames * N + patch_time - 2 * pad_frames, " | |
| f"got {raw_frames=}, {chunk_raw_frames=}, {self.pad_frames=}, {patch_time=}." | |
| ) | |
| if height % patch_height != 0 or width % patch_width != 0: | |
| raise ValueError( | |
| "Dense runtime requires spatial dimensions divisible by patch size " | |
| f"{(patch_height, patch_width)}, got {(height, width)}." | |
| ) | |
| pad_frames = self.pad_frames | |
| if not is_image: | |
| latents_per_boundary = pad_frames // patch_time | |
| else: | |
| latents_per_boundary = 0 | |
| del batch_size | |
| # preserve the chunk size to reduce number of captured cuda graphs | |
| if self.backend == "batched_with_padding" and pad_to is None and self.cg_compiled: | |
| width_patches = width // patch_width | |
| height_patches = height // patch_height | |
| padded_chunk_frames = chunk_raw_frames + 2 * pad_frames | |
| temporal_patches = padded_chunk_frames // patch_time | |
| pad_to = width_patches * height_patches * temporal_patches | |
| use_chunk_batching = self.backend == "batched" and encode_chunk_batch_size > 1 | |
| def _pad_video_chunk(video_chunk: torch.Tensor) -> torch.Tensor: | |
| if pad_frames > 0 and not is_image: | |
| # UniAE chunk-wise encoding suffers a PSNR dip at chunk boundaries | |
| # because the non-causal encoder lacks context beyond the chunk edges. | |
| # Padding each chunk with pad_frames replicated boundary frames on both | |
| # sides gives the encoder that context, eliminating the boundary dip. | |
| # In practice pad_frames=4 (one temporal patch) is used. | |
| # The corresponding boundary latents are trimmed after decoding | |
| # (see pixel_trim / latents_per_boundary below). | |
| pre = video_chunk[:, 0:1].expand(-1, pad_frames, -1, -1, -1) # [B,pad,H,W,3] | |
| post = video_chunk[:, -1:].expand(-1, pad_frames, -1, -1, -1) # [B,pad,H,W,3] | |
| video_chunk = torch.cat([pre, video_chunk, post], dim=1) # [B,t+2*pad,H,W,3] | |
| return video_chunk | |
| def _trim_boundary_latents(encoded_chunk: torch.Tensor) -> torch.Tensor: | |
| if latents_per_boundary > 0 and not self.pixel_trim: | |
| t_latent = encoded_chunk.shape[1] | |
| encoded_chunk = encoded_chunk[:, latents_per_boundary : t_latent - latents_per_boundary] | |
| return encoded_chunk | |
| def _encode_padded_chunks(padded_chunks: list[torch.Tensor]) -> list[torch.Tensor]: | |
| if len(padded_chunks) == 1: | |
| encoded = self._encode_video_chunk(padded_chunks[0], pad_to=pad_to) | |
| return [_trim_boundary_latents(encoded)] | |
| batched_video = cat_with_bounded_inputs(padded_chunks, dim=0) # [B*G,t_pad,H,W,3] | |
| encoded = self._encode_video_chunk(batched_video, pad_to=pad_to) # [B*G,T_lat,Hp,Wp,2C] | |
| per_video_batch = padded_chunks[0].shape[0] | |
| return list(_trim_boundary_latents(encoded).split(per_video_batch, dim=0)) | |
| encoded_chunks: list[torch.Tensor] = [] | |
| if not is_image: | |
| # Noncausal first chunk: encode frame 0 alone, padded to patch_time copies | |
| # at the head so the encoder sees exactly patch_time frames → 1 latent L₁. | |
| # pad_to=None: this chunk has 1 temporal patch, not the regular chunk shape. | |
| first_frame = video[:, 0:1] | |
| first_chunk = first_frame.expand(-1, patch_time, -1, -1, -1).contiguous() | |
| encoded_chunks.append(self._encode_video_chunk(first_chunk, pad_to=None)) | |
| chunk_specs = [ | |
| ( | |
| start_frame, | |
| end_frame := min(start_frame + chunk_raw_frames, raw_frames), | |
| end_frame - start_frame == chunk_raw_frames, | |
| ) | |
| for start_frame in range(0 if is_image else 1, raw_frames, chunk_raw_frames) | |
| ] | |
| pending_full_chunks: list[torch.Tensor] = [] | |
| for start_frame, end_frame, is_full_chunk in chunk_specs: | |
| padded_chunk = _pad_video_chunk(video[:, start_frame:end_frame]) # [B,t_pad,H,W,3] | |
| if not use_chunk_batching or not is_full_chunk: | |
| if pending_full_chunks: | |
| encoded_chunks.extend(_encode_padded_chunks(pending_full_chunks)) | |
| pending_full_chunks = [] | |
| encoded_chunks.extend(_encode_padded_chunks([padded_chunk])) | |
| continue | |
| pending_full_chunks.append(padded_chunk) | |
| if len(pending_full_chunks) == encode_chunk_batch_size: | |
| encoded_chunks.extend(_encode_padded_chunks(pending_full_chunks)) | |
| pending_full_chunks = [] | |
| if pending_full_chunks: | |
| encoded_chunks.extend(_encode_padded_chunks(pending_full_chunks)) | |
| return cat_with_bounded_inputs(encoded_chunks, dim=1) | |
| def decode( | |
| self, | |
| dense_latent: torch.Tensor, | |
| chunk_raw_frames: int | None = None, | |
| ) -> torch.Tensor: | |
| """Decode a dense latent grid into a dense channels-last video tensor. | |
| When ``pixel_trim`` is enabled and ``pad_frames > 0``, the latent | |
| contains boundary tokens from encoding. After decoding, the | |
| corresponding boundary pixel frames are trimmed from each chunk. | |
| **Output shape contract**: | |
| - Video (``temporal_patches > 1``): ``[B, T, H, W, C]`` where T is the | |
| total number of decoded pixel frames across all chunks (after trim). | |
| - Image (``temporal_patches == 1``): ``[B, 1, H, W, C]``. The image | |
| latent decodes to ``patch_time`` frames. Sparse-compatible zero | |
| padding keeps the first frame; legacy repeat padding keeps the last. | |
| This differs from pre-``dense_runtime`` behavior where the full | |
| ``[B, patch_time, H, W, C]`` was returned. | |
| """ | |
| if self.decoder_cache_spec.patch_frames != 0: | |
| raise NotImplementedError("Dense runtime decoder V1 does not support KV cache.") | |
| latent = self._canonicalize_dense_latent(dense_latent) | |
| temporal_patches = latent.shape[1] | |
| if chunk_raw_frames is None: | |
| chunk_patch_frames = self.decoder_window_spec.patch_frames | |
| else: | |
| if chunk_raw_frames <= 0: | |
| raise ValueError(f"chunk_raw_frames must be positive, got {chunk_raw_frames}.") | |
| if chunk_raw_frames % self.patch_size[0] != 0: | |
| raise ValueError( | |
| f"chunk_raw_frames must be divisible by patch_size[0]={self.patch_size[0]}, got {chunk_raw_frames}." | |
| ) | |
| chunk_patch_frames = chunk_raw_frames // self.patch_size[0] | |
| pad_frames = self.pad_frames | |
| trim_pixel = self.pixel_trim and pad_frames > 0 | |
| # Images were encoded as a single latent (no noncausal first chunk). | |
| # Videos have temporal_patches > 1: latent[0] is the noncausal first frame. | |
| is_image = temporal_patches == 1 | |
| # Patch 0 is always a single-latent chunk — either the noncausal first | |
| # frame (video) or the sole image latent. The first video frame was | |
| # repeated across one temporal patch, so keep the last decoded frame. | |
| # Images keep either the last frame for the deployed repeat contract or | |
| # the first frame for sparse-compatible zero padding. For images | |
| # temporal_patches == 1, so the loop below is empty. | |
| decoded_chunks: list[torch.Tensor] = [] | |
| decoded_first = self._decode_latent_chunk(latent[:, 0:1]) # [B, patch_time, H, W, C] | |
| if is_image and self.image_temporal_padding == "zero": | |
| decoded_chunks.append(decoded_first[:, :1]) # [B,1,H,W,C] | |
| else: | |
| decoded_chunks.append(decoded_first[:, -1:]) # [B,1,H,W,C] | |
| for start_patch in range(1, temporal_patches, chunk_patch_frames): | |
| end_patch = min(start_patch + chunk_patch_frames, temporal_patches) | |
| latent_chunk = latent[:, start_patch:end_patch] | |
| decoded_chunk = self._decode_latent_chunk(latent_chunk) | |
| # Images have no boundary padding, so pixel trim only applies to video chunks. | |
| if trim_pixel and not is_image: | |
| decoded_chunk = decoded_chunk[:, pad_frames:-pad_frames] | |
| decoded_chunks.append(decoded_chunk) | |
| return cat_with_bounded_inputs(decoded_chunks, dim=1) | |
| def _metadata_cache_key( | |
| self, | |
| module_name: str, | |
| module: SparseTransformerBase, | |
| batch_size: int, | |
| temporal_patches: int, | |
| height_patches: int, | |
| width_patches: int, | |
| device: torch.device, | |
| dtype: torch.dtype, | |
| ) -> DenseGridMetadataKey: | |
| """Build a stable metadata-cache key for one dense grid shape.""" | |
| if isinstance(module.pos_embedder, LearnedPositionEmbedder): | |
| position_embedding_pointer = module.pos_embedder.position_embedding.weight.data_ptr() | |
| else: | |
| position_embedding_pointer = -1 | |
| return ( | |
| module_name, | |
| int(batch_size), | |
| int(temporal_patches), | |
| int(height_patches), | |
| int(width_patches), | |
| str(device), | |
| str(dtype), | |
| position_embedding_pointer, | |
| ) | |
| def train(self, mode: bool = True) -> "DenseAutoencoderRuntime": | |
| """Set module mode and discard metadata tied to prior parameter state.""" | |
| super().train(mode) | |
| self._metadata_cache.clear() | |
| return self | |
| def _raw_frames_to_patch_frames(self, raw_frames: int) -> int: | |
| """Convert raw video frames into temporal patch steps.""" | |
| patch_time = self.patch_size[0] | |
| if raw_frames % patch_time != 0: | |
| raise ValueError( | |
| f"Dense runtime requires raw frame counts divisible by patch_size[0]={patch_time}, got {raw_frames}." | |
| ) | |
| return raw_frames // patch_time | |
| def _canonicalize_dense_latent(self, dense_latent: torch.Tensor) -> torch.Tensor: | |
| """Normalize latent tensors to channels-last `[B, T_p, H_p, W_p, C]` format.""" | |
| expected_channels = self.autoencoder.latent_channels | |
| patch_time = self.patch_size[0] | |
| if dense_latent.ndim == 5: | |
| channels_last_match = dense_latent.shape[-1] == expected_channels | |
| channels_first_match = dense_latent.shape[1] == expected_channels | |
| if channels_last_match and channels_first_match: | |
| raise ValueError( | |
| "Dense runtime cannot infer 5D latent layout when both the channel-last and " | |
| "channel-first dimensions match the expected channel count " | |
| f"{expected_channels}; got shape {tuple(dense_latent.shape)}." | |
| ) | |
| if channels_last_match: | |
| latent = dense_latent | |
| elif channels_first_match: | |
| latent = rearrange(dense_latent, "b c t h w -> b t h w c") | |
| else: | |
| raise ValueError( | |
| "Dense runtime expects 5D latents in `[B, T, H, W, C]` or `[B, C, T, H, W]` format, " | |
| f"got shape {tuple(dense_latent.shape)} with expected channels={expected_channels}." | |
| ) | |
| elif dense_latent.ndim == 4: | |
| if patch_time != 1: | |
| raise ValueError( | |
| "Dense runtime image latents are only supported when patch_size[0] == 1, " | |
| f"got patch_size[0]={patch_time}." | |
| ) | |
| channels_last_match = dense_latent.shape[-1] == expected_channels | |
| channels_first_match = dense_latent.shape[1] == expected_channels | |
| if channels_last_match and channels_first_match: | |
| raise ValueError( | |
| "Dense runtime cannot infer 4D latent layout when both the channel-last and " | |
| "channel-first dimensions match the expected channel count " | |
| f"{expected_channels}; got shape {tuple(dense_latent.shape)}." | |
| ) | |
| if channels_last_match: | |
| latent = dense_latent.unsqueeze(1) | |
| elif channels_first_match: | |
| latent = rearrange(dense_latent, "b c h w -> b 1 h w c") | |
| else: | |
| raise ValueError( | |
| "Dense runtime expects 4D latents in `[B, H, W, C]` or `[B, C, H, W]` format, " | |
| f"got shape {tuple(dense_latent.shape)} with expected channels={expected_channels}." | |
| ) | |
| else: | |
| raise ValueError( | |
| "Dense runtime expects latent inputs with 4 or 5 dimensions, " | |
| f"got tensor with shape {tuple(dense_latent.shape)}." | |
| ) | |
| return latent.contiguous() | |
| def _encode_video_chunk( | |
| self, | |
| dense_video_chunk: torch.Tensor, | |
| pad_to: int | None = None, | |
| ) -> torch.Tensor: | |
| """Encode one dense video chunk into projected latent moments.""" | |
| assert pad_to is None or self.backend == "batched_with_padding", ( | |
| "pad_to is only supported for batched_with_padding backend" | |
| ) | |
| batch_size, raw_frames, height, width, _ = dense_video_chunk.shape | |
| patch_time, patch_height, patch_width = self.patch_size | |
| temporal_patches = raw_frames // patch_time | |
| height_patches = height // patch_height | |
| width_patches = width // patch_width | |
| seq_len = temporal_patches * height_patches * width_patches | |
| patch_feats = self._patchify_dense_video(dense_video_chunk) | |
| metadata = self._get_or_build_grid_metadata( | |
| module_name="encoder", | |
| module=self.autoencoder.encoder, | |
| batch_size=batch_size, | |
| temporal_patches=temporal_patches, | |
| height_patches=height_patches, | |
| width_patches=width_patches, | |
| device=patch_feats.device, | |
| dtype=self.autoencoder.encoder.input_layer.weight.dtype, | |
| ) | |
| learned_pe = metadata.learned_pe | |
| rope_freqs_cis = metadata.rope_freqs_cis | |
| needs_padding = pad_to is not None and pad_to > seq_len | |
| if pad_to is not None and pad_to < seq_len: | |
| raise ValueError(f"pad_to ({pad_to}) must be >= sequence length ({seq_len}).") | |
| if needs_padding: | |
| if batch_size != 1: | |
| raise ValueError( | |
| f"pad_to requires batch_size=1 for correct varlen masking, got batch_size={batch_size}." | |
| ) | |
| pad_amount = pad_to - seq_len | |
| patch_feats = F.pad(patch_feats, (0, 0, 0, pad_amount)) | |
| if learned_pe is not None: | |
| learned_pe = F.pad(learned_pe, (0, 0, 0, pad_amount)) | |
| if rope_freqs_cis is not None: | |
| rope_pad = torch.zeros( | |
| pad_amount, | |
| rope_freqs_cis.shape[-1], | |
| dtype=rope_freqs_cis.dtype, | |
| device=rope_freqs_cis.device, | |
| ) | |
| rope_freqs_cis = torch.cat([rope_freqs_cis, rope_pad], dim=0) | |
| moments = self._encode_chunk_core( | |
| patch_feats, | |
| learned_pe=learned_pe, | |
| rope_freqs_cis=rope_freqs_cis, | |
| q_seqlen=metadata.q_seqlen, | |
| cu_seqlens_q=metadata.cu_seqlens, | |
| max_q_seqlen=metadata.max_seq_len if not needs_padding else pad_to, | |
| ) | |
| if needs_padding: | |
| moments = moments[:, :seq_len] | |
| if self.cg_compiled: | |
| moments = moments.clone() | |
| return moments.reshape(batch_size, temporal_patches, height_patches, width_patches, -1) | |
| def _encode_chunk_core( | |
| self, | |
| patch_feats: torch.Tensor, | |
| learned_pe: torch.Tensor | None, | |
| rope_freqs_cis: torch.Tensor | None, | |
| q_seqlen: list[int] | None = None, | |
| cu_seqlens_q: torch.Tensor | None = None, | |
| max_q_seqlen: int | None = None, | |
| ) -> torch.Tensor: | |
| """Encode one dense `[B, S, patch_dim]` chunk into projected latent moments.""" | |
| encoder = self.autoencoder.encoder | |
| input_dtype = encoder.input_layer.weight.dtype | |
| if patch_feats.dtype != input_dtype: | |
| patch_feats = patch_feats.to(input_dtype) | |
| feats = F.linear(patch_feats, encoder.input_layer.weight, encoder.input_layer.bias) | |
| if learned_pe is not None: | |
| feats = feats + learned_pe | |
| block_param = next(encoder.blocks.parameters(), None) | |
| block_dtype = block_param.dtype if block_param is not None else feats.dtype | |
| if feats.dtype != block_dtype: | |
| feats = feats.to(block_dtype) | |
| feats = self._run_block_stack( | |
| blocks=encoder.blocks, | |
| feats=feats, | |
| q_seqlen=q_seqlen, | |
| cu_seqlens_q=cu_seqlens_q, | |
| max_q_seqlen=max_q_seqlen, | |
| rope_freqs_cis=rope_freqs_cis, | |
| ) | |
| feats = encoder.post_layernorm(feats) | |
| return F.linear(feats, self.autoencoder.proj.weight, self.autoencoder.proj.bias) | |
| def _patchify_dense_video(self, dense_video: torch.Tensor) -> torch.Tensor: | |
| """Patchify a dense channels-last video chunk into `[B, S, patch_dim]`.""" | |
| batch_size, raw_frames, height, width, channels = dense_video.shape | |
| patch_time, patch_height, patch_width = self.patch_size | |
| temporal_patches = raw_frames // patch_time | |
| height_patches = height // patch_height | |
| width_patches = width // patch_width | |
| return rearrange( | |
| dense_video, | |
| "b (nt pt) (nh ph) (nw pw) c -> b (nt nh nw) (pt ph pw c)", | |
| b=batch_size, | |
| nt=temporal_patches, | |
| nh=height_patches, | |
| nw=width_patches, | |
| pt=patch_time, | |
| ph=patch_height, | |
| pw=patch_width, | |
| c=channels, | |
| ) | |
| def _decode_latent_chunk(self, dense_latent_chunk: torch.Tensor) -> torch.Tensor: | |
| """Decode one dense latent chunk into a dense channels-last video chunk.""" | |
| batch_size, temporal_patches, height_patches, width_patches, _ = dense_latent_chunk.shape | |
| feats = dense_latent_chunk.reshape(batch_size, temporal_patches * height_patches * width_patches, -1) | |
| metadata = self._get_or_build_grid_metadata( | |
| module_name="decoder", | |
| module=self.autoencoder.decoder, | |
| batch_size=batch_size, | |
| temporal_patches=temporal_patches, | |
| height_patches=height_patches, | |
| width_patches=width_patches, | |
| device=feats.device, | |
| dtype=self.autoencoder.decoder.input_layer.weight.dtype, | |
| ) | |
| patch_feats = self._decode_chunk_core( | |
| feats, | |
| learned_pe=metadata.learned_pe, | |
| rope_freqs_cis=metadata.rope_freqs_cis, | |
| q_seqlen=metadata.q_seqlen, | |
| cu_seqlens_q=metadata.cu_seqlens, | |
| max_q_seqlen=metadata.max_seq_len, | |
| ) | |
| return self._unpatchify_dense_video_chunk( | |
| patch_feats, | |
| batch_size=batch_size, | |
| temporal_patches=temporal_patches, | |
| height_patches=height_patches, | |
| width_patches=width_patches, | |
| ) | |
| def _decode_chunk_core( | |
| self, | |
| feats: torch.Tensor, | |
| learned_pe: torch.Tensor | None, | |
| rope_freqs_cis: torch.Tensor | None, | |
| q_seqlen: list[int] | None = None, | |
| cu_seqlens_q: torch.Tensor | None = None, | |
| max_q_seqlen: int | None = None, | |
| ) -> torch.Tensor: | |
| """Decode one dense `[B, S, latent_dim]` chunk into patch-space features.""" | |
| decoder = self.autoencoder.decoder | |
| input_dtype = decoder.input_layer.weight.dtype | |
| if feats.dtype != input_dtype: | |
| feats = feats.to(input_dtype) | |
| feats = F.linear(feats, decoder.input_layer.weight, decoder.input_layer.bias) | |
| if learned_pe is not None: | |
| feats = feats + learned_pe | |
| block_param = next(decoder.blocks.parameters(), None) | |
| block_dtype = block_param.dtype if block_param is not None else feats.dtype | |
| if feats.dtype != block_dtype: | |
| feats = feats.to(block_dtype) | |
| feats = self._run_block_stack( | |
| blocks=decoder.blocks, | |
| feats=feats, | |
| q_seqlen=q_seqlen, | |
| cu_seqlens_q=cu_seqlens_q, | |
| max_q_seqlen=max_q_seqlen, | |
| rope_freqs_cis=rope_freqs_cis, | |
| ) | |
| feats = decoder.out_norm(feats) | |
| return F.linear(feats, decoder.out_layer.weight, decoder.out_layer.bias) | |
| def _unpatchify_dense_video_chunk( | |
| self, | |
| patch_feats: torch.Tensor, | |
| batch_size: int, | |
| temporal_patches: int, | |
| height_patches: int, | |
| width_patches: int, | |
| ) -> torch.Tensor: | |
| """Unpatchify dense decoder outputs into channels-last video chunks.""" | |
| patch_time, patch_height, patch_width = self.patch_size | |
| patch_volume = patch_time * patch_height * patch_width | |
| if self.autoencoder.out_channels % patch_volume != 0: | |
| raise ValueError( | |
| f"Autoencoder out_channels={self.autoencoder.out_channels} is not divisible by patch volume {patch_volume}." | |
| ) | |
| output_channels = self.autoencoder.out_channels // patch_volume | |
| return rearrange( | |
| patch_feats, | |
| "b (nt nh nw) (pt ph pw c) -> b (nt pt) (nh ph) (nw pw) c", | |
| b=batch_size, | |
| nt=temporal_patches, | |
| nh=height_patches, | |
| nw=width_patches, | |
| pt=patch_time, | |
| ph=patch_height, | |
| pw=patch_width, | |
| c=output_channels, | |
| ) | |
| def _run_block_stack( | |
| self, | |
| blocks: nn.ModuleList, | |
| feats: torch.Tensor, | |
| q_seqlen: list[int] | None, | |
| cu_seqlens_q: torch.Tensor | None, | |
| max_q_seqlen: int | None, | |
| rope_freqs_cis: torch.Tensor | None, | |
| ) -> torch.Tensor: | |
| """Run one backend-specific transformer block stack over `[B, S, D]` features.""" | |
| backend = self.resolve_backend(use_compile=torch.compiler.is_compiling()) | |
| if backend == "varlen": | |
| if q_seqlen is None or cu_seqlens_q is None or max_q_seqlen is None: | |
| raise ValueError("Varlen dense backend requires q_seqlen, cu_seqlens_q, and max_q_seqlen metadata.") | |
| return run_varlen_block_stack( | |
| blocks, | |
| feats, | |
| q_seqlen=q_seqlen, | |
| cu_seqlens_q=cu_seqlens_q, | |
| max_q_seqlen=max_q_seqlen, | |
| q_freqs_cis=rope_freqs_cis, | |
| ) | |
| if backend == "batched": | |
| return run_batched_block_stack( | |
| blocks, | |
| feats, | |
| q_freqs_cis=rope_freqs_cis, | |
| ) | |
| if backend == "batched_with_padding": | |
| assert feats.shape[0] == 1, ( | |
| "batched_with_padding backend only supports batch_size=1, due to varlen kernel requirements." | |
| ) | |
| return run_batched_block_stack( | |
| blocks, | |
| feats, | |
| cu_seqlens_q=cu_seqlens_q, | |
| max_q_seqlen=max_q_seqlen, | |
| q_freqs_cis=rope_freqs_cis, | |
| ) | |
| raise ValueError(f"Unsupported dense runtime backend: {backend}") | |
| def _get_or_build_grid_metadata( | |
| self, | |
| module_name: str, | |
| module: SparseTransformerBase, | |
| batch_size: int, | |
| temporal_patches: int, | |
| height_patches: int, | |
| width_patches: int, | |
| device: torch.device, | |
| dtype: torch.dtype, | |
| ) -> DenseGridMetadata: | |
| """Fetch or create dense-grid metadata for one uniform chunk shape.""" | |
| key = self._metadata_cache_key( | |
| module_name, | |
| module, | |
| batch_size, | |
| temporal_patches, | |
| height_patches, | |
| width_patches, | |
| device, | |
| dtype, | |
| ) | |
| learned_position_requires_grad = bool( | |
| module.pe_mode in {"joint", "learned"} | |
| and isinstance(module.pos_embedder, LearnedPositionEmbedder) | |
| and module.pos_embedder.position_embedding.weight.requires_grad | |
| ) | |
| # Trainable learned positions are cacheable only during eval without | |
| # gradients. Entering either train or eval mode clears prior metadata, | |
| # while the storage pointer above detects parameter replacement. | |
| cache_enabled = self.metadata_cache_max_entries > 0 and not ( | |
| learned_position_requires_grad and (self.training or torch.is_grad_enabled()) | |
| ) | |
| if not cache_enabled: | |
| return self._build_grid_metadata( | |
| module=module, | |
| batch_size=batch_size, | |
| temporal_patches=temporal_patches, | |
| height_patches=height_patches, | |
| width_patches=width_patches, | |
| device=device, | |
| ) | |
| cached = self._metadata_cache.get(key) | |
| if cached is not None: | |
| self._metadata_cache.move_to_end(key) | |
| return cached | |
| metadata = self._build_grid_metadata( | |
| module=module, | |
| batch_size=batch_size, | |
| temporal_patches=temporal_patches, | |
| height_patches=height_patches, | |
| width_patches=width_patches, | |
| device=device, | |
| ) | |
| self._metadata_cache[key] = metadata | |
| if len(self._metadata_cache) > self.metadata_cache_max_entries: | |
| self._metadata_cache.popitem(last=False) | |
| return metadata | |
| def _build_grid_metadata( | |
| self, | |
| module: SparseTransformerBase, | |
| batch_size: int, | |
| temporal_patches: int, | |
| height_patches: int, | |
| width_patches: int, | |
| device: torch.device, | |
| ) -> DenseGridMetadata: | |
| """Precompute dense-grid metadata for one uniform chunk shape.""" | |
| seq_len = temporal_patches * height_patches * width_patches | |
| q_seqlen = [seq_len] * batch_size | |
| cu_seqlens = torch.arange( | |
| 0, | |
| (batch_size + 1) * seq_len, | |
| seq_len, | |
| dtype=torch.int32, | |
| device=device, | |
| ) | |
| learned_pe = self._build_learned_position_embeddings( | |
| module, | |
| temporal_patches=temporal_patches, | |
| height_patches=height_patches, | |
| width_patches=width_patches, | |
| device=device, | |
| ) | |
| rope_freqs_cis = self._build_rope_freqs_cis( | |
| module, | |
| batch_size=batch_size, | |
| temporal_patches=temporal_patches, | |
| height_patches=height_patches, | |
| width_patches=width_patches, | |
| device=device, | |
| ) | |
| return DenseGridMetadata( | |
| batch_size=batch_size, | |
| temporal_patches=temporal_patches, | |
| height_patches=height_patches, | |
| width_patches=width_patches, | |
| learned_pe=learned_pe, | |
| rope_freqs_cis=rope_freqs_cis, | |
| cu_seqlens=cu_seqlens, | |
| q_seqlen=q_seqlen, | |
| max_seq_len=seq_len, | |
| ) | |
| def _build_learned_position_embeddings( | |
| self, | |
| module: SparseTransformerBase, | |
| temporal_patches: int, | |
| height_patches: int, | |
| width_patches: int, | |
| device: torch.device, | |
| ) -> torch.Tensor | None: | |
| """Build broadcastable learned spatial embeddings for one uniform chunk.""" | |
| if module.pe_mode not in {"joint", "learned"}: | |
| return None | |
| if not isinstance(module.pos_embedder, LearnedPositionEmbedder): | |
| raise ValueError( | |
| "Dense runtime V1 expects LearnedPositionEmbedder for learned/joint PE, " | |
| f"got {type(module.pos_embedder).__name__}." | |
| ) | |
| pos_embedder = module.pos_embedder | |
| positional_embeddings = pos_embedder.position_embedding.weight.reshape( | |
| pos_embedder.position_embedding_size, | |
| pos_embedder.position_embedding_size, | |
| -1, | |
| ) | |
| spatial_embeddings = pos_embedder._get_interpolated_position_embedding( | |
| positional_embeddings, | |
| target_height=height_patches, | |
| target_width=width_patches, | |
| target_device=device, | |
| ).to(dtype=positional_embeddings.dtype) | |
| spatial_flat = spatial_embeddings.reshape(height_patches * width_patches, -1) | |
| temporal_flat = spatial_flat.repeat(temporal_patches, 1) | |
| return temporal_flat.unsqueeze(0) | |
| def _build_rope_freqs_cis( | |
| self, | |
| module: SparseTransformerBase, | |
| batch_size: int, | |
| temporal_patches: int, | |
| height_patches: int, | |
| width_patches: int, | |
| device: torch.device, | |
| ) -> torch.Tensor | None: | |
| """Build RoPE frequencies for one regular dense patch grid.""" | |
| blocks_with_rope = [block for block in module.blocks if getattr(block.attn, "use_rope", False)] | |
| if not blocks_with_rope: | |
| return None | |
| rope_configs = { | |
| ( | |
| block.attn.rope.head_dim, | |
| block.attn.rope.pos_cls_token, | |
| ) | |
| for block in blocks_with_rope | |
| } | |
| if len(rope_configs) != 1: | |
| raise ValueError("Dense runtime V1 requires uniform RoPE configuration across blocks.") | |
| positions = self._build_regular_patch_positions( | |
| temporal_patches=temporal_patches, | |
| height_patches=height_patches, | |
| width_patches=width_patches, | |
| device=device, | |
| ) | |
| if batch_size > 1: | |
| positions = positions.unsqueeze(0).expand(batch_size, -1, -1).reshape(batch_size * positions.shape[0], -1) | |
| return blocks_with_rope[0].attn.rope.compute_freqs_cis(positions, has_special_tokens=False) | |
| def _build_regular_patch_positions( | |
| self, | |
| temporal_patches: int, | |
| height_patches: int, | |
| width_patches: int, | |
| device: torch.device, | |
| ) -> torch.Tensor: | |
| """Build regular `[S, 4]` patch coordinates in `(t, h, w, z)` order.""" | |
| return ( | |
| torch.stack( | |
| torch.meshgrid( | |
| torch.arange(temporal_patches, device=device), | |
| torch.arange(height_patches, device=device), | |
| torch.arange(width_patches, device=device), | |
| torch.arange(1, device=device), | |
| indexing="ij", | |
| ), | |
| dim=-1, | |
| ) | |
| .reshape(-1, 4) | |
| .to(dtype=torch.int32) | |
| ) | |
| def _sample_dense_posterior(self, moments: torch.Tensor) -> torch.Tensor: | |
| """Sample the dense latent posterior from `[mean, logvar]` moments.""" | |
| original_dtype = moments.dtype | |
| mean, logvar = torch.chunk(moments.to(torch.float32), 2, dim=-1) | |
| std = torch.exp(0.5 * logvar) | |
| sample = mean + std * torch.randn_like(std) | |
| return sample.to(original_dtype) | |