[diffusion] feat: support distilled vae generic (#14195)
Co-authored-by: BBuf <1182563586@qq.com> Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Dongjie Zou
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user