From b5d9fc873bf238ff1172b8050dfa1fbb665f0f4d Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 29 Dec 2025 23:37:59 +0800 Subject: [PATCH] [diffusion] chore: minor refactor by streamlining the VAE class hierarchy (#16069) --- .../runtime/loader/component_loader.py | 4 ++-- .../runtime/models/vaes/autoencoder_kl_flux2.py | 4 ++-- .../models/vaes/autoencoder_kl_qwenimage.py | 16 +++++++--------- .../multimodal_gen/runtime/models/vaes/common.py | 8 +++----- .../runtime/models/vaes/hunyuanvae.py | 2 +- .../multimodal_gen/runtime/models/vaes/wanvae.py | 2 +- .../pipelines_core/stages/image_encoding.py | 8 ++++++-- 7 files changed, 22 insertions(+), 22 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index 918dbb37a..98f597541 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_flux2.py b/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_flux2.py index 9df72c09c..7410358fb 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_flux2.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py b/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py index 01d3c0114..42c5426d7 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/common.py b/python/sglang/multimodal_gen/runtime/models/vaes/common.py index 58dad2507..d3011b623 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/common.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py b/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py index a9de61a54..972967fa1 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py +++ b/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). diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py b/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py index 1018d43be..336c2fb5c 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py +++ b/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]. diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py index 6b22afc49..91f637e66 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py +++ b/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