[diffusion] chore: minor refactor by streamlining the VAE class hierarchy (#16069)

python/sglang/multimodal_gen/runtime/loader/component_loader.py

@@ -192,11 +192,11 @@ class ComponentLoader(ABC):
             if consumed is None or consumed == 0.0:
                 consumed = gpu_mem_before_loading - current_gpu_mem
             logger.info(
-                f"Loaded %s: %s from {source}. avail mem: %.2f GB, %.2f GB consumed",
+                f"Loaded %s: %s from {source}. model size: %.2f GB, avail mem: %.2f GB",
                 module_name,
                 component.__class__.__name__,
-                current_gpu_mem,
                 consumed,
+                current_gpu_mem,
             )
         return component, consumed
 

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

@@ -22,7 +22,7 @@ from sglang.multimodal_gen.configs.models.vaes.flux import Flux2VAEConfig
 from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE
 
 
-class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE):
+class AutoencoderKLFlux2(ParallelTiledVAE):
     r"""
     A VAE model with KL loss for encoding images into latents and decoding latent representations into images.
 
@@ -39,7 +39,7 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE):
         self,
         config: Flux2VAEConfig,
     ):
-        super().__init__()
+        super().__init__(config=config)
 
         self.config = config
         arch_config = config.arch_config

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

@@ -14,6 +14,7 @@ from diffusers.models.modeling_outputs import AutoencoderKLOutput
 
 from sglang.multimodal_gen.configs.models.vaes.qwenimage import QwenImageVAEConfig
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)  # pylint: disable=invalid-name
@@ -757,7 +758,7 @@ class QwenImageDecoder3d(nn.Module):
         return x
 
 
-class AutoencoderKLQwenImage(nn.Module):
+class AutoencoderKLQwenImage(ParallelTiledVAE):
     r"""
     A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos.
 
@@ -773,7 +774,7 @@ class AutoencoderKLQwenImage(nn.Module):
         config: QwenImageVAEConfig,
     ) -> None:
         # fmt: on
-        super().__init__()
+        super().__init__(config=config)
         base_dim = config.arch_config.base_dim
         z_dim = config.arch_config.z_dim
         dim_mult = config.arch_config.dim_mult
@@ -789,19 +790,18 @@ class AutoencoderKLQwenImage(nn.Module):
         self.latents_mean = config.arch_config.latents_mean
         self.config = config.arch_config
 
-
         self.encoder = QwenImageEncoder3d(
-            base_dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout, input_channels=self.input_channels
+            base_dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout,
+            input_channels=self.input_channels
         )
         self.quant_conv = QwenImageCausalConv3d(z_dim * 2, z_dim * 2, 1)
         self.post_quant_conv = QwenImageCausalConv3d(z_dim, z_dim, 1)
 
         self.decoder = QwenImageDecoder3d(
-            base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout, input_channels=self.input_channels
+            base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout,
+            input_channels=self.input_channels
         )
 
-        self.spatial_compression_ratio = 2 ** len(self.temperal_downsample)
-
         # When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension
         # to perform decoding of a single video latent at a time.
         self.use_slicing = False
@@ -840,8 +840,6 @@ class AutoencoderKLQwenImage(nn.Module):
             .view(1, latent_channels, 1, 1, 1)
             .to(cuda_device, dtype)
         )
-        latents_std_tensor = torch.tensor(config.arch_config.latents_std, dtype=dtype, device=cuda_device)
-        self.scaling_factor = (1.0 / latents_std_tensor).view(1, latent_channels, 1, 1, 1)
 
     def enable_tiling(
         self,

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

@@ -12,6 +12,7 @@ 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.runtime.distributed import (
@@ -20,7 +21,7 @@ from sglang.multimodal_gen.runtime.distributed import (
 )
 
 
-class ParallelTiledVAE(ABC):
+class ParallelTiledVAE(ABC, nn.Module):
     tile_sample_min_height: int
     tile_sample_min_width: int
     tile_sample_min_num_frames: int
@@ -33,6 +34,7 @@ class ParallelTiledVAE(ABC):
     use_parallel_tiling: bool
 
     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
@@ -45,10 +47,6 @@ class ParallelTiledVAE(ABC):
         self.use_temporal_tiling = config.use_temporal_tiling
         self.use_parallel_tiling = config.use_parallel_tiling
 
-    def to(self, device) -> "ParallelTiledVAE":
-        # TODO: implement this
-        return self
-
     @property
     def device(self):
         return next(self.parameters()).device

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

@@ -760,7 +760,7 @@ class HunyuanVideoDecoder3D(nn.Module):
         return hidden_states
 
 
-class AutoencoderKLHunyuanVideo(nn.Module, ParallelTiledVAE):
+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).

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

@@ -1125,7 +1125,7 @@ def unpatchify(x, patch_size):
     return x
 
 
-class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
+class AutoencoderKLWan(ParallelTiledVAE):
     r"""
     A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos.
     Introduced in [Wan 2.1].

python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py

@@ -10,6 +10,7 @@ This module contains implementations of image encoding stages for diffusion pipe
 import PIL
 import torch
 from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
 
 from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
     qwen_image_postprocess_text,
@@ -279,9 +280,12 @@ class ImageVAEEncodingStage(PipelineStage):
                 #     self.vae.enable_parallel()
                 if not vae_autocast_enabled:
                     video_condition = video_condition.to(vae_dtype)
-                encoder_output: DiagonalGaussianDistribution = self.vae.encode(
+                latent_dist: DiagonalGaussianDistribution = self.vae.encode(
                     video_condition
                 )
+                # for auto_encoder from diffusers
+                if isinstance(latent_dist, AutoencoderKLOutput):
+                    latent_dist = latent_dist.latent_dist
 
             generator = batch.generator
             if generator is None:
@@ -290,7 +294,7 @@ class ImageVAEEncodingStage(PipelineStage):
             sample_mode = server_args.pipeline_config.vae_config.encode_sample_mode()
 
             latent_condition = self.retrieve_latents(
-                encoder_output, generator, sample_mode=sample_mode
+                latent_dist, generator, sample_mode=sample_mode
             )
             latent_condition = server_args.pipeline_config.postprocess_vae_encode(
                 latent_condition, self.vae