[diffusion] feat: support distilled vae generic (#14195)

Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/configs/models/vaes/flux.py

@@ -47,9 +47,21 @@ class FluxVAEConfig(VAEConfig):
         ) * 2
 
     def post_init(self):
-        self.arch_config.vae_scale_factor = 2 ** (
-            len(self.arch_config.block_out_channels) - 1
-        )
+        # Calculate vae_scale_factor: prefer block_out_channels, fallback to dim_mult or scale_factor_spatial
+        if (
+            hasattr(self.arch_config, "block_out_channels")
+            and self.arch_config.block_out_channels
+        ):
+            self.arch_config.vae_scale_factor = 2 ** (
+                len(self.arch_config.block_out_channels) - 1
+            )
+        elif self.arch_config.dim_mult:
+            self.arch_config.vae_scale_factor = 2 ** (
+                len(self.arch_config.dim_mult) - 1
+            )
+        else:
+            self.arch_config.vae_scale_factor = self.arch_config.scale_factor_spatial
+
         self.arch_config.spatial_compression_ratio = self.arch_config.vae_scale_factor
 
 

python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py

@@ -5,6 +5,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     SlidingTileAttnConfig,
 )
 from sglang.multimodal_gen.configs.pipeline_configs.flux import FluxPipelineConfig
+from sglang.multimodal_gen.configs.pipeline_configs.flux_finetuned import (
+    Flux2FinetunedPipelineConfig,
+)
 from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import (
     FastHunyuanConfig,
     HunyuanConfig,
@@ -23,6 +26,7 @@ __all__ = [
     "HunyuanConfig",
     "FastHunyuanConfig",
     "FluxPipelineConfig",
+    "Flux2FinetunedPipelineConfig",
     "PipelineConfig",
     "SlidingTileAttnConfig",
     "WanT2V480PConfig",

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -281,7 +281,7 @@ class PipelineConfig:
         return image_latents
 
     # called after scale_and_shift, before vae decoding
-    def preprocess_decoding(self, latents):
+    def preprocess_decoding(self, latents, server_args=None, vae=None):
         return latents
 
     def gather_latents_for_sp(self, latents):
@@ -487,9 +487,30 @@ class PipelineConfig:
             raise ValueError("model_path is required in kwargs")
 
         # 1. Get the pipeline config class from the registry
+        from sglang.multimodal_gen.configs.pipeline_configs.flux import (
+            Flux2PipelineConfig,
+        )
+
         model_info = get_model_info(model_path)
 
-        pipeline_config = model_info.pipeline_config_cls()
+        # 1.5. Adjust pipeline config for fine-tuned VAE if needed
+        pipeline_config_cls = model_info.pipeline_config_cls
+        vae_path = kwargs.get(prefix_with_dot + "vae_path") or kwargs.get("vae_path")
+
+        # Check if this is a Flux2 model with fal/FLUX.2-Tiny-AutoEncoder
+        if (
+            isinstance(pipeline_config_cls, type)
+            and issubclass(pipeline_config_cls, Flux2PipelineConfig)
+            and vae_path is not None
+            and "FLUX.2-Tiny-AutoEncoder" in vae_path
+        ):
+            from sglang.multimodal_gen.configs.pipeline_configs.flux_finetuned import (
+                Flux2FinetunedPipelineConfig,
+            )
+
+            pipeline_config_cls = Flux2FinetunedPipelineConfig
+
+        pipeline_config = pipeline_config_cls()
 
         # 2. Load PipelineConfig from a json file or a PipelineConfig object if provided
         if isinstance(pipeline_config_or_path, str):

python/sglang/multimodal_gen/configs/pipeline_configs/flux.py

@@ -215,10 +215,10 @@ def _prepare_latent_ids(
     t = torch.arange(1)  # [0] - time dimension
     h = torch.arange(height)
     w = torch.arange(width)
-    l = torch.arange(1)  # [0] - layer dimension
+    layer = torch.arange(1)  # [0] - layer dimension
 
     # Create position IDs: (H*W, 4)
-    latent_ids = torch.cartesian_prod(t, h, w, l)
+    latent_ids = torch.cartesian_prod(t, h, w, layer)
 
     # Expand to batch: (B, H*W, 4)
     latent_ids = latent_ids.unsqueeze(0).expand(batch_size, -1, -1)
@@ -289,9 +289,9 @@ def _prepare_text_ids(
         t = torch.arange(1) if t_coord is None else t_coord[i]
         h = torch.arange(1)
         w = torch.arange(1)
-        l = torch.arange(L)
+        layer = torch.arange(L)
 
-        coords = torch.cartesian_prod(t, h, w, l)
+        coords = torch.cartesian_prod(t, h, w, layer)
         out_ids.append(coords)
 
     return torch.stack(out_ids)
@@ -500,7 +500,6 @@ class Flux2PipelineConfig(FluxPipelineConfig):
         return image_latents
 
     def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch):
-
         txt_ids = _prepare_text_ids(prompt_embeds).to(device=device)
 
         img_ids = batch.latent_ids
@@ -550,19 +549,51 @@ class Flux2PipelineConfig(FluxPipelineConfig):
         image_latents = _patchify_latents(image_latents)
         return image_latents
 
-    def preprocess_decoding(self, latents):
-        latents = _unpatchify_latents(latents)
+    def _check_vae_has_bn(self, vae):
+        """Check if VAE has bn attribute (cached check to avoid repeated hasattr calls)."""
+        if not hasattr(self, "_vae_has_bn_cache"):
+            self._vae_has_bn_cache = hasattr(vae, "bn") and vae.bn is not None
+        return self._vae_has_bn_cache
+
+    def preprocess_decoding(self, latents, server_args=None, vae=None):
+        """Preprocess latents before decoding.
+
+        Dynamically adapts based on VAE type:
+        - Standard Flux2 VAE (has bn): needs unpatchify (128 channels -> 32 channels)
+        - Distilled VAE (no bn): keeps patchified latents (128 channels)
+        """
+        if vae is not None and self._check_vae_has_bn(vae):
+            return _unpatchify_latents(latents)
         return latents
 
     def get_decode_scale_and_shift(self, device, dtype, vae):
+        """Get scale and shift for decoding.
+
+        Dynamically adapts based on VAE type:
+        - Standard Flux2 VAE (has bn): uses BatchNorm statistics
+        - Distilled VAE (no bn): uses scaling_factor from config
+        """
         vae_arch_config = self.vae_config.arch_config
-        latents_bn_mean = (
-            vae.bn.running_mean.view(1, -1, 1, 1).to(device=device).to(device, dtype)
+
+        if self._check_vae_has_bn(vae):
+            # Standard Flux2 VAE: use BatchNorm statistics
+            latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(device, dtype)
+            latents_bn_std = torch.sqrt(
+                vae.bn.running_var.view(1, -1, 1, 1) + vae_arch_config.batch_norm_eps
+            ).to(device, dtype)
+            return 1 / latents_bn_std, latents_bn_mean
+
+        # Distilled VAE or unknown: use scaling_factor
+        scaling_factor = (
+            getattr(vae.config, "scaling_factor", None)
+            if hasattr(vae, "config")
+            else getattr(vae, "scaling_factor", None)
+        ) or getattr(vae_arch_config, "scaling_factor", 0.13025)
+
+        scale = torch.tensor(scaling_factor, device=device, dtype=dtype).view(
+            1, 1, 1, 1
         )
-        latents_bn_std = torch.sqrt(
-            vae.bn.running_var.view(1, -1, 1, 1) + vae_arch_config.batch_norm_eps
-        ).to(device, dtype)
-        return 1 / latents_bn_std, latents_bn_mean
+        return 1 / scale, None
 
     def post_denoising_loop(self, latents, batch):
         latent_ids = batch.latent_ids

python/sglang/multimodal_gen/configs/pipeline_configs/flux_finetuned.py

@@ -0,0 +1,103 @@
+"""
+Pipeline configuration for Flux fine-tuned/distilled models.
+
+This module provides specialized handling for Flux fine-tuned models from HuggingFace,
+such as fal/FLUX.2-Tiny-AutoEncoder and other community fine-tuned variants.
+
+Key differences from standard Flux2PipelineConfig:
+- Handles custom VAE architectures loaded via auto_map
+- Supports both patchified (128 channels) and unpatchified (32 channels) latents
+- Dynamically adapts scale/shift based on VAE type
+- Properly handles 5D latents (batch, channels, frames, height, width) for decoding
+"""
+
+from dataclasses import dataclass
+
+import torch
+
+from sglang.multimodal_gen.configs.pipeline_configs.flux import (
+    Flux2PipelineConfig,
+    _unpatchify_latents,
+)
+
+
+@dataclass
+class Flux2FinetunedPipelineConfig(Flux2PipelineConfig):
+    """
+    Pipeline configuration for Flux fine-tuned/distilled models.
+
+    This configuration automatically detects and handles custom VAE architectures
+    (e.g., Flux2TinyAutoEncoder) loaded via HuggingFace's auto_map mechanism.
+
+    Features:
+    - Automatic VAE type detection (standard vs. distilled)
+    - Proper handling of patchified/unpatchified latents
+    - Support for custom scaling factors from fine-tuned models
+    - 5D latents support for both single-frame and multi-frame generation
+    """
+
+    def preprocess_decoding(
+        self, latents: torch.Tensor, server_args=None, vae=None
+    ) -> torch.Tensor:
+        """
+        Preprocess latents before decoding.
+
+        Handles both standard Flux2 VAE and fine-tuned/distilled VAEs:
+        - Standard Flux2 VAE (has bn): needs unpatchify (128 channels -> 32 channels)
+        - Distilled/Finetuned VAE (no bn): keeps patchified latents (128 channels)
+
+        Also handles 5D latents (batch, channels, frames, height, width) by converting
+        to 4D (batch, channels, height, width) for single-frame cases.
+
+        Args:
+            latents: Input latents tensor, can be 4D or 5D
+            server_args: Server arguments (optional, for compatibility)
+            vae: VAE model instance for dynamic type detection
+
+        Returns:
+            Preprocessed latents ready for VAE decoding
+        """
+        # Handle 5D latents (batch, channels, frames, height, width)
+        if latents.ndim == 5:
+            batch_size, channels, frames, height, width = latents.shape
+            if frames == 1:
+                latents = latents.squeeze(2)
+            else:
+                latents = latents.permute(0, 2, 1, 3, 4).contiguous()
+                latents = latents.view(batch_size * frames, channels, height, width)
+
+        if vae is not None and self._check_vae_has_bn(vae):
+            latents = _unpatchify_latents(latents)
+        return latents
+
+    def get_decode_scale_and_shift(self, device, dtype, vae):
+        """
+        Get scale and shift for decoding.
+
+        Dynamically adapts based on VAE type:
+        - Standard Flux2 VAE (has bn): uses BatchNorm statistics
+        - Distilled/Finetuned VAE (no bn): uses scaling_factor from config
+
+        Args:
+            device: Target device for tensors
+            dtype: Target dtype for tensors
+            vae: VAE model instance
+
+        Returns:
+            Tuple of (scaling_factor, shift_factor)
+            - scaling_factor: Tensor or scalar to divide latents by
+            - shift_factor: Tensor or scalar to add to latents (None for distilled VAEs)
+        """
+        vae_arch_config = self.vae_config.arch_config
+
+        if self._check_vae_has_bn(vae):
+            # Standard Flux2 VAE: use BatchNorm statistics
+            latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(device, dtype)
+            latents_bn_std = torch.sqrt(
+                vae.bn.running_var.view(1, -1, 1, 1) + vae_arch_config.batch_norm_eps
+            ).to(device, dtype)
+            return 1 / latents_bn_std, latents_bn_mean
+
+        # Distilled/Finetuned VAE: Flux2TinyAutoEncoder doesn't need external scaling
+        scale = torch.tensor(1.0, device=device, dtype=dtype).view(1, 1, 1, 1)
+        return scale, None

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

@@ -4,10 +4,11 @@
 
 import dataclasses
 import glob
+import importlib.util
 import json
 import os
 import time
-from abc import ABC, abstractmethod
+from abc import ABC
 from collections.abc import Generator, Iterable
 from copy import deepcopy
 from typing import Any, cast
@@ -101,17 +102,13 @@ def load_native(library, component_module_path: str, server_args: ServerArgs):
             revision=server_args.revision,
         )
     elif library == "diffusers":
-        import diffusers
+        from diffusers import AutoModel
 
-        config = get_diffusers_component_config(model_path=component_module_path)
-        class_name = config.pop("_class_name", None)
-        if class_name:
-            cls = getattr(diffusers, class_name)
-            return cls.from_pretrained(
-                component_module_path, revision=server_args.revision, **config
-            )
-        else:
-            raise ValueError("Cannot determine class name for generic diffusers loader")
+        return AutoModel.from_pretrained(
+            component_module_path,
+            revision=server_args.revision,
+            trust_remote_code=server_args.trust_remote_code,
+        )
     else:
         raise ValueError(f"Unsupported library: {library}")
 
@@ -202,13 +199,6 @@ class ComponentLoader(ABC):
             f"load_customized not implemented for {self.__class__.__name__}"
         )
 
-    @abstractmethod
-    def load_customized(
-        self, model_path: str, server_args: ServerArgs, module_name: str
-    ) -> Any:
-        """Implement the minimal core load logic in subclasses."""
-        raise NotImplementedError
-
     @classmethod
     def for_module_type(
         cls, module_type: str, transformers_or_diffusers: str
@@ -503,7 +493,6 @@ class ImageEncoderLoader(TextEncoderLoader):
         encoder_config.update_model_arch(model_config)
 
         # Always start with local device; load_model will adjust for offload if needed
-        should_offload = self.should_offload(server_args)
         # TODO(will): add support for other dtypes
         return self.load_model(
             component_model_path,
@@ -555,7 +544,7 @@ class VAELoader(ComponentLoader):
     ):
         """Load the VAE based on the model path, and inference args."""
         config = get_diffusers_component_config(model_path=component_model_path)
-        class_name = config.pop("_class_name")
+        class_name = config.pop("_class_name", None)
         assert (
             class_name is not None
         ), "Model config does not contain a _class_name attribute. Only diffusers format is supported."
@@ -571,23 +560,42 @@ class VAELoader(ComponentLoader):
 
         target_device = self.target_device(server_args.vae_cpu_offload)
 
-        with set_default_torch_dtype(
-            PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
-        ), skip_init_modules():
+        # Check for auto_map first (custom VAE classes)
+        auto_map = config.get("auto_map", {})
+        auto_model_map = auto_map.get("AutoModel")
+        if auto_model_map:
+            module_path, cls_name = auto_model_map.rsplit(".", 1)
+            custom_module_file = os.path.join(component_model_path, f"{module_path}.py")
+            spec = importlib.util.spec_from_file_location("_custom", custom_module_file)
+            custom_module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(custom_module)
+            vae_cls = getattr(custom_module, cls_name)
+            vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
+            with set_default_torch_dtype(vae_dtype):
+                vae = vae_cls.from_pretrained(
+                    component_model_path,
+                    revision=server_args.revision,
+                    trust_remote_code=server_args.trust_remote_code,
+                )
+            vae = vae.to(device=target_device, dtype=vae_dtype)
+            return vae.eval()
+
+        # Load from ModelRegistry (standard VAE classes)
+        with (
+            set_default_torch_dtype(
+                PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
+            ),
+            skip_init_modules(),
+        ):
             vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
             vae = vae_cls(vae_config).to(target_device)
 
-        # Find all safetensors files
         safetensors_list = _list_safetensors_files(component_model_path)
-        # TODO(PY)
         assert (
             len(safetensors_list) == 1
         ), f"Found {len(safetensors_list)} safetensors files in {component_model_path}"
         loaded = safetensors_load_file(safetensors_list[0])
-        vae.load_state_dict(
-            loaded, strict=False
-        )  # We might only load encoder or decoder
-
+        vae.load_state_dict(loaded, strict=False)
         return vae.eval()
 
 

python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py

@@ -281,7 +281,6 @@ class ComposedPipelineBase(ABC):
             transformers_or_diffusers,
             architecture,
         ) in tqdm(iterable=model_index.items(), desc="Loading required modules"):
-
             if transformers_or_diffusers is None:
                 logger.warning(
                     "Module %s in model_index.json has null value, removing from required_config_modules",
@@ -304,7 +303,19 @@ class ComposedPipelineBase(ABC):
             else:
                 load_module_name = module_name
 
-            component_model_path = os.path.join(self.model_path, load_module_name)
+            # Use custom VAE path if provided, otherwise use default path
+            if module_name == "vae" and server_args.vae_path is not None:
+                component_model_path = server_args.vae_path
+                # Download from HuggingFace Hub if path doesn't exist locally
+                if not os.path.exists(component_model_path):
+                    component_model_path = maybe_download_model(component_model_path)
+                logger.info(
+                    "Using custom VAE path: %s instead of default path: %s",
+                    component_model_path,
+                    os.path.join(self.model_path, load_module_name),
+                )
+            else:
+                component_model_path = os.path.join(self.model_path, load_module_name)
             module = PipelineComponentLoader.load_module(
                 module_name=load_module_name,
                 component_model_path=component_model_path,

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

@@ -27,6 +27,26 @@ from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 logger = init_logger(__name__)
 
 
+def _ensure_tensor_decode_output(decode_output):
+    """
+    Ensure VAE decode output is a tensor.
+
+    Some VAE implementations return DecoderOutput objects with a .sample attribute,
+    tuples, or tensors directly. This function normalizes the output to always be a tensor.
+
+    Args:
+        decode_output: Output from VAE.decode(), can be DecoderOutput, tuple, or torch.Tensor
+
+    Returns:
+        torch.Tensor: The decoded image tensor
+    """
+    if isinstance(decode_output, tuple):
+        return decode_output[0]
+    if hasattr(decode_output, "sample"):
+        return decode_output.sample
+    return decode_output
+
+
 class DecodingStage(PipelineStage):
     """
     Stage for decoding latent representations into pixel space.
@@ -106,7 +126,10 @@ class DecodingStage(PipelineStage):
 
         # scale and shift
         latents = self.scale_and_shift(latents, server_args)
-        latents = server_args.pipeline_config.preprocess_decoding(latents)
+        # Preprocess latents before decoding (e.g., unpatchify for standard Flux2 VAE)
+        latents = server_args.pipeline_config.preprocess_decoding(
+            latents, server_args, vae=self.vae
+        )
 
         # Decode latents
         with torch.autocast(
@@ -120,7 +143,8 @@ class DecodingStage(PipelineStage):
                 pass
             if not vae_autocast_enabled:
                 latents = latents.to(vae_dtype)
-            image = self.vae.decode(latents)
+            decode_output = self.vae.decode(latents)
+            image = _ensure_tensor_decode_output(decode_output)
 
         # De-normalize image to [0, 1] range
         image = (image / 2 + 0.5).clamp(0, 1)

python/sglang/multimodal_gen/runtime/server_args.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py
 """The arguments of sglang-diffusion Inference."""
+
 import argparse
 import dataclasses
 import inspect
@@ -252,6 +253,9 @@ class ServerArgs:
     # (Wenxuan) prefer to keep it here instead of in pipeline config to not make it complicated.
     lora_path: str | None = None
     lora_nickname: str = "default"  # for swapping adapters in the pipeline
+
+    # VAE parameters
+    vae_path: str | None = None  # Custom VAE path (e.g., for distilled autoencoder)
     # can restrict layers to adapt, e.g. ["q_proj"]
     # Will adapt only q, k, v, o by default.
     lora_target_modules: list[str] | None = None
@@ -374,6 +378,12 @@ class ServerArgs:
             type=str,
             help="Directory containing StepVideo model",
         )
+        parser.add_argument(
+            "--vae-path",
+            type=str,
+            default=ServerArgs.vae_path,
+            help="Custom path to VAE model (e.g., for distilled autoencoder). If not specified, VAE will be loaded from the main model path.",
+        )
 
         # attention
         parser.add_argument(
@@ -825,17 +835,15 @@ class ServerArgs:
         if self.ulysses_degree is None:
             self.ulysses_degree = 1
             logger.info(
-                f"Ulysses degree not set, " f"using default value {self.ulysses_degree}"
+                f"Ulysses degree not set, using default value {self.ulysses_degree}"
             )
 
         if self.ring_degree is None:
             self.ring_degree = 1
-            logger.info(
-                f"Ring degree not set, " f"using default value {self.ring_degree}"
-            )
+            logger.info(f"Ring degree not set, using default value {self.ring_degree}")
 
         if self.ring_degree > 1:
-            if self.attention_backend != None and self.attention_backend != "fa":
+            if self.attention_backend is not None and self.attention_backend != "fa":
                 raise ValueError(
                     "Ring Attention is only supported for flash attention backend for now"
                 )

python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py

@@ -135,33 +135,33 @@ def get_diffusers_component_config(
     """Gets a configuration of a submodule for the given diffusers model.
 
     Args:
-        model_path: the path of the submodule
+        model_path: the path of the submodule (can be local path or HuggingFace model ID)
 
     Returns:
         The loaded configuration.
     """
 
-    # Check if the model path exists
-    if os.path.exists(model_path):
-        # tokenizer
-        config_names = ["generation_config.json"]
-        # By default, we load config.json, but scheduler_config.json for scheduler
-        if "scheduler" in model_path:
-            config_names.append("scheduler_config.json")
-        else:
-            config_names.append("config.json")
+    # Download from HuggingFace Hub if path doesn't exist locally
+    if not os.path.exists(model_path):
+        model_path = maybe_download_model(model_path)
 
-        config_file_paths = [
-            os.path.join(model_path, config_name) for config_name in config_names
-        ]
-
-        combined_config = reduce(
-            lambda acc, path: acc | load_dict(path), config_file_paths, {}
-        )
-
-        return combined_config
+    # tokenizer
+    config_names = ["generation_config.json"]
+    # By default, we load config.json, but scheduler_config.json for scheduler
+    if "scheduler" in model_path:
+        config_names.append("scheduler_config.json")
     else:
-        raise RuntimeError(f"Diffusers config file not found at {model_path}")
+        config_names.append("config.json")
+
+    config_file_paths = [
+        os.path.join(model_path, config_name) for config_name in config_names
+    ]
+
+    combined_config = reduce(
+        lambda acc, path: acc | load_dict(path), config_file_paths, {}
+    )
+
+    return combined_config
 
 
 # Models don't use the same configuration key for determining the maximum
@@ -390,9 +390,10 @@ def maybe_download_model(
         logger.info(
             "Downloading model snapshot from HF Hub for %s...", model_name_or_path
         )
-        with get_lock(model_name_or_path).acquire(
-            poll_interval=2
-        ), suppress_other_loggers(not_suppress_on_main_rank=True):
+        with (
+            get_lock(model_name_or_path).acquire(poll_interval=2),
+            suppress_other_loggers(not_suppress_on_main_rank=True),
+        ):
             local_path = snapshot_download(
                 repo_id=model_name_or_path,
                 ignore_patterns=["*.onnx", "*.msgpack"],