chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/runtime/models/vaes/autoencoder.py

@@ -0,0 +1,602 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+from typing import Dict, Optional, Tuple, Union
+
+import torch
+from diffusers.models.attention_processor import (
+    ADDED_KV_ATTENTION_PROCESSORS,
+    CROSS_ATTENTION_PROCESSORS,
+    Attention,
+    AttentionProcessor,
+    AttnAddedKVProcessor,
+    AttnProcessor,
+    FusedAttnProcessor2_0,
+)
+from diffusers.models.autoencoders.vae import (
+    Decoder,
+    DecoderOutput,
+    DiagonalGaussianDistribution,
+    Encoder,
+)
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.vaes.flux import FluxVAEConfig
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.vaes.common import (
+    can_install_spatial_shard_parallel_decode,
+)
+from sglang.multimodal_gen.runtime.models.vaes.parallel.diffusers_spatial import (
+    enable_diffusers_decoder_spatial_parallel,
+    spatial_parallel_diffusers_decode,
+)
+
+
+class AutoencoderKL(nn.Module, LayerwiseOffloadableModuleMixin):
+    r"""
+    A VAE model with KL loss for encoding images into latents and decoding latent representations into images.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+
+    Parameters:
+        in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
+        out_channels (int,  *optional*, defaults to 3): Number of channels in the output.
+        down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
+            Tuple of downsample block types.
+        up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
+            Tuple of upsample block types.
+        block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
+            Tuple of block output channels.
+        act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
+        latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space.
+        sample_size (`int`, *optional*, defaults to `32`): Sample input size.
+        scaling_factor (`float`, *optional*, defaults to 0.18215):
+            The component-wise standard deviation of the trained latent space computed using the first batch of the
+            training set. This is used to scale the latent space to have unit variance when training the diffusion
+            model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the
+            diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1
+            / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image
+            Synthesis with Latent Diffusion Models](https://huggingface.co/papers/2112.10752) paper.
+        force_upcast (`bool`, *optional*, default to `True`):
+            If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE
+            can be fine-tuned / trained to a lower range without losing too much precision in which case `force_upcast`
+            can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
+        mid_block_add_attention (`bool`, *optional*, default to `True`):
+            If enabled, the mid_block of the Encoder and Decoder will have attention blocks. If set to false, the
+            mid_block will only have resnet blocks
+    """
+
+    layerwise_offload_dit_group_enabled = False
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"]
+    layer_names = ["encoder.down_blocks", "decoder.up_blocks"]
+
+    def __init__(
+        self,
+        config: FluxVAEConfig,
+    ):
+        super().__init__()
+        self.config = config
+        arch_config = config.arch_config
+
+        in_channels = arch_config.in_channels
+        out_channels = arch_config.out_channels
+        down_block_types = arch_config.down_block_types
+        up_block_types = arch_config.up_block_types
+        block_out_channels = arch_config.block_out_channels
+        layers_per_block = arch_config.layers_per_block
+        act_fn = arch_config.act_fn
+        latent_channels = arch_config.latent_channels
+        norm_num_groups = arch_config.norm_num_groups
+        sample_size = arch_config.sample_size
+        use_quant_conv = arch_config.use_quant_conv
+        use_post_quant_conv = arch_config.use_post_quant_conv
+        mid_block_add_attention = arch_config.mid_block_add_attention
+
+        # pass init params to Encoder
+        self.encoder = Encoder(
+            in_channels=in_channels,
+            out_channels=latent_channels,
+            down_block_types=down_block_types,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            act_fn=act_fn,
+            norm_num_groups=norm_num_groups,
+            double_z=True,
+            mid_block_add_attention=mid_block_add_attention,
+        )
+
+        # pass init params to Decoder
+        self.decoder = Decoder(
+            in_channels=latent_channels,
+            out_channels=out_channels,
+            up_block_types=up_block_types,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            norm_num_groups=norm_num_groups,
+            act_fn=act_fn,
+            mid_block_add_attention=mid_block_add_attention,
+        )
+
+        self.quant_conv = (
+            nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1)
+            if use_quant_conv
+            else None
+        )
+        self.post_quant_conv = (
+            nn.Conv2d(latent_channels, latent_channels, 1)
+            if use_post_quant_conv
+            else None
+        )
+
+        self.use_slicing = False
+        self.use_tiling = False
+        self.use_parallel_decode = config.use_parallel_decode
+        self.parallel_decode_mode = config.parallel_decode_mode
+        self._spatial_parallel_decode_enabled = False
+        self._spatial_parallel_upsample_count = 0
+        if can_install_spatial_shard_parallel_decode(self.config):
+            self._spatial_parallel_upsample_count = (
+                enable_diffusers_decoder_spatial_parallel(self.decoder)
+            )
+            self._spatial_parallel_decode_enabled = True
+
+        # only relevant if vae tiling is enabled
+        self.tile_sample_min_size = sample_size
+        sample_size = (
+            self.config.sample_size[0]
+            if isinstance(self.config.sample_size, (list, tuple))
+            else self.config.sample_size
+        )
+        self.tile_latent_min_size = int(
+            sample_size / (2 ** (len(self.config.block_out_channels) - 1))
+        )
+        self.tile_overlap_factor = 0.25
+
+    def enable_tiling(self, use_tiling: bool = True):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        self.use_tiling = use_tiling
+
+    def disable_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.enable_tiling(False)
+
+    def enable_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.use_slicing = True
+
+    def disable_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.use_slicing = False
+
+    @property
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(
+            name: str,
+            module: torch.nn.Module,
+            processors: Dict[str, AttentionProcessor],
+        ):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor()
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(
+        self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]
+    ):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
+    def set_default_attn_processor(self):
+        """
+        Disables custom attention processors and sets the default attention implementation.
+        """
+        if all(
+            proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS
+            for proc in self.attn_processors.values()
+        ):
+            processor = AttnAddedKVProcessor()
+        elif all(
+            proc.__class__ in CROSS_ATTENTION_PROCESSORS
+            for proc in self.attn_processors.values()
+        ):
+            processor = AttnProcessor()
+        else:
+            raise ValueError(
+                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+            )
+
+        self.set_attn_processor(processor)
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, height, width = x.shape
+
+        if self.use_tiling and (
+            width > self.tile_sample_min_size or height > self.tile_sample_min_size
+        ):
+            return self._tiled_encode(x)
+
+        enc = self.encoder(x)
+        if self.quant_conv is not None:
+            enc = self.quant_conv(enc)
+
+        return enc
+
+    def encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+        """
+        Encode a batch of images into latents.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+                The latent representations of the encoded images. If `return_dict` is True, a
+                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
+        """
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self._encode(x)
+
+        posterior = DiagonalGaussianDistribution(h)
+
+        if not return_dict:
+            return (posterior,)
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def _decode(
+        self, z: torch.Tensor, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        if self.use_tiling and (
+            z.shape[-1] > self.tile_latent_min_size
+            or z.shape[-2] > self.tile_latent_min_size
+        ):
+            return self.tiled_decode(z, return_dict=return_dict)
+
+        if self.post_quant_conv is not None:
+            z = self.post_quant_conv(z)
+
+        if self._spatial_parallel_decode_enabled:
+            dec = spatial_parallel_diffusers_decode(
+                self.decoder, z, self._spatial_parallel_upsample_count
+            )
+        else:
+            dec = self.decoder(z)
+
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    def decode(self, z: torch.FloatTensor) -> Union[DecoderOutput, torch.FloatTensor]:
+        """
+        Decode a batch of images.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+
+        """
+
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z).sample
+
+        return decoded
+
+    def blend_v(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[2], b.shape[2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[
+                :, :, y, :
+            ] * (y / blend_extent)
+        return b
+
+    def blend_h(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[3], b.shape[3], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[
+                :, :, :, x
+            ] * (x / blend_extent)
+        return b
+
+    def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Encode a batch of images using a tiled encoder.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
+        steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
+        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
+        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
+        output, but they should be much less noticeable.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+
+        Returns:
+            `torch.Tensor`:
+                The latent representation of the encoded videos.
+        """
+
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        # Split the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[2], overlap_size):
+            row = []
+            for j in range(0, x.shape[3], overlap_size):
+                tile = x[
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_size,
+                    j : j + self.tile_sample_min_size,
+                ]
+                tile = self.encoder(tile)
+                if self.config.use_quant_conv:
+                    tile = self.quant_conv(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        enc = torch.cat(result_rows, dim=2)
+        return enc
+
+    def tiled_encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> AutoencoderKLOutput:
+        r"""Encode a batch of images using a tiled encoder.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
+        steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
+        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
+        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
+        output, but they should be much less noticeable.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`:
+                If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain
+                `tuple` is returned.
+        """
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        # Split the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[2], overlap_size):
+            row = []
+            for j in range(0, x.shape[3], overlap_size):
+                tile = x[
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_size,
+                    j : j + self.tile_sample_min_size,
+                ]
+                tile = self.encoder(tile)
+                if self.config.use_quant_conv:
+                    tile = self.quant_conv(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        moments = torch.cat(result_rows, dim=2)
+        posterior = DiagonalGaussianDistribution(moments)
+
+        if not return_dict:
+            return (posterior,)
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def tiled_decode(
+        self, z: torch.Tensor, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        r"""
+        Decode a batch of images using a tiled decoder.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+        """
+        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_sample_min_size - blend_extent
+
+        # Split z into overlapping 64x64 tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, z.shape[2], overlap_size):
+            row = []
+            for j in range(0, z.shape[3], overlap_size):
+                tile = z[
+                    :,
+                    :,
+                    i : i + self.tile_latent_min_size,
+                    j : j + self.tile_latent_min_size,
+                ]
+                if self.config.use_post_quant_conv:
+                    tile = self.post_quant_conv(tile)
+                decoded = self.decoder(tile)
+                row.append(decoded)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        dec = torch.cat(result_rows, dim=2)
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        generator: Optional[torch.Generator] = None,
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+        """
+        x = sample
+        posterior = self.encode(x).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z).sample
+
+        return dec
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections
+    def fuse_qkv_projections(self):
+        """
+        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
+        are fused. For cross-attention modules, key and value projection matrices are fused.
+
+        > [!WARNING] > This API is 🧪 experimental.
+        """
+        self.original_attn_processors = None
+
+        for _, attn_processor in self.attn_processors.items():
+            if "Added" in str(attn_processor.__class__.__name__):
+                raise ValueError(
+                    "`fuse_qkv_projections()` is not supported for models having added KV projections."
+                )
+
+        self.original_attn_processors = self.attn_processors
+
+        for module in self.modules():
+            if isinstance(module, Attention):
+                module.fuse_projections(fuse=True)
+
+        self.set_attn_processor(FusedAttnProcessor2_0())
+
+
+EntryClass = AutoencoderKL

python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_dc.py

@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.vaes.sana import SanaVAEConfig
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_decode_parallel_rank,
+    get_decode_parallel_world_size,
+)
+from sglang.multimodal_gen.runtime.layers.parallel_conv import (
+    gather_and_trim_height,
+    split_height_for_parallel_decode,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.vaes.common import (
+    can_install_spatial_shard_parallel_decode,
+)
+from sglang.multimodal_gen.runtime.models.vaes.parallel.diffusers_spatial import (
+    enable_diffusers_decoder_spatial_parallel,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class AutoencoderDC(nn.Module, LayerwiseOffloadableModuleMixin):
+    """Deep Compression Autoencoder wrapper with 32x spatial compression."""
+
+    layerwise_offload_dit_group_enabled = False
+    layer_names = ["_inner_model.encoder.down_blocks", "_inner_model.decoder.up_blocks"]
+
+    def __init__(self, config: SanaVAEConfig = None, **kwargs):
+        super().__init__()
+        self._config = config
+        self._inner_model = None
+        self._loaded_state_dict: dict[str, torch.Tensor] = {}
+        self._spatial_parallel_decode_enabled = False
+
+    def _ensure_inner_model(self, state_dict: dict[str, torch.Tensor] | None = None):
+        if self._inner_model is not None:
+            return
+
+        from diffusers import AutoencoderDC as DiffusersAutoencoderDC
+
+        device = "cpu"
+        state_to_load = (
+            state_dict if state_dict is not None else self._loaded_state_dict
+        )
+        if state_to_load:
+            first_tensor = next(iter(state_to_load.values()))
+            device = first_tensor.device
+        hf_config = {}
+        if self._config is not None:
+            arch = self._config.arch_config
+            for key, value in vars(arch).items():
+                if key == "extra_attrs" and isinstance(value, dict):
+                    for ek, ev in value.items():
+                        hf_config[ek] = ev
+                elif not key.startswith("_") and not callable(value):
+                    hf_config[key] = value
+
+        self._inner_model = DiffusersAutoencoderDC.from_config(hf_config)
+
+        if state_to_load:
+            missing, unexpected = self._inner_model.load_state_dict(
+                state_to_load, strict=False
+            )
+            if missing:
+                logger.warning(
+                    "AutoencoderDC missing keys when loading: %d keys", len(missing)
+                )
+                if len(missing) > 10:
+                    logger.debug("First 10 missing keys: %s", list(missing)[:10])
+                else:
+                    logger.debug("Missing keys: %s", list(missing))
+            if unexpected:
+                logger.debug(
+                    "AutoencoderDC unexpected keys when loading: %d keys",
+                    len(unexpected),
+                )
+            if state_dict is None:
+                self._loaded_state_dict.clear()
+
+        self._inner_model = self._inner_model.to(device)
+        if can_install_spatial_shard_parallel_decode(self._config):
+            enable_diffusers_decoder_spatial_parallel(self._inner_model.decoder)
+            self._spatial_parallel_decode_enabled = True
+
+    @property
+    def config(self):
+        if self._inner_model is not None:
+            return self._inner_model.config
+        return self._config
+
+    @property
+    def dtype(self):
+        if self._inner_model is not None:
+            return next(self._inner_model.parameters()).dtype
+        return torch.float32
+
+    @property
+    def device(self):
+        if self._inner_model is not None:
+            return next(self._inner_model.parameters()).device
+        return torch.device("cpu")
+
+    def encode(self, x: torch.Tensor, **kwargs):
+        self._ensure_inner_model()
+        return self._inner_model.encode(x, **kwargs)
+
+    def decode(self, z: torch.Tensor, **kwargs):
+        self._ensure_inner_model()
+        z = z.to(dtype=self.dtype)
+        if not self._spatial_parallel_decode_enabled:
+            return self._inner_model.decode(z, **kwargs)
+
+        expected_height = (
+            z.shape[-2] * self._config.arch_config.spatial_compression_ratio
+        )
+        z, expected_height = split_height_for_parallel_decode(
+            z,
+            expected_height=expected_height,
+            world_size=get_decode_parallel_world_size(),
+            rank=get_decode_parallel_rank(),
+        )
+        decoded = self._inner_model.decode(z, **kwargs)
+        if isinstance(decoded, tuple):
+            sample = gather_and_trim_height(decoded[0], expected_height)
+            return (sample, *decoded[1:])
+        sample = gather_and_trim_height(decoded.sample, expected_height)
+        return decoded.__class__(sample=sample)
+
+    def forward(self, x: torch.Tensor, **kwargs):
+        self._ensure_inner_model()
+        return self._inner_model(x, **kwargs)
+
+    def load_state_dict(
+        self,
+        state_dict: dict[str, torch.Tensor],
+        strict: bool = True,
+        assign: bool = False,
+    ):
+        """Intercept load_state_dict to route weights into the inner diffusers model."""
+        self._ensure_inner_model(state_dict=state_dict)
+
+    def state_dict(self, *args, **kwargs) -> dict[str, torch.Tensor]:
+        self._ensure_inner_model()
+        return self._inner_model.state_dict(*args, **kwargs)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Buffer weights for deferred loading. The inner model is built lazily."""
+        loaded_params: set[str] = set()
+        for name, weight in weights:
+            self._loaded_state_dict[name] = weight
+            loaded_params.add(name)
+        return loaded_params
+
+    def to(self, *args, **kwargs):
+        if self._inner_model is not None:
+            self._inner_model = self._inner_model.to(*args, **kwargs)
+        return super().to(*args, **kwargs)
+
+
+EntryClass = AutoencoderDC

python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_flux2.py

@@ -0,0 +1,539 @@
+import math
+from typing import Dict, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from diffusers.models.attention_processor import (
+    ADDED_KV_ATTENTION_PROCESSORS,
+    CROSS_ATTENTION_PROCESSORS,
+    AttentionProcessor,
+    AttnAddedKVProcessor,
+    AttnProcessor,
+)
+from diffusers.models.autoencoders.vae import (
+    Decoder,
+    DecoderOutput,
+    DiagonalGaussianDistribution,
+    Encoder,
+)
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+
+from sglang.multimodal_gen.configs.models.vaes.flux import Flux2VAEConfig
+from sglang.multimodal_gen.runtime.models.vaes.common import (
+    ParallelTiledVAE,
+    can_install_spatial_shard_parallel_decode,
+)
+from sglang.multimodal_gen.runtime.models.vaes.parallel.diffusers_spatial import (
+    enable_diffusers_decoder_spatial_parallel,
+    spatial_parallel_diffusers_decode,
+)
+
+
+class AutoencoderKLFlux2(ParallelTiledVAE):
+    r"""
+    A VAE model with KL loss for encoding images into latents and decoding latent representations into images.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+
+    Parameters:
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"]
+
+    def __init__(
+        self,
+        config: Flux2VAEConfig,
+    ):
+        super().__init__(config=config)
+
+        self.config = config
+        arch_config = config.arch_config
+
+        in_channels: int = arch_config.in_channels
+        out_channels: int = arch_config.out_channels
+        down_block_types: Tuple[str, ...] = arch_config.down_block_types
+        up_block_types: Tuple[str, ...] = arch_config.up_block_types
+        block_out_channels: Tuple[int, ...] = arch_config.block_out_channels
+        decoder_block_out_channels: Optional[Tuple[int, ...]] = getattr(
+            arch_config, "decoder_block_out_channels", None
+        )
+        layers_per_block: int = arch_config.layers_per_block
+        act_fn: str = arch_config.act_fn
+        latent_channels: int = arch_config.latent_channels
+        norm_num_groups: int = arch_config.norm_num_groups
+        sample_size: int = arch_config.sample_size
+        use_quant_conv: bool = arch_config.use_quant_conv
+        use_post_quant_conv: bool = arch_config.use_post_quant_conv
+        mid_block_add_attention: bool = arch_config.mid_block_add_attention
+        batch_norm_eps: float = arch_config.batch_norm_eps
+        batch_norm_momentum: float = arch_config.batch_norm_momentum
+        patch_size: Tuple[int, int] = arch_config.patch_size
+        # pass init params to Encoder
+        self.encoder = Encoder(
+            in_channels=in_channels,
+            out_channels=latent_channels,
+            down_block_types=down_block_types,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            act_fn=act_fn,
+            norm_num_groups=norm_num_groups,
+            double_z=True,
+            mid_block_add_attention=mid_block_add_attention,
+        )
+
+        # pass init params to Decoder
+        self.decoder = Decoder(
+            in_channels=latent_channels,
+            out_channels=out_channels,
+            up_block_types=up_block_types,
+            block_out_channels=decoder_block_out_channels or block_out_channels,
+            layers_per_block=layers_per_block,
+            norm_num_groups=norm_num_groups,
+            act_fn=act_fn,
+            mid_block_add_attention=mid_block_add_attention,
+        )
+
+        self.quant_conv = (
+            nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1)
+            if use_quant_conv
+            else None
+        )
+        self.post_quant_conv = (
+            nn.Conv2d(latent_channels, latent_channels, 1)
+            if use_post_quant_conv
+            else None
+        )
+
+        self.bn = nn.BatchNorm2d(
+            math.prod(patch_size) * latent_channels,
+            eps=batch_norm_eps,
+            momentum=batch_norm_momentum,
+            affine=False,
+            track_running_stats=True,
+        )
+
+        self.use_slicing = False
+        self.use_tiling = False
+        self._spatial_parallel_decode_enabled = False
+        self._spatial_parallel_upsample_count = 0
+        if can_install_spatial_shard_parallel_decode(self.config):
+            self._spatial_parallel_upsample_count = (
+                enable_diffusers_decoder_spatial_parallel(self.decoder)
+            )
+            self._spatial_parallel_decode_enabled = True
+
+        # only relevant if vae tiling is enabled
+        self.tile_sample_min_size = self.config.sample_size
+        sample_size = (
+            self.config.sample_size[0]
+            if isinstance(self.config.sample_size, (list, tuple))
+            else self.config.sample_size
+        )
+        self.tile_latent_min_size = int(
+            sample_size / (2 ** (len(self.config.block_out_channels) - 1))
+        )
+        self.tile_overlap_factor = 0.25
+
+    @property
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(
+            name: str,
+            module: torch.nn.Module,
+            processors: Dict[str, AttentionProcessor],
+        ):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor()
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(
+        self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]
+    ):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
+    def set_default_attn_processor(self):
+        """
+        Disables custom attention processors and sets the default attention implementation.
+        """
+        if all(
+            proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS
+            for proc in self.attn_processors.values()
+        ):
+            processor = AttnAddedKVProcessor()
+        elif all(
+            proc.__class__ in CROSS_ATTENTION_PROCESSORS
+            for proc in self.attn_processors.values()
+        ):
+            processor = AttnProcessor()
+        else:
+            raise ValueError(
+                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+            )
+
+        self.set_attn_processor(processor)
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, height, width = x.shape
+
+        if self.use_tiling and (
+            width > self.tile_sample_min_size or height > self.tile_sample_min_size
+        ):
+            return self._tiled_encode(x)
+
+        enc = self.encoder(x)
+        if self.quant_conv is not None:
+            enc = self.quant_conv(enc)
+
+        return enc
+
+    def encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> Union[DiagonalGaussianDistribution]:
+        """
+        Encode a batch of images into latents.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+                The latent representations of the encoded images. If `return_dict` is True, a
+                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
+        """
+
+        if x.ndim == 5:
+            assert x.shape[2] == 1
+            x = x.squeeze(2)
+
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self._encode(x)
+
+        posterior = DiagonalGaussianDistribution(h)
+        return posterior
+
+    def _decode(
+        self, z: torch.Tensor, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        if self.use_tiling and (
+            z.shape[-1] > self.tile_latent_min_size
+            or z.shape[-2] > self.tile_latent_min_size
+        ):
+            return self.tiled_decode(z, return_dict=return_dict)
+
+        if self.post_quant_conv is not None:
+            z = self.post_quant_conv(z)
+
+        if self._spatial_parallel_decode_enabled:
+            dec = spatial_parallel_diffusers_decode(
+                self.decoder, z, self._spatial_parallel_upsample_count
+            )
+        else:
+            dec = self.decoder(z)
+
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    def decode(
+        self, z: torch.FloatTensor, return_dict: bool = True, generator=None
+    ) -> Union[DecoderOutput, torch.FloatTensor]:
+        """
+        Decode a batch of images.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+
+        """
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z).sample
+
+        return decoded
+
+    def blend_v(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[2], b.shape[2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[
+                :, :, y, :
+            ] * (y / blend_extent)
+        return b
+
+    def blend_h(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[3], b.shape[3], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[
+                :, :, :, x
+            ] * (x / blend_extent)
+        return b
+
+    def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Encode a batch of images using a tiled encoder.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
+        steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
+        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
+        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
+        output, but they should be much less noticeable.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+
+        Returns:
+            `torch.Tensor`:
+                The latent representation of the encoded videos.
+        """
+
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        # Split the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[2], overlap_size):
+            row = []
+            for j in range(0, x.shape[3], overlap_size):
+                tile = x[
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_size,
+                    j : j + self.tile_sample_min_size,
+                ]
+                tile = self.encoder(tile)
+                if self.config.use_quant_conv:
+                    tile = self.quant_conv(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        enc = torch.cat(result_rows, dim=2)
+        return enc
+
+    def tiled_encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> AutoencoderKLOutput:
+        r"""Encode a batch of images using a tiled encoder.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
+        steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
+        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
+        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
+        output, but they should be much less noticeable.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`:
+                If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain
+                `tuple` is returned.
+        """
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        # Split the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[2], overlap_size):
+            row = []
+            for j in range(0, x.shape[3], overlap_size):
+                tile = x[
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_size,
+                    j : j + self.tile_sample_min_size,
+                ]
+                tile = self.encoder(tile)
+                if self.config.use_quant_conv:
+                    tile = self.quant_conv(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        moments = torch.cat(result_rows, dim=2)
+        posterior = DiagonalGaussianDistribution(moments)
+
+        if not return_dict:
+            return (posterior,)
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def tiled_decode(
+        self, z: torch.Tensor, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        r"""
+        Decode a batch of images using a tiled decoder.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+        """
+        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_sample_min_size - blend_extent
+
+        # Split z into overlapping 64x64 tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, z.shape[2], overlap_size):
+            row = []
+            for j in range(0, z.shape[3], overlap_size):
+                tile = z[
+                    :,
+                    :,
+                    i : i + self.tile_latent_min_size,
+                    j : j + self.tile_latent_min_size,
+                ]
+                if self.config.use_post_quant_conv:
+                    tile = self.post_quant_conv(tile)
+                decoded = self.decoder(tile)
+                row.append(decoded)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        dec = torch.cat(result_rows, dim=2)
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        return_dict: bool = True,
+        generator: Optional[torch.Generator] = None,
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+        """
+        x = sample
+        posterior = self.encode(x).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z).sample
+
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+
+EntryClass = AutoencoderKLFlux2

python/sglang/multimodal_gen/runtime/models/vaes/common.py

@@ -0,0 +1,810 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+from functools import lru_cache
+from math import isqrt, prod
+from typing import Optional, cast
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+
+from sglang.multimodal_gen.configs.models import VAEConfig
+from sglang.multimodal_gen.configs.models.vaes.base import (
+    should_use_spatial_shard_parallel_decode,
+)
+from sglang.multimodal_gen.runtime.distributed import (
+    get_decode_parallel_group_coordinator,
+    get_decode_parallel_world_size,
+    get_sp_parallel_rank,
+    get_sp_world_size,
+    model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+
+
+@lru_cache(maxsize=1)
+def _cached_decode_parallel_world_size(
+    is_dist_initialized: bool, is_model_parallel_initialized: bool, group_id: int
+) -> int:
+    if not is_dist_initialized or not is_model_parallel_initialized:
+        return 1
+    return get_decode_parallel_world_size()
+
+
+def _decode_parallel_world_size() -> int:
+    is_dist_initialized = dist.is_initialized()
+    is_model_parallel_initialized = model_parallel_is_initialized()
+    if not is_dist_initialized or not is_model_parallel_initialized:
+        return _cached_decode_parallel_world_size(
+            is_dist_initialized, is_model_parallel_initialized, 0
+        )
+    return _cached_decode_parallel_world_size(
+        is_dist_initialized,
+        is_model_parallel_initialized,
+        id(get_decode_parallel_group_coordinator()),
+    )
+
+
+def has_decode_parallel_world() -> bool:
+    return _decode_parallel_world_size() > 1
+
+
+def can_install_spatial_shard_parallel_decode(config: VAEConfig | None) -> bool:
+    world_size = _decode_parallel_world_size()
+    return (
+        config is not None
+        and world_size > 1
+        and should_use_spatial_shard_parallel_decode(config, world_size=world_size)
+    )
+
+
+def should_run_spatial_shard_parallel_decode(
+    config: VAEConfig, z: torch.Tensor
+) -> bool:
+    world_size = _decode_parallel_world_size()
+    return world_size > 1 and should_use_spatial_shard_parallel_decode(
+        config, z, world_size
+    )
+
+
+class ParallelTiledVAE(ABC, nn.Module, LayerwiseOffloadableModuleMixin):
+    layerwise_offload_dit_group_enabled = False
+    layer_names = [
+        "encoder.down_blocks",
+        "decoder.up_blocks",
+        "encoder.down",
+        "decoder.up",
+    ]
+    tile_sample_min_height: int
+    tile_sample_min_width: int
+    tile_sample_min_num_frames: int
+    tile_sample_stride_height: int
+    tile_sample_stride_width: int
+    tile_sample_stride_num_frames: int
+    blend_num_frames: int
+    use_tiling: bool
+    use_temporal_tiling: bool
+    use_parallel_tiling: bool
+    use_parallel_decode: bool
+    parallel_decode_mode: str
+
+    def __init__(self, config: VAEConfig, **kwargs) -> None:
+        super().__init__()
+        self.config = config
+        self.tile_sample_min_height = config.tile_sample_min_height
+        self.tile_sample_min_width = config.tile_sample_min_width
+        self.tile_sample_min_num_frames = config.tile_sample_min_num_frames
+        self.tile_sample_stride_height = config.tile_sample_stride_height
+        self.tile_sample_stride_width = config.tile_sample_stride_width
+        self.tile_sample_stride_num_frames = config.tile_sample_stride_num_frames
+        self.blend_num_frames = config.blend_num_frames
+        self.use_tiling = config.use_tiling
+        self.use_temporal_tiling = config.use_temporal_tiling
+        self.use_parallel_tiling = config.use_parallel_tiling
+        self.use_parallel_decode = config.use_parallel_decode
+        self.parallel_decode_mode = config.parallel_decode_mode
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def temporal_compression_ratio(self) -> int:
+        return cast(int, self.config.temporal_compression_ratio)
+
+    @property
+    def spatial_compression_ratio(self) -> int:
+        return cast(int, self.config.spatial_compression_ratio)
+
+    @property
+    def scaling_factor(self) -> float | torch.Tensor:
+        return cast(float | torch.Tensor, self.config.scaling_factor)
+
+    @abstractmethod
+    def _encode(self, *args, **kwargs) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def _decode(self, *args, **kwargs) -> torch.Tensor:
+        pass
+
+    def encode(self, x: torch.Tensor) -> DiagonalGaussianDistribution:
+        batch_size, num_channels, num_frames, height, width = x.shape
+        latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1
+
+        if (
+            self.use_tiling
+            and self.use_temporal_tiling
+            and num_frames > self.tile_sample_min_num_frames
+        ):
+            latents = self.tiled_encode(x)[:, :, :latent_num_frames]
+        elif self.use_tiling and (
+            width > self.tile_sample_min_width or height > self.tile_sample_min_height
+        ):
+            latents = self.spatial_tiled_encode(x)[:, :, :latent_num_frames]
+        else:
+            latents = self._encode(x)[:, :, :latent_num_frames]
+        return DiagonalGaussianDistribution(latents)
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = z.shape
+        tile_latent_min_height = (
+            self.tile_sample_min_height // self.spatial_compression_ratio
+        )
+        tile_latent_min_width = (
+            self.tile_sample_stride_width // self.spatial_compression_ratio
+        )
+        tile_latent_min_num_frames = (
+            self.tile_sample_min_num_frames // self.temporal_compression_ratio
+        )
+        num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
+
+        if should_run_spatial_shard_parallel_decode(self.config, z):
+            return self._decode(z)[:, :, :num_sample_frames]
+
+        if (
+            self.parallel_decode_mode == "tiled"
+            and self.use_tiling
+            and self.use_parallel_tiling
+            and get_sp_world_size() > 1
+        ):
+            return self.parallel_tiled_decode(z)[:, :, :num_sample_frames]
+        if (
+            self.use_tiling
+            and self.use_temporal_tiling
+            and num_frames > tile_latent_min_num_frames
+        ):
+            return self.tiled_decode(z)[:, :, :num_sample_frames]
+
+        if self.use_tiling and (
+            width > tile_latent_min_width or height > tile_latent_min_height
+        ):
+            return self.spatial_tiled_decode(z)[:, :, :num_sample_frames]
+
+        return self._decode(z)[:, :, :num_sample_frames]
+
+    def blend_v(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (
+                1 - y / blend_extent
+            ) + b[:, :, :, y, :] * (y / blend_extent)
+        return b
+
+    def blend_h(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (
+                1 - x / blend_extent
+            ) + b[:, :, :, :, x] * (x / blend_extent)
+        return b
+
+    def blend_t(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[-3], b.shape[-3], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (
+                1 - x / blend_extent
+            ) + b[:, :, x, :, :] * (x / blend_extent)
+        return b
+
+    def spatial_tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Encode a batch of images using a tiled encoder.
+
+        Args:
+            x (`torch.Tensor`): Input batch of videos.
+
+        Returns:
+            `torch.Tensor`:
+                The latent representation of the encoded videos.
+        """
+        _, _, _, height, width = x.shape
+        # latent_height = height // self.spatial_compression_ratio
+        # latent_width = width // self.spatial_compression_ratio
+
+        tile_latent_min_height = (
+            self.tile_sample_min_height // self.spatial_compression_ratio
+        )
+        tile_latent_min_width = (
+            self.tile_sample_min_width // self.spatial_compression_ratio
+        )
+        tile_latent_stride_height = (
+            self.tile_sample_stride_height // self.spatial_compression_ratio
+        )
+        tile_latent_stride_width = (
+            self.tile_sample_stride_width // self.spatial_compression_ratio
+        )
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        # Split x into overlapping tiles and encode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, self.tile_sample_stride_height):
+            row = []
+            for j in range(0, width, self.tile_sample_stride_width):
+                tile = x[
+                    :,
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_height,
+                    j : j + self.tile_sample_min_width,
+                ]
+                tile = self._encode(tile)
+                row.append(tile)
+            rows.append(row)
+
+        return self._merge_spatial_tiles(
+            rows,
+            blend_height,
+            blend_width,
+            tile_latent_stride_height,
+            tile_latent_stride_width,
+        )
+
+    def parallel_tiled_decode(self, z: torch.FloatTensor) -> torch.FloatTensor:
+        """
+        Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs
+        """
+        world_size, rank = get_sp_world_size(), get_sp_parallel_rank()
+        _, _, T, H, W = z.shape
+
+        tile_latent_min_height = (
+            self.tile_sample_min_height // self.spatial_compression_ratio
+        )
+        tile_latent_min_width = (
+            self.tile_sample_min_width // self.spatial_compression_ratio
+        )
+        tile_latent_min_num_frames = (
+            self.tile_sample_min_num_frames // self.temporal_compression_ratio
+        )
+        tile_latent_stride_height = (
+            self.tile_sample_stride_height // self.spatial_compression_ratio
+        )
+        tile_latent_stride_width = (
+            self.tile_sample_stride_width // self.spatial_compression_ratio
+        )
+        tile_latent_stride_num_frames = (
+            self.tile_sample_stride_num_frames // self.temporal_compression_ratio
+        )
+
+        blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+        blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+        # Calculate tile dimensions
+        num_t_tiles = (
+            T + tile_latent_stride_num_frames - 1
+        ) // tile_latent_stride_num_frames
+        num_h_tiles = (H + tile_latent_stride_height - 1) // tile_latent_stride_height
+        num_w_tiles = (W + tile_latent_stride_width - 1) // tile_latent_stride_width
+        total_spatial_tiles = num_h_tiles * num_w_tiles
+        total_tiles = num_t_tiles * total_spatial_tiles
+        tiles_per_rank = (total_tiles + world_size - 1) // world_size
+        start_tile_idx = rank * tiles_per_rank
+        end_tile_idx = min((rank + 1) * tiles_per_rank, total_tiles)
+
+        local_results = []
+        local_dim_metadata = []
+        for global_idx in range(start_tile_idx, end_tile_idx):
+            t_idx = global_idx // total_spatial_tiles
+            spatial_idx = global_idx % total_spatial_tiles
+            h_idx = spatial_idx // num_w_tiles
+            w_idx = spatial_idx % num_w_tiles
+
+            t_start = t_idx * tile_latent_stride_num_frames
+            h_start = h_idx * tile_latent_stride_height
+            w_start = w_idx * tile_latent_stride_width
+
+            tile = z[
+                :,
+                :,
+                t_start : t_start + tile_latent_min_num_frames + 1,
+                h_start : h_start + tile_latent_min_height,
+                w_start : w_start + tile_latent_min_width,
+            ]
+            decoded_tile = self._decode(tile)
+            if t_start > 0:
+                decoded_tile = decoded_tile[:, :, 1:, :, :]
+            local_results.append(decoded_tile.reshape(-1))
+            local_dim_metadata.append(decoded_tile.shape)
+
+        if local_results:
+            results = torch.cat(local_results, dim=0).contiguous()
+        else:
+            results = z.new_empty((0,), dtype=z.dtype)
+        del local_results
+
+        local_size = torch.tensor(
+            [results.size(0)], device=results.device, dtype=torch.int64
+        )
+        all_sizes = [
+            torch.zeros(1, device=results.device, dtype=torch.int64)
+            for _ in range(world_size)
+        ]
+        dist.all_gather(all_sizes, local_size)
+        max_size = max(size.item() for size in all_sizes)
+
+        padded_results = torch.zeros(
+            max_size, device=results.device, dtype=results.dtype
+        )
+        padded_results[: results.size(0)] = results
+
+        gathered_dim_metadata = [None] * world_size
+        gathered_results = (
+            torch.zeros_like(padded_results)
+            .repeat(world_size, *[1] * len(padded_results.shape))
+            .contiguous()
+        )
+        dist.all_gather_into_tensor(gathered_results, padded_results)
+        dist.all_gather_object(gathered_dim_metadata, local_dim_metadata)
+        gathered_dim_metadata = cast(list[list[torch.Size]], gathered_dim_metadata)
+
+        data: list = [
+            [[[] for _ in range(num_w_tiles)] for _ in range(num_h_tiles)]
+            for _ in range(num_t_tiles)
+        ]
+        global_idx = 0
+        for i, per_rank_metadata in enumerate(gathered_dim_metadata):
+            start_shape = 0
+            for shape in per_rank_metadata:
+                mul_shape = prod(shape)
+                current_data = gathered_results[
+                    i, start_shape : start_shape + mul_shape
+                ].reshape(shape)
+                t_idx = global_idx // total_spatial_tiles
+                spatial_idx = global_idx % total_spatial_tiles
+                h_idx = spatial_idx // num_w_tiles
+                w_idx = spatial_idx % num_w_tiles
+                data[t_idx][h_idx][w_idx] = current_data
+                start_shape += mul_shape
+                global_idx += 1
+
+        result_slices = []
+        last_slice_data = None
+        for i, tem_data in enumerate(data):
+            slice_data = self._merge_spatial_tiles(
+                tem_data,
+                blend_height,
+                blend_width,
+                self.tile_sample_stride_height,
+                self.tile_sample_stride_width,
+            )
+            if i > 0:
+                slice_data = self.blend_t(
+                    last_slice_data, slice_data, self.blend_num_frames
+                )
+                result_slices.append(
+                    slice_data[:, :, : self.tile_sample_stride_num_frames, :, :]
+                )
+            else:
+                result_slices.append(
+                    slice_data[:, :, : self.tile_sample_stride_num_frames + 1, :, :]
+                )
+            last_slice_data = slice_data
+        return torch.cat(result_slices, dim=2)
+
+    def parallel_patch_decode(self, z: torch.FloatTensor) -> torch.FloatTensor:
+        world_size, rank = get_sp_world_size(), get_sp_parallel_rank()
+        if world_size <= 1:
+            return self._decode(z)
+
+        tile_latent_min_height = (
+            self.tile_sample_min_height // self.spatial_compression_ratio
+        )
+        tile_latent_min_width = (
+            self.tile_sample_min_width // self.spatial_compression_ratio
+        )
+        tile_latent_stride_height = (
+            self.tile_sample_stride_height // self.spatial_compression_ratio
+        )
+        tile_latent_stride_width = (
+            self.tile_sample_stride_width // self.spatial_compression_ratio
+        )
+        overlap_h = max(0, tile_latent_min_height - tile_latent_stride_height)
+        overlap_w = max(0, tile_latent_min_width - tile_latent_stride_width)
+        halo_h = overlap_h // 2
+        halo_w = overlap_w // 2
+
+        _, _, _, latent_h, latent_w = z.shape
+        scale = self.spatial_compression_ratio
+        out_h = latent_h * scale
+        out_w = latent_w * scale
+        root = isqrt(world_size)
+        grid_rows, grid_cols = 1, world_size
+        for rows in range(root, 0, -1):
+            if world_size % rows == 0:
+                grid_rows, grid_cols = rows, world_size // rows
+                break
+        patch_id = rank
+        patch_row = patch_id // grid_cols
+        patch_col = patch_id % grid_cols
+
+        h0 = (patch_row * latent_h) // grid_rows
+        h1 = ((patch_row + 1) * latent_h) // grid_rows
+        w0 = (patch_col * latent_w) // grid_cols
+        w1 = ((patch_col + 1) * latent_w) // grid_cols
+
+        ext_h0 = max(0, h0 - halo_h)
+        ext_h1 = min(latent_h, h1 + halo_h)
+        ext_w0 = max(0, w0 - halo_w)
+        ext_w1 = min(latent_w, w1 + halo_w)
+
+        local_patch = z[:, :, :, ext_h0:ext_h1, ext_w0:ext_w1]
+        decoded_patch = self._decode(local_patch)
+
+        crop_top = (h0 - ext_h0) * scale
+        crop_bottom = crop_top + (h1 - h0) * scale
+        crop_left = (w0 - ext_w0) * scale
+        crop_right = crop_left + (w1 - w0) * scale
+        decoded_core = decoded_patch[
+            :, :, :, crop_top:crop_bottom, crop_left:crop_right
+        ].contiguous()
+
+        local_result = decoded_core.reshape(-1)
+        local_dim_metadata = torch.tensor(
+            decoded_core.shape, device=z.device, dtype=torch.int64
+        )
+        local_position = torch.tensor(
+            [h0 * scale, h1 * scale, w0 * scale, w1 * scale],
+            device=z.device,
+            dtype=torch.int64,
+        )
+        gathered_positions = [
+            torch.empty_like(local_position) for _ in range(world_size)
+        ]
+        dist.all_gather(gathered_positions, local_position)
+
+        local_size = torch.tensor(
+            [local_result.size(0)], device=z.device, dtype=torch.int64
+        )
+        gathered_dim_metadata = [
+            torch.empty_like(local_dim_metadata) for _ in range(world_size)
+        ]
+        dist.all_gather(gathered_dim_metadata, local_dim_metadata)
+
+        all_sizes = [
+            torch.zeros(1, device=z.device, dtype=torch.int64)
+            for _ in range(world_size)
+        ]
+        dist.all_gather(all_sizes, local_size)
+        max_size = max(size.item() for size in all_sizes)
+
+        padded_results = torch.zeros(max_size, device=z.device, dtype=z.dtype)
+        padded_results[: local_result.size(0)] = local_result
+        gathered_results = torch.empty(
+            (world_size, *padded_results.shape),
+            device=padded_results.device,
+            dtype=padded_results.dtype,
+        )
+        dist.all_gather_into_tensor(gathered_results, padded_results)
+
+        dec = z.new_empty(
+            (
+                decoded_core.shape[0],
+                decoded_core.shape[1],
+                decoded_core.shape[2],
+                out_h,
+                out_w,
+            )
+        )
+        for src_rank, positions in enumerate(gathered_positions):
+            h_start, h_end, w_start, w_end = [int(x.item()) for x in positions]
+            shape = tuple(int(x.item()) for x in gathered_dim_metadata[src_rank])
+            patch = gathered_results[src_rank][: prod(shape)].reshape(shape)
+            dec[:, :, :, h_start:h_end, w_start:w_end] = patch
+        return dec
+
+    def _merge_spatial_tiles(
+        self, tiles, blend_height, blend_width, stride_height, stride_width
+    ) -> torch.Tensor:
+        """Helper function to merge spatial tiles with blending"""
+        result_rows = []
+        for i, row in enumerate(tiles):
+            result_row = []
+            for j, tile in enumerate(row):
+                if i > 0:
+                    tile = self.blend_v(tiles[i - 1][j], tile, blend_height)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_width)
+                result_row.append(tile[:, :, :, :stride_height, :stride_width])
+            result_rows.append(torch.cat(result_row, dim=-1))
+        return torch.cat(result_rows, dim=-2)
+
+    def spatial_tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
+        r"""
+        Decode a batch of images using a tiled decoder.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+
+        Returns:
+            `torch.Tensor`:
+                The decoded images.
+        """
+
+        _, _, _, height, width = z.shape
+        # sample_height = height * self.spatial_compression_ratio
+        # sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = (
+            self.tile_sample_min_height // self.spatial_compression_ratio
+        )
+        tile_latent_min_width = (
+            self.tile_sample_min_width // self.spatial_compression_ratio
+        )
+        tile_latent_stride_height = (
+            self.tile_sample_stride_height // self.spatial_compression_ratio
+        )
+        tile_latent_stride_width = (
+            self.tile_sample_stride_width // self.spatial_compression_ratio
+        )
+
+        blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+        blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+        # Split z into overlapping tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, tile_latent_stride_height):
+            row = []
+            for j in range(0, width, tile_latent_stride_width):
+                tile = z[
+                    :,
+                    :,
+                    :,
+                    i : i + tile_latent_min_height,
+                    j : j + tile_latent_min_width,
+                ]
+                decoded = self._decode(tile)
+                row.append(decoded)
+            rows.append(row)
+        return self._merge_spatial_tiles(
+            rows,
+            blend_height,
+            blend_width,
+            self.tile_sample_stride_height,
+            self.tile_sample_stride_width,
+        )
+
+    def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        _, _, num_frames, height, width = x.shape
+
+        # tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
+        tile_latent_stride_num_frames = (
+            self.tile_sample_stride_num_frames // self.temporal_compression_ratio
+        )
+
+        row = []
+        for i in range(0, num_frames, self.tile_sample_stride_num_frames):
+            tile = x[:, :, i : i + self.tile_sample_min_num_frames + 1, :, :]
+            if self.use_tiling and (
+                height > self.tile_sample_min_height
+                or width > self.tile_sample_min_width
+            ):
+                tile = self.spatial_tiled_encode(tile)
+            else:
+                tile = self._encode(tile)
+            if i > 0:
+                tile = tile[:, :, 1:, :, :]
+            row.append(tile)
+        result_row = []
+        for i, tile in enumerate(row):
+            if i > 0:
+                tile = self.blend_t(row[i - 1], tile, self.blend_num_frames)
+                result_row.append(tile[:, :, :tile_latent_stride_num_frames, :, :])
+            else:
+                result_row.append(tile[:, :, : tile_latent_stride_num_frames + 1, :, :])
+        enc = torch.cat(result_row, dim=2)
+        return enc
+
+    def tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = z.shape
+
+        tile_latent_min_height = (
+            self.tile_sample_min_height // self.spatial_compression_ratio
+        )
+        tile_latent_min_width = (
+            self.tile_sample_min_width // self.spatial_compression_ratio
+        )
+        tile_latent_min_num_frames = (
+            self.tile_sample_min_num_frames // self.temporal_compression_ratio
+        )
+        tile_latent_stride_num_frames = (
+            self.tile_sample_stride_num_frames // self.temporal_compression_ratio
+        )
+
+        row = []
+        for i in range(0, num_frames, tile_latent_stride_num_frames):
+            tile = z[:, :, i : i + tile_latent_min_num_frames + 1, :, :]
+            if self.use_tiling and (
+                tile.shape[-1] > tile_latent_min_width
+                or tile.shape[-2] > tile_latent_min_height
+            ):
+                decoded = self.spatial_tiled_decode(tile)
+            else:
+                decoded = self._decode(tile)
+            if i > 0:
+                decoded = decoded[:, :, 1:, :, :]
+            row.append(decoded)
+        result_row = []
+        for i, tile in enumerate(row):
+            if i > 0:
+                tile = self.blend_t(row[i - 1], tile, self.blend_num_frames)
+                result_row.append(
+                    tile[:, :, : self.tile_sample_stride_num_frames, :, :]
+                )
+            else:
+                result_row.append(
+                    tile[:, :, : self.tile_sample_stride_num_frames + 1, :, :]
+                )
+
+        dec = torch.cat(result_row, dim=2)
+        return dec
+
+    def enable_tiling(
+        self,
+        tile_sample_min_height: int | None = None,
+        tile_sample_min_width: int | None = None,
+        tile_sample_min_num_frames: int | None = None,
+        tile_sample_stride_height: int | None = None,
+        tile_sample_stride_width: int | None = None,
+        tile_sample_stride_num_frames: int | None = None,
+        blend_num_frames: int | None = None,
+        use_tiling: bool | None = None,
+        use_temporal_tiling: bool | None = None,
+        use_parallel_tiling: bool | None = None,
+    ) -> None:
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+
+        Args:
+            tile_sample_min_height (`int`, *optional*):
+                The minimum height required for a sample to be separated into tiles across the height dimension.
+            tile_sample_min_width (`int`, *optional*):
+                The minimum width required for a sample to be separated into tiles across the width dimension.
+            tile_sample_min_num_frames (`int`, *optional*):
+                The minimum number of frames required for a sample to be separated into tiles across the frame
+                dimension.
+            tile_sample_stride_height (`int`, *optional*):
+                The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are
+                no tiling artifacts produced across the height dimension.
+            tile_sample_stride_width (`int`, *optional*):
+                The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling
+                artifacts produced across the width dimension.
+            tile_sample_stride_num_frames (`int`, *optional*):
+                The stride between two consecutive frame tiles. This is to ensure that there are no tiling artifacts
+                produced across the frame dimension.
+        """
+        self.use_tiling = True
+        self.tile_sample_min_height = (
+            tile_sample_min_height or self.tile_sample_min_height
+        )
+        self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width
+        self.tile_sample_min_num_frames = (
+            tile_sample_min_num_frames or self.tile_sample_min_num_frames
+        )
+        self.tile_sample_stride_height = (
+            tile_sample_stride_height or self.tile_sample_stride_height
+        )
+        self.tile_sample_stride_width = (
+            tile_sample_stride_width or self.tile_sample_stride_width
+        )
+        self.tile_sample_stride_num_frames = (
+            tile_sample_stride_num_frames or self.tile_sample_stride_num_frames
+        )
+        if blend_num_frames is not None:
+            self.blend_num_frames = blend_num_frames
+        else:
+            self.blend_num_frames = (
+                self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
+            )
+        self.use_tiling = use_tiling or self.use_tiling
+        self.use_temporal_tiling = use_temporal_tiling or self.use_temporal_tiling
+        self.use_parallel_tiling = use_parallel_tiling or self.use_parallel_tiling
+
+    def disable_tiling(self) -> None:
+        r"""
+        Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.use_tiling = False
+
+
+# adapted from https://github.com/huggingface/diffusers/blob/e7ffeae0a191f710881d1fbde00cd6ff025e81f2/src/diffusers/models/autoencoders/vae.py#L691
+class DiagonalGaussianDistribution:
+
+    def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        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:
+        # make sure sample is on the same device as the parameters and has same dtype
+        sample = randn_tensor(
+            self.mean.shape,
+            generator=generator,
+            device=self.parameters.device,
+            dtype=self.parameters.dtype,
+        )
+        x = self.mean + self.std * sample
+        return x
+
+    def kl(
+        self,
+        other: Optional["DiagonalGaussianDistribution"] = None,
+        dims: tuple[int, ...] = (1, 2, 3),
+    ) -> torch.Tensor:
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        else:
+            if other is None:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
+                    dim=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=dims,
+                )
+
+    def nll(
+        self, sample: torch.Tensor, dims: tuple[int, ...] = (1, 2, 3)
+    ) -> torch.Tensor:
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(
+            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+            dim=dims,
+        )
+
+    def mode(self) -> torch.Tensor:
+        return self.mean

python/sglang/multimodal_gen/runtime/models/vaes/cosmos3_avae.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Decoder-only audio tokenizer for the Cosmos3 sound modality."""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+from torch import nn
+from torch.nn.utils import weight_norm
+
+
+class Snake1d(nn.Module):
+    def __init__(self, hidden_dim: int, logscale: bool = True) -> None:
+        super().__init__()
+        self.alpha = nn.Parameter(torch.zeros(1, hidden_dim, 1))
+        self.beta = nn.Parameter(torch.zeros(1, hidden_dim, 1))
+        self.logscale = logscale
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shape = x.shape
+        alpha = torch.exp(self.alpha) if self.logscale else self.alpha
+        beta = torch.exp(self.beta) if self.logscale else self.beta
+        x = x.reshape(shape[0], shape[1], -1)
+        x = x + (beta + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+        return x.reshape(shape)
+
+
+class OobleckResidualUnit(nn.Module):
+    def __init__(self, dim: int, dilation: int = 1) -> None:
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.snake1 = Snake1d(dim)
+        self.conv1 = weight_norm(
+            nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad)
+        )
+        self.snake2 = Snake1d(dim)
+        self.conv2 = weight_norm(nn.Conv1d(dim, dim, kernel_size=1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = self.conv1(self.snake1(x))
+        y = self.conv2(self.snake2(y))
+        pad = (x.shape[-1] - y.shape[-1]) // 2
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return x + y
+
+
+class OobleckDecoderBlock(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        stride: int,
+        output_padding: int,
+    ) -> None:
+        super().__init__()
+        self.snake1 = Snake1d(input_dim)
+        self.conv_t1 = weight_norm(
+            nn.ConvTranspose1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+                output_padding=output_padding,
+            )
+        )
+        self.res_unit1 = OobleckResidualUnit(output_dim, dilation=1)
+        self.res_unit2 = OobleckResidualUnit(output_dim, dilation=3)
+        self.res_unit3 = OobleckResidualUnit(output_dim, dilation=9)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.snake1(x)
+        x = self.conv_t1(x)
+        x = self.res_unit1(x)
+        x = self.res_unit2(x)
+        return self.res_unit3(x)
+
+
+class OobleckDecoder(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        input_channels: int,
+        audio_channels: int,
+        upsampling_ratios: list[int],
+        channel_multiples: list[int],
+    ) -> None:
+        super().__init__()
+        strides = upsampling_ratios
+        mults = [1] + list(channel_multiples)
+
+        self.conv1 = weight_norm(
+            nn.Conv1d(input_channels, channels * mults[-1], kernel_size=7, padding=3)
+        )
+
+        blocks = []
+        for i, stride in enumerate(strides):
+            blocks.append(
+                OobleckDecoderBlock(
+                    input_dim=channels * mults[len(strides) - i],
+                    output_dim=channels * mults[len(strides) - i - 1],
+                    stride=stride,
+                    output_padding=stride % 2,
+                )
+            )
+        self.block = nn.ModuleList(blocks)
+        self.snake1 = Snake1d(channels)
+        self.conv2 = weight_norm(
+            nn.Conv1d(channels, audio_channels, kernel_size=7, padding=3, bias=False)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv1(x)
+        for layer in self.block:
+            x = layer(x)
+        x = self.snake1(x)
+        return self.conv2(x)
+
+
+def _cfg(config: dict[str, Any], *keys: str, default: Any = None) -> Any:
+    for k in keys:
+        v = config.get(k)
+        if v is not None:
+            return v
+    return default
+
+
+class Cosmos3AVAEAudioTokenizer(nn.Module):
+    """Cosmos3 audio tokenizer: latents → waveform via an Oobleck decoder stack."""
+
+    def __init__(self, config: dict[str, Any]) -> None:
+        super().__init__()
+        self.sample_rate = int(
+            _cfg(config, "sampling_rate", "sample_rate", default=48000)
+        )
+        self.audio_channels = int(
+            _cfg(
+                config,
+                "dec_out_channels",
+                "audio_channels",
+                default=2 if bool(config.get("stereo", True)) else 1,
+            )
+        )
+        self.latent_channels = int(
+            _cfg(config, "vocoder_input_dim", "io_channels", "latent_ch", default=64)
+        )
+        dec_strides = [
+            int(s) for s in _cfg(config, "dec_strides", default=[2, 4, 5, 6, 8])
+        ]
+        self.hop_size = int(
+            _cfg(
+                config,
+                "hop_size",
+                default=math.prod(dec_strides) if dec_strides else 1920,
+            )
+        )
+        stride_product = math.prod(dec_strides)
+        if stride_product != self.hop_size:
+            raise ValueError(
+                "Cosmos3 AVAE dec_strides product must equal hop_size: "
+                f"product={stride_product}, hop_size={self.hop_size}."
+            )
+
+        norm = str(_cfg(config, "normalization_type", default="none"))
+        if bool(_cfg(config, "normalize_latents", default=False)) and norm == "none":
+            norm = "tanh"
+        self.normalization_type = norm
+        self.tanh_input_scale = float(_cfg(config, "tanh_input_scale", default=1.5))
+        self.tanh_output_scale = float(_cfg(config, "tanh_output_scale", default=3.5))
+        self.tanh_clamp = float(_cfg(config, "tanh_clamp", default=0.995))
+
+        self.decoder = OobleckDecoder(
+            channels=int(_cfg(config, "dec_dim", default=320)),
+            input_channels=self.latent_channels,
+            audio_channels=self.audio_channels,
+            upsampling_ratios=list(reversed(dec_strides)),
+            channel_multiples=list(
+                _cfg(config, "dec_c_mults", default=[1, 2, 4, 8, 16])
+            ),
+        )
+
+    @property
+    def temporal_compression_factor(self) -> int:
+        return self.hop_size
+
+    def get_latent_num_samples(self, num_audio_samples: int) -> int:
+        return int(num_audio_samples) // self.hop_size
+
+    def get_audio_num_samples(self, num_latent_samples: int) -> int:
+        return int(num_latent_samples) * self.hop_size
+
+    def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor:
+        if self.normalization_type == "tanh":
+            in_dtype = latent.dtype
+            x = torch.clamp(
+                latent.float() / self.tanh_output_scale,
+                -self.tanh_clamp,
+                self.tanh_clamp,
+            )
+            return (torch.atanh(x) * self.tanh_input_scale).to(in_dtype)
+        if self.normalization_type != "none":
+            raise ValueError(
+                f"Unsupported AVAE normalization_type={self.normalization_type!r}."
+            )
+        return latent
+
+    @torch.no_grad()
+    def decode(self, latent: torch.Tensor) -> torch.Tensor:
+        squeeze = latent.ndim == 2
+        if squeeze:
+            latent = latent.unsqueeze(0)
+        decoder_dtype = next(self.decoder.parameters()).dtype
+        decoder_device = next(self.decoder.parameters()).device
+        z = self._denormalize_latent(latent.to(decoder_device)).to(decoder_dtype)
+        audio = self.decoder(z).clamp(-1.0, 1.0).to(latent.dtype)
+        return audio.squeeze(0) if squeeze else audio
+
+
+EntryClass = Cosmos3AVAEAudioTokenizer

python/sglang/multimodal_gen/runtime/models/vaes/dac.py

@@ -0,0 +1,646 @@
+# Copied and adapted from: https://github.com/descriptinc/descript-audio-codec
+
+# SPDX-License-Identifier: MIT
+
+import math
+from bisect import bisect_right
+from typing import Union
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.vaes.dac import DacVAEConfig
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.vaes.common import (
+    DiagonalGaussianDistribution,
+)
+
+
+def _snake(x, alpha):
+    shape = x.shape
+    x = x.reshape(shape[0], shape[1], -1)
+    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+    x = x.reshape(shape)
+    return x
+
+
+# Scripting this brings model speed up 1.4x
+snake = torch.jit.script(_snake)
+
+
+# ROCm HIPRTC can fail to compile the scripted bf16 Snake kernel.
+def _should_use_eager_snake_on_rocm_bf16(x: torch.Tensor, alpha: torch.Tensor) -> bool:
+    return (
+        torch.version.hip is not None
+        and (x.is_cuda or alpha.is_cuda)
+        and (x.dtype == torch.bfloat16 or alpha.dtype == torch.bfloat16)
+    )
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x):
+        if _should_use_eager_snake_on_rocm_bf16(x, self.alpha):
+            return _snake(x, self.alpha)
+        return snake(x, self.alpha)
+
+
+class VectorQuantize(nn.Module):
+    """
+    Implementation of VQ similar to Karpathy's repo:
+    https://github.com/karpathy/deep-vector-quantization
+    Additionally uses following tricks from Improved VQGAN
+    (https://arxiv.org/pdf/2110.04627.pdf):
+        1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space
+            for improved codebook usage
+        2. l2-normalized codes: Converts euclidean distance to cosine similarity which
+            improves training stability
+    """
+
+    def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int):
+        super().__init__()
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+
+        self.in_proj = nn.Conv1d(input_dim, codebook_dim, kernel_size=1)
+        self.out_proj = nn.Conv1d(codebook_dim, input_dim, kernel_size=1)
+        self.codebook = nn.Embedding(codebook_size, codebook_dim)
+
+    def forward(self, z):
+        """Quantize the input tensor using a fixed codebook and return the corresponding codebook vectors.
+
+        Args:
+            z (torch.Tensor): Input tensor with shape ``[B, D, T]``.
+
+        Returns:
+            tuple: A tuple containing:
+                - z_q (torch.Tensor): Quantized continuous representation with shape ``[B, D, T]``.
+                - commitment_loss (torch.Tensor): Commitment loss scalar to train encoder to predict
+                  vectors closer to codebook entries.
+                - codebook_loss (torch.Tensor): Codebook loss scalar to update the codebook.
+                - indices (torch.Tensor): Codebook indices (quantized discrete representation) with shape ``[B, T]``.
+                - z_e (torch.Tensor): Projected latents (continuous representation before quantization) with shape ``[B, D, T]``.
+        """
+
+        # Factorized codes (ViT-VQGAN) Project input into low-dimensional space
+        z_e = self.in_proj(z)  # z_e : (B x D x T)
+        z_q, indices = self.decode_latents(z_e)
+
+        commitment_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])
+        codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])
+
+        z_q = (
+            z_e + (z_q - z_e).detach()
+        )  # noop in forward pass, straight-through gradient estimator in backward pass
+
+        z_q = self.out_proj(z_q)
+
+        return z_q, commitment_loss, codebook_loss, indices, z_e
+
+    def embed_code(self, embed_id):
+        return F.embedding(embed_id, self.codebook.weight)
+
+    def decode_code(self, embed_id):
+        return self.embed_code(embed_id).transpose(1, 2)
+
+    def decode_latents(self, latents):
+        encodings = rearrange(latents, "b d t -> (b t) d")
+        codebook = self.codebook.weight  # codebook: (N x D)
+
+        # L2 normalize encodings and codebook (ViT-VQGAN)
+        encodings = F.normalize(encodings)
+        codebook = F.normalize(codebook)
+
+        # Compute euclidean distance with codebook
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )
+        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+        z_q = self.decode_code(indices)
+        return z_q, indices
+
+
+class ResidualVectorQuantize(nn.Module):
+    """
+    Introduced in SoundStream: An end2end neural audio codec
+    https://arxiv.org/abs/2107.03312
+    """
+
+    def __init__(
+        self,
+        input_dim: int = 512,
+        n_codebooks: int = 9,
+        codebook_size: int = 1024,
+        codebook_dim: Union[int, list] = 8,
+        quantizer_dropout: float = 0.0,
+    ):
+        super().__init__()
+        if isinstance(codebook_dim, int):
+            codebook_dim = [codebook_dim for _ in range(n_codebooks)]
+
+        self.n_codebooks = n_codebooks
+        self.codebook_dim = codebook_dim
+        self.codebook_size = codebook_size
+        dim_offsets = [0]
+        for dim in self.codebook_dim:
+            dim_offsets.append(dim_offsets[-1] + dim)
+        self._codebook_dim_offsets = tuple(dim_offsets)
+
+        self.quantizers = nn.ModuleList(
+            [
+                VectorQuantize(input_dim, codebook_size, codebook_dim[i])
+                for i in range(n_codebooks)
+            ]
+        )
+        self.quantizer_dropout = quantizer_dropout
+
+    def forward(self, z, n_quantizers: int = None):
+        """Quantize the input tensor using a fixed set of codebooks and return the corresponding codebook vectors.
+
+        Args:
+            z (torch.Tensor): Input tensor with shape ``[B, D, T]``.
+            n_quantizers (int, optional): Number of quantizers to use. If ``None``,
+                all quantizers are used. When ``n_quantizers`` < ``self.n_codebooks``,
+                quantizer dropout is applied. Note: if ``self.quantizer_dropout`` > 0
+                and in training mode, this argument is ignored and a random number of
+                quantizers is used.
+
+        Returns:
+            tuple: A tuple containing:
+                - z_q (torch.Tensor): Quantized continuous representation with shape ``[B, D, T]``.
+                - codes (torch.Tensor): Codebook indices for each codebook with shape ``[B, N, T]``
+                  (quantized discrete representation of input).
+                - latents (torch.Tensor): Projected latents with shape ``[B, N*D, T]``
+                  (continuous representation before quantization).
+                - commitment_loss (torch.Tensor): Commitment loss scalar to train encoder to predict
+                  vectors closer to codebook entries.
+                - codebook_loss (torch.Tensor): Codebook loss scalar to update the codebook.
+        """
+        z_q = 0
+        residual = z
+        commitment_loss = 0
+        codebook_loss = 0
+
+        codebook_indices = []
+        latents = []
+
+        if n_quantizers is None:
+            n_quantizers = self.n_codebooks
+        quantizers = self.quantizers
+        if self.training:
+            batch_size = z.shape[0]
+            device = z.device
+            n_quantizers = torch.full(
+                (batch_size,),
+                self.n_codebooks + 1,
+                device=device,
+                dtype=torch.long,
+            )
+            if self.quantizer_dropout > 0:
+                dropout = torch.randint(
+                    1,
+                    self.n_codebooks + 1,
+                    (batch_size,),
+                    device=device,
+                )
+                n_dropout = int(batch_size * self.quantizer_dropout)
+                if n_dropout > 0:
+                    n_quantizers[:n_dropout] = dropout[:n_dropout]
+
+            for i, quantizer in enumerate(quantizers):
+                z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(
+                    residual
+                )
+
+                # Create mask to apply quantizer dropout
+                mask = i < n_quantizers
+                z_q = z_q + z_q_i * mask[:, None, None]
+                residual = residual - z_q_i
+
+                # Sum losses
+                commitment_loss += (commitment_loss_i * mask).mean()
+                codebook_loss += (codebook_loss_i * mask).mean()
+
+                codebook_indices.append(indices_i)
+                latents.append(z_e_i)
+        else:
+            for i, quantizer in enumerate(quantizers):
+                if i >= n_quantizers:
+                    break
+                z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(
+                    residual
+                )
+                z_q = z_q + z_q_i
+                residual = residual - z_q_i
+
+                commitment_loss += commitment_loss_i.mean()
+                codebook_loss += codebook_loss_i.mean()
+
+                codebook_indices.append(indices_i)
+                latents.append(z_e_i)
+
+        codes = torch.stack(codebook_indices, dim=1)
+        latents = torch.cat(latents, dim=1)
+
+        return z_q, codes, latents, commitment_loss, codebook_loss
+
+    def from_codes(self, codes: torch.Tensor):
+        """Reconstruct the continuous representation from quantized codes.
+
+        Args:
+            codes (torch.Tensor): Quantized discrete representation with shape ``[B, N, T]``.
+
+        Returns:
+            tuple: A tuple containing:
+                - z_q (torch.Tensor): Quantized continuous representation with shape ``[B, D, T]``.
+                - z_p (torch.Tensor): Concatenated latent space representation with shape ``[B, N*D, T]``.
+                - codes (torch.Tensor): Original input codebook indices with shape ``[B, N, T]``.
+        """
+        z_q = 0.0
+        z_p = []
+        n_codebooks = codes.shape[1]
+        for i in range(n_codebooks):
+            z_p_i = self.quantizers[i].decode_code(codes[:, i, :])
+            z_p.append(z_p_i)
+
+            z_q_i = self.quantizers[i].out_proj(z_p_i)
+            z_q = z_q + z_q_i
+        return z_q, torch.cat(z_p, dim=1), codes
+
+    def from_latents(self, latents: torch.Tensor):
+        """Reconstruct the continuous representation from unquantized latents.
+
+        Args:
+            latents (torch.Tensor): Continuous representation after projection with shape ``[B, N*D, T]``.
+
+        Returns:
+            tuple: A tuple containing:
+                - z_q (torch.Tensor): Quantized representation of full-projected space with shape ``[B, D, T]``.
+                - z_p (torch.Tensor): Quantized representation of latent space with shape ``[B, N*D, T]``.
+                - codes (torch.Tensor): Codebook indices with shape ``[B, N, T]``.
+        """
+        z_q = 0
+        z_p = []
+        codes = []
+        dims = self._codebook_dim_offsets
+        n_codebooks = bisect_right(dims, latents.shape[1]) - 1
+        for i in range(n_codebooks):
+            j, k = dims[i], dims[i + 1]
+            z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :])
+            z_p.append(z_p_i)
+            codes.append(codes_i)
+
+            z_q_i = self.quantizers[i].out_proj(z_p_i)
+            z_q = z_q + z_q_i
+
+        return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1)
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, dim: int = 16, dilation: int = 1):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            Snake1d(dim),
+            nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+            Snake1d(dim),
+            nn.Conv1d(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x):
+        y = self.block(x)
+        pad = (x.shape[-1] - y.shape[-1]) // 2
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return x + y
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, dim: int = 16, stride: int = 1):
+        super().__init__()
+        self.block = nn.Sequential(
+            ResidualUnit(dim // 2, dilation=1),
+            ResidualUnit(dim // 2, dilation=3),
+            ResidualUnit(dim // 2, dilation=9),
+            Snake1d(dim // 2),
+            nn.Conv1d(
+                dim // 2,
+                dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+            ),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        d_model: int = 64,
+        strides: list = [2, 4, 8, 8],
+        d_latent: int = 64,
+    ):
+        super().__init__()
+        # Create first convolution
+        self.block = [nn.Conv1d(1, d_model, kernel_size=7, padding=3)]
+
+        # Create EncoderBlocks that double channels as they downsample by `stride`
+        for stride in strides:
+            d_model *= 2
+            self.block += [EncoderBlock(d_model, stride=stride)]
+
+        # Create last convolution
+        self.block += [
+            Snake1d(d_model),
+            nn.Conv1d(d_model, d_latent, kernel_size=3, padding=1),
+        ]
+
+        # Wrap black into nn.Sequential
+        self.block = nn.Sequential(*self.block)
+        self.enc_dim = d_model
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1):
+        super().__init__()
+        self.block = nn.Sequential(
+            Snake1d(input_dim),
+            nn.ConvTranspose1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+                output_padding=stride % 2,
+            ),
+            ResidualUnit(output_dim, dilation=1),
+            ResidualUnit(output_dim, dilation=3),
+            ResidualUnit(output_dim, dilation=9),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        input_channel,
+        channels,
+        rates,
+        d_out: int = 1,
+    ):
+        super().__init__()
+
+        # Add first conv layer
+        layers = [nn.Conv1d(input_channel, channels, kernel_size=7, padding=3)]
+
+        # Add upsampling + MRF blocks
+        for i, stride in enumerate(rates):
+            input_dim = channels // 2**i
+            output_dim = channels // 2 ** (i + 1)
+            layers += [DecoderBlock(input_dim, output_dim, stride)]
+
+        # Add final conv layer
+        layers += [
+            Snake1d(output_dim),
+            nn.Conv1d(output_dim, d_out, kernel_size=7, padding=3),
+            nn.Tanh(),
+        ]
+
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class DAC(nn.Module, LayerwiseOffloadableModuleMixin):
+    layerwise_offload_dit_group_enabled = False
+    layer_names = ["encoder.block", "decoder.model"]
+
+    def __init__(
+        self,
+        config: DacVAEConfig,
+    ):
+        super().__init__()
+
+        self.continuous = config.continuous
+        self.decoder_dim = config.decoder_dim
+        self.decoder_rates = config.decoder_rates
+        self.encoder_dim = config.encoder_dim
+        self.encoder_rates = config.encoder_rates
+        self.hop_length = math.prod(config.encoder_rates)
+        self.sample_rate = config.sample_rate
+
+        if config.latent_dim is None:
+            latent_dim = config.encoder_dim * (2 ** len(config.encoder_rates))
+        else:
+            latent_dim = config.latent_dim
+
+        self.latent_dim = latent_dim
+
+        if config.load_encoder:
+            self.encoder = Encoder(config.encoder_dim, config.encoder_rates, latent_dim)
+
+        if not config.continuous:
+            self.n_codebooks = config.n_codebooks
+            self.codebook_size = config.codebook_size
+            self.codebook_dim = config.codebook_dim
+            self.quantizer = ResidualVectorQuantize(
+                input_dim=latent_dim,
+                n_codebooks=config.n_codebooks,
+                codebook_size=config.codebook_size,
+                codebook_dim=config.codebook_dim,
+                quantizer_dropout=config.quantizer_dropout,
+            )
+        else:
+            self.quant_conv = torch.nn.Conv1d(latent_dim, 2 * latent_dim, 1)
+            self.post_quant_conv = torch.nn.Conv1d(latent_dim, latent_dim, 1)
+
+        if config.load_decoder:
+            self.decoder = Decoder(
+                latent_dim,
+                config.decoder_dim,
+                config.decoder_rates,
+            )
+
+        self.apply(self.init_weights)
+
+    @staticmethod
+    def init_weights(m):
+        if isinstance(m, nn.Conv1d):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            nn.init.constant_(m.bias, 0)
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def preprocess(self, audio_data, sample_rate):
+        if sample_rate is None:
+            sample_rate = self.sample_rate
+        assert sample_rate == self.sample_rate
+
+        length = audio_data.shape[-1]
+        right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
+        audio_data = nn.functional.pad(audio_data, (0, right_pad))
+
+        return audio_data
+
+    def encode(
+        self,
+        audio_data: torch.Tensor,
+        n_quantizers: int = None,
+    ):
+        """Encode audio data into latent representations.
+
+        This method processes audio through the encoder network and optionally applies
+        vector quantization (in VQ mode) or projects to a Gaussian distribution (in
+        continuous mode) to produce latent representations.
+
+        Args:
+            audio_data (torch.Tensor): Audio data to encode, with shape ``[B, 1, T]``.
+            n_quantizers (int, optional): Number of quantizers to use. If ``None``,
+                all quantizers are used. Only applicable in VQ mode (``continuous=False``).
+
+        Returns:
+            tuple: A tuple containing:
+                - z (torch.Tensor): Encoded representation. In VQ mode, this is the
+                  quantized continuous representation with shape ``[B, D, T]``. In
+                  continuous mode, this is a ``DiagonalGaussianDistribution`` object.
+                - codes (torch.Tensor or None): Codebook indices with shape ``[B, N, T]``
+                  in VQ mode, ``None`` in continuous mode.
+                - latents (torch.Tensor or None): Projected latents with shape ``[B, N*D, T]``
+                  in VQ mode, ``None`` in continuous mode.
+                - commitment_loss (torch.Tensor): Commitment loss scalar.
+                - codebook_loss (torch.Tensor): Codebook loss scalar.
+
+        Note:
+            In continuous mode, the encoded representation is projected through a
+            quantization convolution layer and wrapped in a ``DiagonalGaussianDistribution``
+            for VAE training.
+        """
+        z = self.encoder(audio_data)  # [B x D x T]
+        if not self.continuous:
+            z, codes, latents, commitment_loss, codebook_loss = self.quantizer(
+                z, n_quantizers
+            )
+        else:
+            z = self.quant_conv(z)  # [B x 2D x T]
+            z = DiagonalGaussianDistribution(z)
+            codes, latents, commitment_loss, codebook_loss = None, None, 0, 0
+
+        return z, codes, latents, commitment_loss, codebook_loss
+
+    def decode(self, z: torch.Tensor):
+        """Decode latent representations back to audio waveforms.
+
+        This method takes latent representations (either quantized from VQ mode or sampled
+        from the posterior in continuous mode) and reconstructs the corresponding audio
+        through the decoder network.
+
+        Args:
+            z (torch.Tensor): Latent representation to decode, with shape ``[B, D, T]``.
+                In VQ mode (``continuous=False``), this is the quantized continuous
+                representation. In continuous mode (``continuous=True``), this is sampled
+                from the posterior distribution.
+
+        Returns:
+            torch.Tensor: Decoded audio data with shape ``[B, 1, T']``. The output length
+            T' is determined by the decoder's upsampling rates and may differ from the
+            input temporal dimension T.
+
+        Note:
+            In continuous mode (``continuous=True``), the input is first passed through
+            a post-quantization convolution layer before being fed to the decoder.
+        """
+        if not self.continuous:
+            audio = self.decoder(z)
+        else:
+            z = self.post_quant_conv(z)
+            audio = self.decoder(z)
+
+        return audio
+
+    def forward(
+        self,
+        audio_data: torch.Tensor,
+        sample_rate: int = None,
+        n_quantizers: int = None,
+    ):
+        """Model forward pass.
+
+        Args:
+            audio_data (torch.Tensor): Audio to encode, shape [B, 1, T].
+            sample_rate (int, optional): Sample rate in Hz. Defaults to
+                ``self.sample_rate`` when ``None``.
+            n_quantizers (int, optional): Number of quantizers to use. When ``None``,
+                all quantizers are used. Only used in VQ mode (``continuous=False``).
+
+        Returns:
+            dict: A dictionary containing different keys depending on the mode:
+
+            **VQ Mode (``continuous=False``):**
+                - "audio" (torch.Tensor): Decoded audio, shape [B, 1, length].
+                - "z" (torch.Tensor): Quantized continuous representation, shape [B, D, T].
+                - "codes" (torch.Tensor): Codebook indices, shape [B, N, T].
+                - "latents" (torch.Tensor): Projected latents, shape [B, N*D, T].
+                - "vq/commitment_loss" (torch.Tensor): Commitment loss.
+                - "vq/codebook_loss" (torch.Tensor): Codebook loss.
+
+            **Continuous Mode (``continuous=True``):**
+                - "audio" (torch.Tensor): Decoded audio, shape [B, 1, length].
+                - "z" (torch.Tensor): Latent representation, shape [B, D, T].
+                - "kl_loss" (torch.Tensor): KL divergence loss (for VAE training).
+        """
+        length = audio_data.shape[-1]
+        audio_data = self.preprocess(audio_data, sample_rate)
+        if not self.continuous:
+            z, codes, latents, commitment_loss, codebook_loss = self.encode(
+                audio_data, n_quantizers
+            )
+
+            x = self.decode(z)
+            return {
+                "audio": x[..., :length],
+                "z": z,
+                "codes": codes,
+                "latents": latents,
+                "vq/commitment_loss": commitment_loss,
+                "vq/codebook_loss": codebook_loss,
+            }
+        else:
+            posterior, _, _, _, _ = self.encode(audio_data, n_quantizers)
+            z = posterior.sample()
+            x = self.decode(z)
+
+            kl_loss = posterior.kl(dims=(1, 2))
+            kl_loss = kl_loss.mean()
+
+            return {
+                "audio": x[..., :length],
+                "z": z,
+                "kl_loss": kl_loss,
+            }
+
+
+EntryClass = DAC

python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py

@@ -0,0 +1,978 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from diffusers
+
+# Copyright 2024 The Hunyuan Team, The HuggingFace Team and The sglang-diffusion Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.jit_kernel.diffusion.group_norm_silu import apply_group_norm_silu
+from sglang.multimodal_gen.configs.models.vaes import HunyuanVAEConfig
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_decode_parallel_rank,
+    get_decode_parallel_world_size,
+)
+from sglang.multimodal_gen.runtime.layers.activation import get_act_fn
+from sglang.multimodal_gen.runtime.layers.parallel_conv import (
+    SpatialParallelConv3d,
+    chunk_height_by_sizes,
+    disable_spatial_parallel_decode,
+    gather_and_trim_height,
+    gather_variable_height,
+    split_height_for_parallel_decode,
+)
+from sglang.multimodal_gen.runtime.models.vaes.common import (
+    ParallelTiledVAE,
+    can_install_spatial_shard_parallel_decode,
+    should_run_spatial_shard_parallel_decode,
+)
+
+
+def prepare_causal_attention_mask(
+    num_frames: int,
+    height_width: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    batch_size: int | None = None,
+) -> torch.Tensor:
+    indices = torch.arange(1, num_frames + 1, dtype=torch.int32, device=device)
+    indices_blocks = indices.repeat_interleave(height_width)
+    x, y = torch.meshgrid(indices_blocks, indices_blocks, indexing="xy")
+    mask = torch.where(x <= y, 0, -float("inf")).to(dtype=dtype)
+
+    if batch_size is not None:
+        mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
+    return mask
+
+
+def _make_spatial_parallel_conv3d(
+    conv: nn.Conv3d,
+    *,
+    height_pad: int,
+    width_pad: int,
+    padding_mode: str,
+) -> SpatialParallelConv3d:
+    spatial_conv = SpatialParallelConv3d(
+        in_channels=conv.in_channels,
+        out_channels=conv.out_channels,
+        kernel_size=conv.kernel_size,
+        stride=conv.stride,
+        padding=(0, height_pad, width_pad),
+        dilation=conv.dilation,
+        groups=conv.groups,
+        bias=conv.bias is not None,
+        padding_mode=padding_mode,
+    )
+    spatial_conv.weight = conv.weight
+    spatial_conv.bias = conv.bias
+    return spatial_conv
+
+
+def _apply_group_norm_silu(
+    hidden_states: torch.Tensor,
+    norm: nn.GroupNorm,
+    nonlinearity: nn.Module,
+    spatial_parallel: bool,
+) -> torch.Tensor:
+    if not spatial_parallel:
+        return apply_group_norm_silu(hidden_states, norm, nonlinearity)
+    hidden_states, heights = gather_variable_height(hidden_states)
+    hidden_states = apply_group_norm_silu(
+        hidden_states.contiguous(), norm, nonlinearity
+    )
+    return chunk_height_by_sizes(hidden_states, heights)
+
+
+class HunyuanVAEAttention(nn.Module):
+
+    def __init__(
+        self, in_channels, heads, dim_head, eps, norm_num_groups, bias
+    ) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.heads = heads
+        self.dim_head = dim_head
+        self.eps = eps
+        self.norm_num_groups = norm_num_groups
+        self.bias = bias
+
+        inner_dim = heads * dim_head
+
+        # Define the projection layers
+        self.to_q = nn.Linear(in_channels, inner_dim, bias=bias)
+        self.to_k = nn.Linear(in_channels, inner_dim, bias=bias)
+        self.to_v = nn.Linear(in_channels, inner_dim, bias=bias)
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, in_channels, bias=bias))
+
+        # Optional normalization layers
+        self.group_norm = nn.GroupNorm(
+            norm_num_groups, in_channels, eps=eps, affine=True
+        )
+
+    def forward(
+        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        residual = hidden_states
+
+        batch_size, sequence_length, _ = hidden_states.shape
+
+        hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        # Project to query, key, value
+        query = self.to_q(hidden_states)
+        key = self.to_k(hidden_states)
+        value = self.to_v(hidden_states)
+
+        # Reshape for multi-head attention
+        head_dim = self.dim_head
+
+        query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+
+        # Perform scaled dot-product attention
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        # Reshape back
+        hidden_states = hidden_states.transpose(1, 2).reshape(
+            batch_size, -1, self.heads * head_dim
+        )
+        hidden_states = hidden_states.to(query.dtype)
+
+        # Linear projection
+        hidden_states = self.to_out(hidden_states)
+
+        # Residual connection and rescale
+        hidden_states = hidden_states + residual
+
+        return hidden_states
+
+
+class HunyuanVideoCausalConv3d(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int | tuple[int, int, int] = 3,
+        stride: int | tuple[int, int, int] = 1,
+        padding: int | tuple[int, int, int] = 0,
+        dilation: int | tuple[int, int, int] = 1,
+        bias: bool = True,
+        pad_mode: str = "replicate",
+    ) -> None:
+        super().__init__()
+
+        kernel_size = (
+            (kernel_size, kernel_size, kernel_size)
+            if isinstance(kernel_size, int)
+            else kernel_size
+        )
+
+        self.pad_mode = pad_mode
+        self.time_causal_padding = (
+            kernel_size[0] // 2,
+            kernel_size[0] // 2,
+            kernel_size[1] // 2,
+            kernel_size[1] // 2,
+            kernel_size[2] - 1,
+            0,
+        )
+        self.spatial_parallel_time_padding = (
+            0,
+            0,
+            0,
+            0,
+            kernel_size[2] - 1,
+            0,
+        )
+        self.spatial_parallel = False
+
+        self.conv = nn.Conv3d(
+            in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias
+        )
+
+    def enable_spatial_parallel(self) -> None:
+        if isinstance(self.conv, SpatialParallelConv3d):
+            self.spatial_parallel = True
+            return
+        if self.conv.kernel_size == (1, 1, 1):
+            self.spatial_parallel = True
+            return
+        self.conv = _make_spatial_parallel_conv3d(
+            self.conv,
+            height_pad=self.time_causal_padding[2],
+            width_pad=self.time_causal_padding[0],
+            padding_mode=self.pad_mode,
+        )
+        self.spatial_parallel = True
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        padding = (
+            self.spatial_parallel_time_padding
+            if self.spatial_parallel
+            else self.time_causal_padding
+        )
+        if any(padding):
+            hidden_states = F.pad(hidden_states, padding, mode=self.pad_mode)
+        return self.conv(hidden_states)
+
+
+class HunyuanVideoUpsampleCausal3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int | None = None,
+        kernel_size: int = 3,
+        stride: int = 1,
+        bias: bool = True,
+        upsample_factor: tuple[int, ...] = (2, 2, 2),
+    ) -> None:
+        super().__init__()
+
+        out_channels = out_channels or in_channels
+        self.upsample_factor = upsample_factor
+
+        self.conv = HunyuanVideoCausalConv3d(
+            in_channels, out_channels, kernel_size, stride, bias=bias
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        num_frames = hidden_states.size(2)
+
+        first_frame, other_frames = hidden_states.split((1, num_frames - 1), dim=2)
+        first_frame = F.interpolate(
+            first_frame.squeeze(2),
+            scale_factor=self.upsample_factor[1:],
+            mode="nearest",
+        ).unsqueeze(2)
+
+        if num_frames > 1:
+            # See: https://github.com/pytorch/pytorch/issues/81665
+            # Unless you have a version of pytorch where non-contiguous implementation of F.interpolate
+            # is fixed, this will raise either a runtime error, or fail silently with bad outputs.
+            # If you are encountering an error here, make sure to try running encoding/decoding with
+            # `vae.enable_tiling()` first. If that doesn't work, open an issue at:
+            # https://github.com/huggingface/diffusers/issues
+            other_frames = other_frames.contiguous()
+            other_frames = F.interpolate(
+                other_frames, scale_factor=self.upsample_factor, mode="nearest"
+            )
+            hidden_states = torch.cat((first_frame, other_frames), dim=2)
+        else:
+            hidden_states = first_frame
+
+        hidden_states = self.conv(hidden_states)
+        return hidden_states
+
+
+class HunyuanVideoDownsampleCausal3D(nn.Module):
+
+    def __init__(
+        self,
+        channels: int,
+        out_channels: int | None = None,
+        padding: int = 1,
+        kernel_size: int = 3,
+        bias: bool = True,
+        stride=2,
+    ) -> None:
+        super().__init__()
+        out_channels = out_channels or channels
+
+        self.conv = HunyuanVideoCausalConv3d(
+            channels, out_channels, kernel_size, stride, padding, bias=bias
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.conv(hidden_states)
+        return hidden_states
+
+
+class HunyuanVideoResnetBlockCausal3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int | None = None,
+        dropout: float = 0.0,
+        groups: int = 32,
+        eps: float = 1e-6,
+        non_linearity: str = "silu",
+    ) -> None:
+        super().__init__()
+        out_channels = out_channels or in_channels
+
+        self.nonlinearity = get_act_fn(non_linearity)
+
+        self.norm1 = nn.GroupNorm(groups, in_channels, eps=eps, affine=True)
+        self.conv1 = HunyuanVideoCausalConv3d(in_channels, out_channels, 3, 1, 0)
+
+        self.norm2 = nn.GroupNorm(groups, out_channels, eps=eps, affine=True)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = HunyuanVideoCausalConv3d(out_channels, out_channels, 3, 1, 0)
+
+        self.conv_shortcut = None
+        if in_channels != out_channels:
+            self.conv_shortcut = HunyuanVideoCausalConv3d(
+                in_channels, out_channels, 1, 1, 0
+            )
+        self.spatial_parallel = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states.contiguous()
+        residual = hidden_states
+
+        hidden_states = _apply_group_norm_silu(
+            hidden_states, self.norm1, self.nonlinearity, self.spatial_parallel
+        )
+        hidden_states = self.conv1(hidden_states)
+
+        hidden_states = _apply_group_norm_silu(
+            hidden_states, self.norm2, self.nonlinearity, self.spatial_parallel
+        )
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+
+        if self.conv_shortcut is not None:
+            residual = self.conv_shortcut(residual)
+
+        hidden_states = hidden_states + residual
+        return hidden_states
+
+
+class HunyuanVideoMidBlock3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_act_fn: str = "silu",
+        resnet_groups: int = 32,
+        add_attention: bool = True,
+        attention_head_dim: int = 1,
+    ) -> None:
+        super().__init__()
+        resnet_groups = (
+            resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+        )
+        self.add_attention = add_attention
+
+        # There is always at least one resnet
+        resnets = [
+            HunyuanVideoResnetBlockCausal3D(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                eps=resnet_eps,
+                groups=resnet_groups,
+                dropout=dropout,
+                non_linearity=resnet_act_fn,
+            )
+        ]
+        attentions: list[HunyuanVAEAttention | None] = []
+
+        for _ in range(num_layers):
+            if self.add_attention:
+                attentions.append(
+                    HunyuanVAEAttention(
+                        in_channels,
+                        heads=in_channels // attention_head_dim,
+                        dim_head=attention_head_dim,
+                        eps=resnet_eps,
+                        norm_num_groups=resnet_groups,
+                        bias=True,
+                    )
+                )
+            else:
+                attentions.append(None)
+
+            resnets.append(
+                HunyuanVideoResnetBlockCausal3D(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    non_linearity=resnet_act_fn,
+                )
+            )
+
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+        self.spatial_parallel = False
+        self.gradient_checkpointing = False
+
+    def _run_attention(
+        self, attn: HunyuanVAEAttention, hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        heights = None
+        if self.spatial_parallel:
+            hidden_states, heights = gather_variable_height(hidden_states)
+
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3)
+        attention_mask = prepare_causal_attention_mask(
+            num_frames,
+            height * width,
+            hidden_states.dtype,
+            hidden_states.device,
+            batch_size=batch_size,
+        )
+        hidden_states = attn(hidden_states, attention_mask=attention_mask)
+        hidden_states = hidden_states.unflatten(1, (num_frames, height, width)).permute(
+            0, 4, 1, 2, 3
+        )
+        if heights is not None:
+            hidden_states = chunk_height_by_sizes(hidden_states, heights)
+        return hidden_states
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            hidden_states = self._gradient_checkpointing_func(
+                self.resnets[0], hidden_states
+            )
+
+            for attn, resnet in zip(self.attentions, self.resnets[1:], strict=True):
+                if attn is not None:
+                    hidden_states = self._run_attention(attn, hidden_states)
+
+                hidden_states = self._gradient_checkpointing_func(resnet, hidden_states)
+
+        else:
+            hidden_states = self.resnets[0](hidden_states)
+
+            for attn, resnet in zip(self.attentions, self.resnets[1:], strict=True):
+                if attn is not None:
+                    hidden_states = self._run_attention(attn, hidden_states)
+
+                hidden_states = resnet(hidden_states)
+
+        return hidden_states
+
+
+class HunyuanVideoDownBlock3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_act_fn: str = "silu",
+        resnet_groups: int = 32,
+        add_downsample: bool = True,
+        downsample_stride: tuple[int, ...] | int = 2,
+        downsample_padding: int = 1,
+    ) -> None:
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                HunyuanVideoResnetBlockCausal3D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    non_linearity=resnet_act_fn,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    HunyuanVideoDownsampleCausal3D(
+                        out_channels,
+                        out_channels=out_channels,
+                        padding=downsample_padding,
+                        stride=downsample_stride,
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for resnet in self.resnets:
+                hidden_states = self._gradient_checkpointing_func(resnet, hidden_states)
+        else:
+            for resnet in self.resnets:
+                hidden_states = resnet(hidden_states)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+        return hidden_states
+
+
+class HunyuanVideoUpBlock3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_act_fn: str = "silu",
+        resnet_groups: int = 32,
+        add_upsample: bool = True,
+        upsample_scale_factor: tuple[int, ...] = (2, 2, 2),
+    ) -> None:
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            input_channels = in_channels if i == 0 else out_channels
+
+            resnets.append(
+                HunyuanVideoResnetBlockCausal3D(
+                    in_channels=input_channels,
+                    out_channels=out_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    non_linearity=resnet_act_fn,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList(
+                [
+                    HunyuanVideoUpsampleCausal3D(
+                        out_channels,
+                        out_channels=out_channels,
+                        upsample_factor=upsample_scale_factor,
+                    )
+                ]
+            )
+        else:
+            self.upsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for resnet in self.resnets:
+                hidden_states = self._gradient_checkpointing_func(resnet, hidden_states)
+
+        else:
+            for resnet in self.resnets:
+                hidden_states = resnet(hidden_states)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states)
+
+        return hidden_states
+
+
+class HunyuanVideoEncoder3D(nn.Module):
+    r"""
+    Causal encoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603).
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        down_block_types: tuple[str, ...] = (
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+            "HunyuanVideoDownBlock3D",
+        ),
+        block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
+        layers_per_block: int = 2,
+        norm_num_groups: int = 32,
+        act_fn: str = "silu",
+        double_z: bool = True,
+        mid_block_add_attention=True,
+        temporal_compression_ratio: int = 4,
+        spatial_compression_ratio: int = 8,
+    ) -> None:
+        super().__init__()
+
+        self.conv_in = HunyuanVideoCausalConv3d(
+            in_channels, block_out_channels[0], kernel_size=3, stride=1
+        )
+        self.mid_block: HunyuanVideoMidBlock3D | None = None
+        self.down_blocks = nn.ModuleList([])
+
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            if down_block_type != "HunyuanVideoDownBlock3D":
+                raise ValueError(f"Unsupported down_block_type: {down_block_type}")
+
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+            num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio))
+            num_time_downsample_layers = int(np.log2(temporal_compression_ratio))
+
+            if temporal_compression_ratio == 4:
+                add_spatial_downsample = bool(i < num_spatial_downsample_layers)
+                add_time_downsample = bool(
+                    i >= (len(block_out_channels) - 1 - num_time_downsample_layers)
+                    and not is_final_block
+                )
+            elif temporal_compression_ratio == 8:
+                add_spatial_downsample = bool(i < num_spatial_downsample_layers)
+                add_time_downsample = bool(i < num_time_downsample_layers)
+            else:
+                raise ValueError(
+                    f"Unsupported time_compression_ratio: {temporal_compression_ratio}"
+                )
+
+            downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1)
+            downsample_stride_T = (2,) if add_time_downsample else (1,)
+            downsample_stride = tuple(downsample_stride_T + downsample_stride_HW)
+
+            down_block = HunyuanVideoDownBlock3D(
+                num_layers=layers_per_block,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                add_downsample=bool(add_spatial_downsample or add_time_downsample),
+                resnet_eps=1e-6,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                downsample_stride=downsample_stride,
+                downsample_padding=0,
+            )
+
+            self.down_blocks.append(down_block)
+
+        self.mid_block = HunyuanVideoMidBlock3D(
+            in_channels=block_out_channels[-1],
+            resnet_eps=1e-6,
+            resnet_act_fn=act_fn,
+            attention_head_dim=block_out_channels[-1],
+            resnet_groups=norm_num_groups,
+            add_attention=mid_block_add_attention,
+        )
+
+        self.conv_norm_out = nn.GroupNorm(
+            num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6
+        )
+        self.conv_act = nn.SiLU()
+
+        conv_out_channels = 2 * out_channels if double_z else out_channels
+        self.conv_out = HunyuanVideoCausalConv3d(
+            block_out_channels[-1], conv_out_channels, kernel_size=3
+        )
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.conv_in(hidden_states)
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for down_block in self.down_blocks:
+                hidden_states = self._gradient_checkpointing_func(
+                    down_block, hidden_states
+                )
+
+            hidden_states = self._gradient_checkpointing_func(
+                self.mid_block, hidden_states
+            )
+        else:
+            for down_block in self.down_blocks:
+                hidden_states = down_block(hidden_states)
+            assert self.mid_block is not None
+            hidden_states = self.mid_block(hidden_states)
+
+        hidden_states = apply_group_norm_silu(
+            hidden_states, self.conv_norm_out, self.conv_act
+        )
+        hidden_states = self.conv_out(hidden_states)
+
+        return hidden_states
+
+
+class HunyuanVideoDecoder3D(nn.Module):
+    r"""
+    Causal decoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603).
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        up_block_types: tuple[str, ...] = (
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+            "HunyuanVideoUpBlock3D",
+        ),
+        block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
+        layers_per_block: int = 2,
+        norm_num_groups: int = 32,
+        act_fn: str = "silu",
+        mid_block_add_attention=True,
+        time_compression_ratio: int = 4,
+        spatial_compression_ratio: int = 8,
+    ):
+        super().__init__()
+        self.layers_per_block = layers_per_block
+
+        self.conv_in = HunyuanVideoCausalConv3d(
+            in_channels, block_out_channels[-1], kernel_size=3, stride=1
+        )
+        self.up_blocks = nn.ModuleList([])
+
+        # mid
+        self.mid_block = HunyuanVideoMidBlock3D(
+            in_channels=block_out_channels[-1],
+            resnet_eps=1e-6,
+            resnet_act_fn=act_fn,
+            attention_head_dim=block_out_channels[-1],
+            resnet_groups=norm_num_groups,
+            add_attention=mid_block_add_attention,
+        )
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            if up_block_type != "HunyuanVideoUpBlock3D":
+                raise ValueError(f"Unsupported up_block_type: {up_block_type}")
+
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+            num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio))
+            num_time_upsample_layers = int(np.log2(time_compression_ratio))
+
+            if time_compression_ratio == 4:
+                add_spatial_upsample = bool(i < num_spatial_upsample_layers)
+                add_time_upsample = bool(
+                    i >= len(block_out_channels) - 1 - num_time_upsample_layers
+                    and not is_final_block
+                )
+            else:
+                raise ValueError(
+                    f"Unsupported time_compression_ratio: {time_compression_ratio}"
+                )
+
+            upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1)
+            upsample_scale_factor_T = (2,) if add_time_upsample else (1,)
+            upsample_scale_factor = tuple(
+                upsample_scale_factor_T + upsample_scale_factor_HW
+            )
+
+            up_block = HunyuanVideoUpBlock3D(
+                num_layers=self.layers_per_block + 1,
+                in_channels=prev_output_channel,
+                out_channels=output_channel,
+                add_upsample=bool(add_spatial_upsample or add_time_upsample),
+                upsample_scale_factor=upsample_scale_factor,
+                resnet_eps=1e-6,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+            )
+
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        self.conv_norm_out = nn.GroupNorm(
+            num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6
+        )
+        self.conv_act = nn.SiLU()
+        self.conv_out = HunyuanVideoCausalConv3d(
+            block_out_channels[0], out_channels, kernel_size=3
+        )
+
+        self.spatial_parallel = False
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.conv_in(hidden_states)
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            hidden_states = self._gradient_checkpointing_func(
+                self.mid_block, hidden_states
+            )
+
+            for up_block in self.up_blocks:
+                hidden_states = self._gradient_checkpointing_func(
+                    up_block, hidden_states
+                )
+        else:
+            hidden_states = self.mid_block(hidden_states)
+
+            for up_block in self.up_blocks:
+                hidden_states = up_block(hidden_states)
+
+        # post-process
+        hidden_states = _apply_group_norm_silu(
+            hidden_states,
+            self.conv_norm_out,
+            self.conv_act,
+            self.spatial_parallel,
+        )
+        hidden_states = self.conv_out(hidden_states)
+
+        return hidden_states
+
+
+def _enable_hunyuan_decoder_spatial_parallel(decoder: nn.Module) -> None:
+    for module in decoder.modules():
+        if isinstance(module, HunyuanVideoCausalConv3d):
+            module.enable_spatial_parallel()
+        elif isinstance(module, HunyuanVideoResnetBlockCausal3D):
+            module.spatial_parallel = True
+        elif isinstance(module, HunyuanVideoMidBlock3D):
+            module.spatial_parallel = True
+        elif isinstance(module, HunyuanVideoDecoder3D):
+            module.spatial_parallel = True
+
+
+class AutoencoderKLHunyuanVideo(ParallelTiledVAE):
+    r"""
+    A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos.
+    Introduced in [HunyuanVideo](https://huggingface.co/papers/2412.03603).
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+    """
+
+    _supports_gradient_checkpointing = True
+
+    def __init__(
+        self,
+        config: HunyuanVAEConfig,
+    ) -> None:
+        nn.Module.__init__(self)
+        ParallelTiledVAE.__init__(self, config)
+
+        # TODO(will): only pass in config. We do this by manually defining a
+        # config for hunyuan vae
+        self.block_out_channels = config.block_out_channels
+
+        if config.load_encoder:
+            self.encoder = HunyuanVideoEncoder3D(
+                in_channels=config.in_channels,
+                out_channels=config.latent_channels,
+                down_block_types=config.down_block_types,
+                block_out_channels=config.block_out_channels,
+                layers_per_block=config.layers_per_block,
+                norm_num_groups=config.norm_num_groups,
+                act_fn=config.act_fn,
+                double_z=True,
+                mid_block_add_attention=config.mid_block_add_attention,
+                temporal_compression_ratio=config.temporal_compression_ratio,
+                spatial_compression_ratio=config.spatial_compression_ratio,
+            )
+            self.quant_conv = nn.Conv3d(
+                2 * config.latent_channels, 2 * config.latent_channels, kernel_size=1
+            )
+
+        if config.load_decoder:
+            self.decoder = HunyuanVideoDecoder3D(
+                in_channels=config.latent_channels,
+                out_channels=config.out_channels,
+                up_block_types=config.up_block_types,
+                block_out_channels=config.block_out_channels,
+                layers_per_block=config.layers_per_block,
+                norm_num_groups=config.norm_num_groups,
+                act_fn=config.act_fn,
+                time_compression_ratio=config.temporal_compression_ratio,
+                spatial_compression_ratio=config.spatial_compression_ratio,
+                mid_block_add_attention=config.mid_block_add_attention,
+            )
+            self.post_quant_conv = nn.Conv3d(
+                config.latent_channels, config.latent_channels, kernel_size=1
+            )
+            self._spatial_parallel_decode_enabled = False
+            if can_install_spatial_shard_parallel_decode(self.config):
+                _enable_hunyuan_decoder_spatial_parallel(self.decoder)
+                self._spatial_parallel_decode_enabled = True
+
+    def _should_use_spatial_parallel_decode(self, z: torch.Tensor) -> bool:
+        return (
+            self._spatial_parallel_decode_enabled
+            and should_run_spatial_shard_parallel_decode(self.config, z)
+        )
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.encoder(x)
+        enc = self.quant_conv(x)
+        return enc
+
+    def _decode(self, z: torch.Tensor) -> torch.Tensor:
+        z = self.post_quant_conv(z)
+        if self._should_use_spatial_parallel_decode(z):
+            expected_height = (
+                z.shape[-2] * self.config.arch_config.spatial_compression_ratio
+            )
+            z, expected_height = split_height_for_parallel_decode(
+                z,
+                expected_height=expected_height,
+                world_size=get_decode_parallel_world_size(),
+                rank=get_decode_parallel_rank(),
+            )
+            dec = gather_and_trim_height(self.decoder(z), expected_height)
+        elif self._spatial_parallel_decode_enabled:
+            with disable_spatial_parallel_decode():
+                dec = self.decoder(z)
+        else:
+            dec = self.decoder(z)
+        return dec
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        generator: torch.Generator | None = None,
+    ) -> torch.Tensor:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+        """
+        x = sample
+        posterior = self.encode(x).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z)
+        return dec
+
+
+EntryClass = AutoencoderKLHunyuanVideo

python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_3_condition_encoder.py

@@ -0,0 +1,210 @@
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.vaes.ltx_2_vae import (
+    LTX2VideoCausalConv3d,
+    LTX2VideoResnetBlock3d,
+    LTXVideoDownsampler3d,
+)
+
+
+def _patchify_video(sample: torch.Tensor, patch_size: int) -> torch.Tensor:
+    if patch_size == 1:
+        return sample
+    batch_size, channels, num_frames, height, width = sample.shape
+    sample = sample.reshape(
+        batch_size,
+        channels,
+        num_frames,
+        1,
+        height // patch_size,
+        patch_size,
+        width // patch_size,
+        patch_size,
+    )
+    return sample.permute(0, 1, 3, 7, 5, 2, 4, 6).flatten(1, 4)
+
+
+class LTX23VideoPixelNorm(nn.Module):
+    def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True)
+        rms = torch.sqrt(mean_sq + self.eps)
+        return x / rms
+
+
+class LTX23PerChannelStatistics(nn.Module):
+    def __init__(self, latent_channels: int) -> None:
+        super().__init__()
+        self.register_buffer("std-of-means", torch.empty(latent_channels))
+        self.register_buffer("mean-of-means", torch.empty(latent_channels))
+
+    def normalize(self, x: torch.Tensor) -> torch.Tensor:
+        mean = self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)
+        std = self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)
+        return (x - mean) / std
+
+
+class LTX23VideoResBlockStack(nn.Module):
+    def __init__(
+        self, channels: int, num_layers: int, spatial_padding_mode: str
+    ) -> None:
+        super().__init__()
+        self.res_blocks = nn.ModuleList(
+            [
+                LTX2VideoResnetBlock3d(
+                    in_channels=channels,
+                    out_channels=channels,
+                    spatial_padding_mode=spatial_padding_mode,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for res_block in self.res_blocks:
+            hidden_states = res_block(hidden_states, causal=True)
+        return hidden_states
+
+
+def _make_ltx23_encoder_block(
+    block_name: str,
+    block_config: dict[str, Any],
+    in_channels: int,
+    spatial_padding_mode: str,
+) -> tuple[nn.Module, int]:
+    if block_name == "res_x":
+        return (
+            LTX23VideoResBlockStack(
+                channels=in_channels,
+                num_layers=int(block_config["num_layers"]),
+                spatial_padding_mode=spatial_padding_mode,
+            ),
+            in_channels,
+        )
+
+    multiplier = int(block_config.get("multiplier", 2))
+    stride_map = {
+        "compress_space_res": (1, 2, 2),
+        "compress_time_res": (2, 1, 1),
+        "compress_all_res": (2, 2, 2),
+    }
+    stride = stride_map.get(block_name)
+    if stride is None:
+        raise ValueError(f"Unsupported LTX-2.3 encoder block: {block_name}")
+    out_channels = in_channels * multiplier
+    return (
+        LTXVideoDownsampler3d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            stride=stride,
+            spatial_padding_mode=spatial_padding_mode,
+        ),
+        out_channels,
+    )
+
+
+class LTX23VideoConditionEncoder(nn.Module, LayerwiseOffloadableModuleMixin):
+    layerwise_offload_dit_group_enabled = False
+    layer_names = ["down_blocks"]
+
+    def __init__(self, config: dict[str, Any]) -> None:
+        super().__init__()
+
+        vae_config = config.get("vae", config)
+        latent_channels = int(vae_config["latent_channels"])
+        patch_size = int(vae_config.get("patch_size", 4))
+        spatial_padding_mode = str(vae_config.get("spatial_padding_mode", "zeros"))
+        encoder_blocks = list(vae_config["encoder_blocks"])
+        latent_log_var = str(vae_config.get("latent_log_var", "uniform"))
+
+        self.patch_size = patch_size
+        self.latency_channels = latent_channels
+        self.latent_log_var = latent_log_var
+        self.per_channel_statistics = LTX23PerChannelStatistics(latent_channels)
+
+        feature_channels = latent_channels
+        self.conv_in = LTX2VideoCausalConv3d(
+            in_channels=int(vae_config.get("in_channels", 3)) * patch_size**2,
+            out_channels=feature_channels,
+            kernel_size=3,
+            stride=1,
+            spatial_padding_mode=spatial_padding_mode,
+        )
+
+        self.down_blocks = nn.ModuleList()
+        for block_name, block_params in encoder_blocks:
+            block_config = (
+                {"num_layers": block_params}
+                if isinstance(block_params, int)
+                else dict(block_params)
+            )
+            block, feature_channels = _make_ltx23_encoder_block(
+                block_name=block_name,
+                block_config=block_config,
+                in_channels=feature_channels,
+                spatial_padding_mode=spatial_padding_mode,
+            )
+            self.down_blocks.append(block)
+
+        self.conv_norm_out = LTX23VideoPixelNorm(dim=1, eps=1e-8)
+        self.conv_act = nn.SiLU()
+
+        conv_out_channels = latent_channels
+        if latent_log_var == "per_channel":
+            conv_out_channels *= 2
+        elif latent_log_var in {"uniform", "constant"}:
+            conv_out_channels += 1
+        elif latent_log_var != "none":
+            raise ValueError(f"Unsupported latent_log_var: {latent_log_var}")
+
+        self.conv_out = LTX2VideoCausalConv3d(
+            in_channels=feature_channels,
+            out_channels=conv_out_channels,
+            kernel_size=3,
+            stride=1,
+            spatial_padding_mode=spatial_padding_mode,
+        )
+
+    def forward(self, sample: torch.Tensor) -> torch.Tensor:
+        frames_count = int(sample.shape[2])
+        if (frames_count - 1) % 8 != 0:
+            frames_to_crop = (frames_count - 1) % 8
+            sample = sample[:, :, :-frames_to_crop, ...]
+
+        hidden_states = _patchify_video(sample, self.patch_size)
+        hidden_states = self.conv_in(hidden_states, causal=True)
+
+        for block in self.down_blocks:
+            hidden_states = block(hidden_states)
+
+        hidden_states = self.conv_norm_out(hidden_states)
+        hidden_states = self.conv_act(hidden_states)
+        hidden_states = self.conv_out(hidden_states, causal=True)
+
+        if self.latent_log_var == "uniform":
+            means = hidden_states[:, :-1, ...]
+            logvar = hidden_states[:, -1:, ...]
+            hidden_states = torch.cat(
+                [
+                    means,
+                    logvar.repeat(1, means.shape[1], *([1] * (means.ndim - 2))),
+                ],
+                dim=1,
+            )
+        elif self.latent_log_var == "constant":
+            means = hidden_states[:, :-1, ...]
+            logvar = torch.full_like(means, -30.0)
+            hidden_states = torch.cat([means, logvar], dim=1)
+
+        means, _ = torch.chunk(hidden_states, 2, dim=1)
+        return self.per_channel_statistics.normalize(means)

python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_audio.py

@@ -0,0 +1,918 @@
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from diffusers.models.autoencoders.vae import (
+    DecoderOutput,
+    DiagonalGaussianDistribution,
+)
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.vaes.ltx_audio import LTXAudioVAEConfig
+from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE
+
+LATENT_DOWNSAMPLE_FACTOR = 4
+
+
+class LTX2AudioCausalConv2d(nn.Module):
+    """
+    A causal 2D convolution that pads asymmetrically along the causal axis.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, Tuple[int, int]],
+        stride: int = 1,
+        dilation: Union[int, Tuple[int, int]] = 1,
+        groups: int = 1,
+        bias: bool = True,
+        causality_axis: str = "height",
+    ) -> None:
+        super().__init__()
+
+        self.causality_axis = causality_axis
+        kernel_size = (
+            (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
+        )
+        dilation = (dilation, dilation) if isinstance(dilation, int) else dilation
+
+        pad_h = (kernel_size[0] - 1) * dilation[0]
+        pad_w = (kernel_size[1] - 1) * dilation[1]
+
+        if self.causality_axis == "none":
+            padding = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)
+        elif self.causality_axis in {"width", "width-compatibility"}:
+            padding = (pad_w, 0, pad_h // 2, pad_h - pad_h // 2)
+        elif self.causality_axis == "height":
+            padding = (pad_w // 2, pad_w - pad_w // 2, pad_h, 0)
+        else:
+            raise ValueError(f"Invalid causality_axis: {causality_axis}")
+
+        self.padding = padding
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=0,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.pad(x, self.padding)
+        return self.conv(x)
+
+
+class LTX2AudioPixelNorm(nn.Module):
+    """
+    Per-pixel (per-location) RMS normalization layer.
+    """
+
+    def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True)
+        rms = torch.sqrt(mean_sq + self.eps)
+        return x / rms
+
+
+class LTX2AudioAttnBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        norm_type: str = "group",
+    ) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+
+        if norm_type == "group":
+            self.norm = nn.GroupNorm(
+                num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+            )
+        elif norm_type == "pixel":
+            self.norm = LTX2AudioPixelNorm(dim=1, eps=1e-6)
+        else:
+            raise ValueError(f"Invalid normalization type: {norm_type}")
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h_ = self.norm(x)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        batch, channels, height, width = q.shape
+        q = q.reshape(batch, channels, height * width).permute(0, 2, 1).contiguous()
+        k = k.reshape(batch, channels, height * width).contiguous()
+        attn = torch.bmm(q, k) * (int(channels) ** (-0.5))
+        attn = torch.nn.functional.softmax(attn, dim=2)
+
+        v = v.reshape(batch, channels, height * width)
+        attn = attn.permute(0, 2, 1).contiguous()
+        h_ = torch.bmm(v, attn).reshape(batch, channels, height, width)
+
+        h_ = self.proj_out(h_)
+        return x + h_
+
+
+class LTX2AudioResnetBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        conv_shortcut: bool = False,
+        dropout: float = 0.0,
+        temb_channels: int = 512,
+        norm_type: str = "group",
+        causality_axis: str = "height",
+    ) -> None:
+        super().__init__()
+        self.causality_axis = causality_axis
+
+        if (
+            self.causality_axis is not None
+            and self.causality_axis != "none"
+            and norm_type == "group"
+        ):
+            raise ValueError("Causal ResnetBlock with GroupNorm is not supported.")
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        if norm_type == "group":
+            self.norm1 = nn.GroupNorm(
+                num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+            )
+        elif norm_type == "pixel":
+            self.norm1 = LTX2AudioPixelNorm(dim=1, eps=1e-6)
+        else:
+            raise ValueError(f"Invalid normalization type: {norm_type}")
+        self.non_linearity = nn.SiLU()
+        if causality_axis is not None:
+            self.conv1 = LTX2AudioCausalConv2d(
+                in_channels,
+                out_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=causality_axis,
+            )
+        else:
+            self.conv1 = nn.Conv2d(
+                in_channels, out_channels, kernel_size=3, stride=1, padding=1
+            )
+        if temb_channels > 0:
+            self.temb_proj = nn.Linear(temb_channels, out_channels)
+        if norm_type == "group":
+            self.norm2 = nn.GroupNorm(
+                num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
+            )
+        elif norm_type == "pixel":
+            self.norm2 = LTX2AudioPixelNorm(dim=1, eps=1e-6)
+        else:
+            raise ValueError(f"Invalid normalization type: {norm_type}")
+        self.dropout = nn.Dropout(dropout)
+        if causality_axis is not None:
+            self.conv2 = LTX2AudioCausalConv2d(
+                out_channels,
+                out_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=causality_axis,
+            )
+        else:
+            self.conv2 = nn.Conv2d(
+                out_channels, out_channels, kernel_size=3, stride=1, padding=1
+            )
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                if causality_axis is not None:
+                    self.conv_shortcut = LTX2AudioCausalConv2d(
+                        in_channels,
+                        out_channels,
+                        kernel_size=3,
+                        stride=1,
+                        causality_axis=causality_axis,
+                    )
+                else:
+                    self.conv_shortcut = nn.Conv2d(
+                        in_channels, out_channels, kernel_size=3, stride=1, padding=1
+                    )
+            else:
+                if causality_axis is not None:
+                    self.nin_shortcut = LTX2AudioCausalConv2d(
+                        in_channels,
+                        out_channels,
+                        kernel_size=1,
+                        stride=1,
+                        causality_axis=causality_axis,
+                    )
+                else:
+                    self.nin_shortcut = nn.Conv2d(
+                        in_channels, out_channels, kernel_size=1, stride=1, padding=0
+                    )
+
+    def forward(
+        self, x: torch.Tensor, temb: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        h = self.norm1(x)
+        h = self.non_linearity(h)
+        h = self.conv1(h)
+
+        if temb is not None:
+            h = h + self.temb_proj(self.non_linearity(temb))[:, :, None, None]
+
+        h = self.norm2(h)
+        h = self.non_linearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = (
+                self.conv_shortcut(x)
+                if self.use_conv_shortcut
+                else self.nin_shortcut(x)
+            )
+
+        return x + h
+
+
+class LTX2AudioDownsample(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        with_conv: bool,
+        causality_axis: Optional[str] = "height",
+    ) -> None:
+        super().__init__()
+        self.with_conv = with_conv
+        self.causality_axis = causality_axis
+
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(
+                in_channels, in_channels, kernel_size=3, stride=2, padding=0
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.with_conv:
+            # Padding tuple is in the order: (left, right, top, bottom).
+            if self.causality_axis == "none":
+                pad = (0, 1, 0, 1)
+            elif self.causality_axis == "width":
+                pad = (2, 0, 0, 1)
+            elif self.causality_axis == "height":
+                pad = (0, 1, 2, 0)
+            elif self.causality_axis == "width-compatibility":
+                pad = (1, 0, 0, 1)
+            else:
+                raise ValueError(
+                    f"Invalid `causality_axis` {self.causality_axis}; supported values are `none`, `width`, `height`,"
+                    f" and `width-compatibility`."
+                )
+
+            x = F.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            # with_conv=False implies that causality_axis is "none"
+            x = F.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+class LTX2AudioUpsample(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        with_conv: bool,
+        causality_axis: Optional[str] = "height",
+    ) -> None:
+        super().__init__()
+        self.with_conv = with_conv
+        self.causality_axis = causality_axis
+        if self.with_conv:
+            if causality_axis is not None:
+                self.conv = LTX2AudioCausalConv2d(
+                    in_channels,
+                    in_channels,
+                    kernel_size=3,
+                    stride=1,
+                    causality_axis=causality_axis,
+                )
+            else:
+                self.conv = nn.Conv2d(
+                    in_channels, in_channels, kernel_size=3, stride=1, padding=1
+                )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+            if self.causality_axis is None or self.causality_axis == "none":
+                pass
+            elif self.causality_axis == "height":
+                x = x[:, :, 1:, :]
+            elif self.causality_axis == "width":
+                x = x[:, :, :, 1:]
+            elif self.causality_axis == "width-compatibility":
+                pass
+            else:
+                raise ValueError(f"Invalid causality_axis: {self.causality_axis}")
+
+        return x
+
+
+class LTX2AudioAudioPatchifier:
+    """
+    Patchifier for spectrogram/audio latents.
+    """
+
+    def __init__(
+        self,
+        patch_size: int,
+        sample_rate: int = 16000,
+        hop_length: int = 160,
+        audio_latent_downsample_factor: int = 4,
+        is_causal: bool = True,
+    ):
+        self.hop_length = hop_length
+        self.sample_rate = sample_rate
+        self.audio_latent_downsample_factor = audio_latent_downsample_factor
+        self.is_causal = is_causal
+        self._patch_size = (1, patch_size, patch_size)
+
+    def patchify(self, audio_latents: torch.Tensor) -> torch.Tensor:
+        batch, channels, time, freq = audio_latents.shape
+        return audio_latents.permute(0, 2, 1, 3).reshape(batch, time, channels * freq)
+
+    def unpatchify(
+        self, audio_latents: torch.Tensor, channels: int, mel_bins: int
+    ) -> torch.Tensor:
+        batch, time, _ = audio_latents.shape
+        return audio_latents.view(batch, time, channels, mel_bins).permute(0, 2, 1, 3)
+
+    @property
+    def patch_size(self) -> Tuple[int, int, int]:
+        return self._patch_size
+
+
+class LTX2AudioEncoder(nn.Module):
+    def __init__(
+        self,
+        base_channels: int = 128,
+        output_channels: int = 1,
+        num_res_blocks: int = 2,
+        attn_resolutions: Optional[Tuple[int, ...]] = None,
+        in_channels: int = 2,
+        resolution: int = 256,
+        latent_channels: int = 8,
+        ch_mult: Tuple[int, ...] = (1, 2, 4),
+        norm_type: str = "group",
+        causality_axis: Optional[str] = "width",
+        dropout: float = 0.0,
+        mid_block_add_attention: bool = False,
+        sample_rate: int = 16000,
+        mel_hop_length: int = 160,
+        is_causal: bool = True,
+        mel_bins: Optional[int] = 64,
+        double_z: bool = True,
+    ):
+        super().__init__()
+
+        self.sample_rate = sample_rate
+        self.mel_hop_length = mel_hop_length
+        self.is_causal = is_causal
+        self.mel_bins = mel_bins
+
+        self.base_channels = base_channels
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.out_ch = output_channels
+        self.give_pre_end = False
+        self.tanh_out = False
+        self.norm_type = norm_type
+        self.latent_channels = latent_channels
+        self.channel_multipliers = ch_mult
+        self.attn_resolutions = attn_resolutions
+        self.causality_axis = causality_axis
+
+        base_block_channels = base_channels
+        base_resolution = resolution
+        self.z_shape = (1, latent_channels, base_resolution, base_resolution)
+
+        if self.causality_axis is not None:
+            self.conv_in = LTX2AudioCausalConv2d(
+                in_channels,
+                base_block_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=self.causality_axis,
+            )
+        else:
+            self.conv_in = nn.Conv2d(
+                in_channels, base_block_channels, kernel_size=3, stride=1, padding=1
+            )
+
+        self.down = nn.ModuleList()
+        block_in = base_block_channels
+        curr_res = self.resolution
+
+        for level in range(self.num_resolutions):
+            stage = nn.Module()
+            stage.block = nn.ModuleList()
+            stage.attn = nn.ModuleList()
+            block_out = self.base_channels * self.channel_multipliers[level]
+
+            for _ in range(self.num_res_blocks):
+                stage.block.append(
+                    LTX2AudioResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                        norm_type=self.norm_type,
+                        causality_axis=self.causality_axis,
+                    )
+                )
+                block_in = block_out
+                if self.attn_resolutions:
+                    if curr_res in self.attn_resolutions:
+                        stage.attn.append(
+                            LTX2AudioAttnBlock(block_in, norm_type=self.norm_type)
+                        )
+
+            if level != self.num_resolutions - 1:
+                stage.downsample = LTX2AudioDownsample(
+                    block_in, True, causality_axis=self.causality_axis
+                )
+                curr_res = curr_res // 2
+
+            self.down.append(stage)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = LTX2AudioResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+        )
+        if mid_block_add_attention:
+            self.mid.attn_1 = LTX2AudioAttnBlock(block_in, norm_type=self.norm_type)
+        else:
+            self.mid.attn_1 = nn.Identity()
+        self.mid.block_2 = LTX2AudioResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+        )
+
+        final_block_channels = block_in
+        z_channels = 2 * latent_channels if double_z else latent_channels
+        if self.norm_type == "group":
+            self.norm_out = nn.GroupNorm(
+                num_groups=32, num_channels=final_block_channels, eps=1e-6, affine=True
+            )
+        elif self.norm_type == "pixel":
+            self.norm_out = LTX2AudioPixelNorm(dim=1, eps=1e-6)
+        else:
+            raise ValueError(f"Invalid normalization type: {self.norm_type}")
+        self.non_linearity = nn.SiLU()
+
+        if self.causality_axis is not None:
+            self.conv_out = LTX2AudioCausalConv2d(
+                final_block_channels,
+                z_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=self.causality_axis,
+            )
+        else:
+            self.conv_out = nn.Conv2d(
+                final_block_channels, z_channels, kernel_size=3, stride=1, padding=1
+            )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # hidden_states expected shape: (batch_size, channels, time, num_mel_bins)
+        hidden_states = self.conv_in(hidden_states)
+
+        for level in range(self.num_resolutions):
+            stage = self.down[level]
+            for block_idx, block in enumerate(stage.block):
+                hidden_states = block(hidden_states, temb=None)
+                if stage.attn:
+                    hidden_states = stage.attn[block_idx](hidden_states)
+
+            if level != self.num_resolutions - 1 and hasattr(stage, "downsample"):
+                hidden_states = stage.downsample(hidden_states)
+
+        hidden_states = self.mid.block_1(hidden_states, temb=None)
+        hidden_states = self.mid.attn_1(hidden_states)
+        hidden_states = self.mid.block_2(hidden_states, temb=None)
+
+        hidden_states = self.norm_out(hidden_states)
+        hidden_states = self.non_linearity(hidden_states)
+        hidden_states = self.conv_out(hidden_states)
+
+        return hidden_states
+
+
+class LTX2AudioDecoder(nn.Module):
+    """
+    Symmetric decoder that reconstructs audio spectrograms from latent features.
+
+    The decoder mirrors the encoder structure with configurable channel multipliers, attention resolutions, and causal
+    convolutions.
+    """
+
+    def __init__(
+        self,
+        base_channels: int = 128,
+        output_channels: int = 1,
+        num_res_blocks: int = 2,
+        attn_resolutions: Optional[Tuple[int, ...]] = None,
+        in_channels: int = 2,
+        resolution: int = 256,
+        latent_channels: int = 8,
+        ch_mult: Tuple[int, ...] = (1, 2, 4),
+        norm_type: str = "group",
+        causality_axis: Optional[str] = "width",
+        dropout: float = 0.0,
+        mid_block_add_attention: bool = False,
+        sample_rate: int = 16000,
+        mel_hop_length: int = 160,
+        is_causal: bool = True,
+        mel_bins: Optional[int] = 64,
+    ) -> None:
+        super().__init__()
+
+        self.sample_rate = sample_rate
+        self.mel_hop_length = mel_hop_length
+        self.is_causal = is_causal
+        self.mel_bins = mel_bins
+        self.patchifier = LTX2AudioAudioPatchifier(
+            patch_size=1,
+            audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
+            sample_rate=sample_rate,
+            hop_length=mel_hop_length,
+            is_causal=is_causal,
+        )
+
+        self.base_channels = base_channels
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.out_ch = output_channels
+        self.give_pre_end = False
+        self.tanh_out = False
+        self.norm_type = norm_type
+        self.latent_channels = latent_channels
+        self.channel_multipliers = ch_mult
+        self.attn_resolutions = attn_resolutions
+        self.causality_axis = causality_axis
+
+        base_block_channels = base_channels * self.channel_multipliers[-1]
+        base_resolution = resolution // (2 ** (self.num_resolutions - 1))
+        self.z_shape = (1, latent_channels, base_resolution, base_resolution)
+
+        if self.causality_axis is not None:
+            self.conv_in = LTX2AudioCausalConv2d(
+                latent_channels,
+                base_block_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=self.causality_axis,
+            )
+        else:
+            self.conv_in = nn.Conv2d(
+                latent_channels, base_block_channels, kernel_size=3, stride=1, padding=1
+            )
+        self.non_linearity = nn.SiLU()
+        self.mid = nn.Module()
+        self.mid.block_1 = LTX2AudioResnetBlock(
+            in_channels=base_block_channels,
+            out_channels=base_block_channels,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+        )
+        if mid_block_add_attention:
+            self.mid.attn_1 = LTX2AudioAttnBlock(
+                base_block_channels, norm_type=self.norm_type
+            )
+        else:
+            self.mid.attn_1 = nn.Identity()
+        self.mid.block_2 = LTX2AudioResnetBlock(
+            in_channels=base_block_channels,
+            out_channels=base_block_channels,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+        )
+
+        self.up = nn.ModuleList()
+        block_in = base_block_channels
+        curr_res = self.resolution // (2 ** (self.num_resolutions - 1))
+
+        for level in reversed(range(self.num_resolutions)):
+            stage = nn.Module()
+            stage.block = nn.ModuleList()
+            stage.attn = nn.ModuleList()
+            block_out = self.base_channels * self.channel_multipliers[level]
+
+            for _ in range(self.num_res_blocks + 1):
+                stage.block.append(
+                    LTX2AudioResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                        norm_type=self.norm_type,
+                        causality_axis=self.causality_axis,
+                    )
+                )
+                block_in = block_out
+                if self.attn_resolutions:
+                    if curr_res in self.attn_resolutions:
+                        stage.attn.append(
+                            LTX2AudioAttnBlock(block_in, norm_type=self.norm_type)
+                        )
+
+            if level != 0:
+                stage.upsample = LTX2AudioUpsample(
+                    block_in, True, causality_axis=self.causality_axis
+                )
+                curr_res *= 2
+
+            self.up.insert(0, stage)
+
+        final_block_channels = block_in
+
+        if self.norm_type == "group":
+            self.norm_out = nn.GroupNorm(
+                num_groups=32, num_channels=final_block_channels, eps=1e-6, affine=True
+            )
+        elif self.norm_type == "pixel":
+            self.norm_out = LTX2AudioPixelNorm(dim=1, eps=1e-6)
+        else:
+            raise ValueError(f"Invalid normalization type: {self.norm_type}")
+
+        if self.causality_axis is not None:
+            self.conv_out = LTX2AudioCausalConv2d(
+                final_block_channels,
+                output_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=self.causality_axis,
+            )
+        else:
+            self.conv_out = nn.Conv2d(
+                final_block_channels,
+                output_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+            )
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+    ) -> torch.Tensor:
+        _, _, frames, mel_bins = sample.shape
+
+        target_frames = frames * LATENT_DOWNSAMPLE_FACTOR
+
+        if self.causality_axis is not None:
+            target_frames = max(target_frames - (LATENT_DOWNSAMPLE_FACTOR - 1), 1)
+
+        target_channels = self.out_ch
+        target_mel_bins = self.mel_bins if self.mel_bins is not None else mel_bins
+
+        hidden_features = self.conv_in(sample)
+        hidden_features = self.mid.block_1(hidden_features, temb=None)
+        hidden_features = self.mid.attn_1(hidden_features)
+        hidden_features = self.mid.block_2(hidden_features, temb=None)
+
+        for level in reversed(range(self.num_resolutions)):
+            stage = self.up[level]
+            for block_idx, block in enumerate(stage.block):
+                hidden_features = block(hidden_features, temb=None)
+                if stage.attn:
+                    hidden_features = stage.attn[block_idx](hidden_features)
+
+            if level != 0 and hasattr(stage, "upsample"):
+                hidden_features = stage.upsample(hidden_features)
+
+        if self.give_pre_end:
+            return hidden_features
+
+        hidden = self.norm_out(hidden_features)
+        hidden = self.non_linearity(hidden)
+        decoded_output = self.conv_out(hidden)
+        decoded_output = torch.tanh(decoded_output) if self.tanh_out else decoded_output
+
+        _, _, current_time, current_freq = decoded_output.shape
+        target_time = target_frames
+        target_freq = target_mel_bins
+
+        decoded_output = decoded_output[
+            :,
+            :target_channels,
+            : min(current_time, target_time),
+            : min(current_freq, target_freq),
+        ]
+
+        time_padding_needed = target_time - decoded_output.shape[2]
+        freq_padding_needed = target_freq - decoded_output.shape[3]
+
+        if time_padding_needed > 0 or freq_padding_needed > 0:
+            padding = (
+                0,
+                max(freq_padding_needed, 0),
+                0,
+                max(time_padding_needed, 0),
+            )
+            decoded_output = F.pad(decoded_output, padding)
+
+        decoded_output = decoded_output[:, :target_channels, :target_time, :target_freq]
+
+        return decoded_output
+
+
+class AutoencoderKLLTX2Audio(ParallelTiledVAE):
+    r"""
+    LTX2 audio VAE for encoding and decoding audio latent representations.
+    """
+
+    _supports_gradient_checkpointing = False
+
+    def __init__(
+        self,
+        config: LTXAudioVAEConfig,
+    ) -> None:
+        super().__init__(config=config)
+
+        causality_axis = config.arch_config.causality_axis
+        attn_resolutions = config.arch_config.attn_resolutions
+        base_channels = config.arch_config.base_channels
+        output_channels = config.arch_config.output_channels
+        ch_mult = config.arch_config.ch_mult
+        num_res_blocks = config.arch_config.num_res_blocks
+        in_channels = config.arch_config.in_channels
+        resolution = config.arch_config.resolution
+        latent_channels = config.arch_config.latent_channels
+        norm_type = config.arch_config.norm_type
+        dropout = config.arch_config.dropout
+        mid_block_add_attention = config.arch_config.mid_block_add_attention
+        sample_rate = config.arch_config.sample_rate
+        mel_hop_length = config.arch_config.mel_hop_length
+        is_causal = config.arch_config.is_causal
+        mel_bins = config.arch_config.mel_bins
+        double_z = config.arch_config.double_z
+
+        supported_causality_axes = {"none", "width", "height", "width-compatibility"}
+        if causality_axis not in supported_causality_axes:
+            raise ValueError(
+                f"{causality_axis=} is not valid. Supported values: {supported_causality_axes}"
+            )
+
+        attn_resolution_set = (
+            set(attn_resolutions) if attn_resolutions else attn_resolutions
+        )
+
+        self.encoder = LTX2AudioEncoder(
+            base_channels=base_channels,
+            output_channels=output_channels,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            attn_resolutions=attn_resolution_set,
+            in_channels=in_channels,
+            resolution=resolution,
+            latent_channels=latent_channels,
+            norm_type=norm_type,
+            causality_axis=causality_axis,
+            dropout=dropout,
+            mid_block_add_attention=mid_block_add_attention,
+            sample_rate=sample_rate,
+            mel_hop_length=mel_hop_length,
+            is_causal=is_causal,
+            mel_bins=mel_bins,
+            double_z=double_z,
+        )
+
+        self.decoder = LTX2AudioDecoder(
+            base_channels=base_channels,
+            output_channels=output_channels,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            attn_resolutions=attn_resolution_set,
+            in_channels=in_channels,
+            resolution=resolution,
+            latent_channels=latent_channels,
+            norm_type=norm_type,
+            causality_axis=causality_axis,
+            dropout=dropout,
+            mid_block_add_attention=mid_block_add_attention,
+            sample_rate=sample_rate,
+            mel_hop_length=mel_hop_length,
+            is_causal=is_causal,
+            mel_bins=mel_bins,
+        )
+
+        # Per-channel statistics for normalizing and denormalizing the latent representation. This statistics is computed over
+        # the entire dataset and stored in model's checkpoint under AudioVAE state_dict
+        latents_std = torch.ones((base_channels,))
+        latents_mean = torch.zeros((base_channels,))
+        self.register_buffer("latents_mean", latents_mean, persistent=True)
+        self.register_buffer("latents_std", latents_std, persistent=True)
+
+        # TODO: confirm whether the mel compression ratio below is correct
+        self.mel_compression_ratio = LATENT_DOWNSAMPLE_FACTOR
+        self.use_slicing = False
+        # stage containers are not called directly, so hooks attach to called lists
+        self.layer_names = [
+            layer_name
+            for prefix, num_resolutions in (
+                ("encoder.down", self.encoder.num_resolutions),
+                ("decoder.up", self.decoder.num_resolutions),
+            )
+            for level in range(num_resolutions)
+            for layer_name in (
+                f"{prefix}.{level}.block",
+                f"{prefix}.{level}.attn",
+            )
+        ]
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        return self.encoder(x)
+
+    def encode(self, x: torch.Tensor, return_dict: bool = True):
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self._encode(x)
+        posterior = DiagonalGaussianDistribution(h)
+
+        if not return_dict:
+            return (posterior,)
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def _decode(self, z: torch.Tensor) -> torch.Tensor:
+        return self.decoder(z)
+
+    def decode(
+        self, z: torch.Tensor, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice) for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z)
+
+        if not return_dict:
+            return (decoded,)
+
+        return DecoderOutput(sample=decoded)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        return_dict: bool = True,
+        generator: Optional[torch.Generator] = None,
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        posterior = self.encode(sample).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z)
+        if not return_dict:
+            return (dec.sample,)
+        return dec
+
+
+EntryClass = AutoencoderKLLTX2Audio

python/sglang/multimodal_gen/runtime/models/vaes/parallel/diffusers_spatial.py

@@ -0,0 +1,109 @@
+import torch
+import torch.nn as nn
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_decode_parallel_rank,
+    get_decode_parallel_world_size,
+)
+from sglang.multimodal_gen.runtime.layers.parallel_conv import (
+    SpatialParallelConv2d,
+    chunk_height_by_sizes,
+    gather_and_trim_height,
+    gather_variable_height,
+    split_for_parallel_decode,
+)
+
+
+def count_decoder_spatial_upsamples(decoder: nn.Module) -> int:
+    return sum(
+        len(upsamplers)
+        for block in getattr(decoder, "up_blocks", [])
+        if (upsamplers := getattr(block, "upsamplers", None)) is not None
+    )
+
+
+def enable_diffusers_decoder_spatial_parallel(decoder: nn.Module) -> int:
+    _replace_conv2d_modules(decoder)
+    _patch_groupnorm_modules(decoder)
+    _patch_attention_modules(decoder)
+    return count_decoder_spatial_upsamples(decoder)
+
+
+def spatial_parallel_diffusers_decode(
+    decoder: nn.Module, z: torch.Tensor, upsample_count: int
+) -> torch.Tensor:
+    z, expected_height = split_for_parallel_decode(
+        z,
+        upsample_count=upsample_count,
+        world_size=get_decode_parallel_world_size(),
+        rank=get_decode_parallel_rank(),
+    )
+    return gather_and_trim_height(decoder(z), expected_height)
+
+
+def _replace_conv2d_modules(module: nn.Module) -> None:
+    for name, child in list(module.named_children()):
+        if type(child) is nn.Conv2d:
+            setattr(module, name, _make_spatial_conv2d(child))
+        else:
+            _replace_conv2d_modules(child)
+
+
+def _make_spatial_conv2d(conv: nn.Conv2d) -> SpatialParallelConv2d:
+    spatial_conv = SpatialParallelConv2d(
+        in_channels=conv.in_channels,
+        out_channels=conv.out_channels,
+        kernel_size=conv.kernel_size,
+        stride=conv.stride,
+        padding=conv.padding,
+        dilation=conv.dilation,
+        groups=conv.groups,
+        bias=conv.bias is not None,
+        padding_mode=conv.padding_mode,
+    )
+    spatial_conv.weight = conv.weight
+    spatial_conv.bias = conv.bias
+    return spatial_conv
+
+
+def _patch_attention_modules(module: nn.Module) -> None:
+    for child in module.children():
+        if child.__class__.__name__ == "Attention":
+            _patch_attention_forward(child)
+        _patch_attention_modules(child)
+
+
+def _patch_groupnorm_modules(module: nn.Module) -> None:
+    for child in module.children():
+        if type(child) is nn.GroupNorm:
+            _patch_groupnorm_forward(child)
+        else:
+            _patch_groupnorm_modules(child)
+
+
+def _patch_groupnorm_forward(norm: nn.GroupNorm) -> None:
+    original_forward = norm.forward
+
+    def spatial_parallel_forward(hidden_states):
+        if hidden_states.dim() < 4:
+            return original_forward(hidden_states)
+        hidden_states, heights = gather_variable_height(hidden_states)
+        hidden_states = hidden_states.contiguous()
+        hidden_states = original_forward(hidden_states)
+        return chunk_height_by_sizes(hidden_states, heights)
+
+    norm.forward = spatial_parallel_forward
+
+
+def _patch_attention_forward(attn: nn.Module) -> None:
+    original_forward = attn.forward
+
+    def spatial_parallel_forward(hidden_states, *args, **kwargs):
+        if hidden_states.dim() != 4:
+            return original_forward(hidden_states, *args, **kwargs)
+        hidden_states, heights = gather_variable_height(hidden_states)
+        hidden_states = hidden_states.contiguous()
+        hidden_states = original_forward(hidden_states, *args, **kwargs)
+        return chunk_height_by_sizes(hidden_states, heights)
+
+    attn.forward = spatial_parallel_forward