diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index a5a98981b..80ee489ad 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/python/sglang/multimodal_gen/envs.py @@ -55,6 +55,7 @@ if TYPE_CHECKING: SGLANG_CACHE_DIT_SECONDARY_TS_ORDER: int = 1 # model loading SGLANG_USE_RUNAI_MODEL_STREAMER: bool = True + SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D: bool = False def get_default_cache_root() -> str: @@ -243,6 +244,9 @@ environment_variables: dict[str, Callable[[], Any]] = { # If set, sgl_diffusion will enable stage logging, which will print the time # taken for each stage "SGLANG_DIFFUSION_STAGE_LOGGING": _lazy_bool("SGLANG_DIFFUSION_STAGE_LOGGING"), + "SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D": _lazy_bool( + "SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D", "false" + ), # ================== cache-dit Env Vars ================== # Enable cache-dit acceleration for DiT inference "SGLANG_CACHE_DIT_ENABLED": _lazy_bool("SGLANG_CACHE_DIT_ENABLED"), diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 9c32f143c..47770a142 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py @@ -1,8 +1,11 @@ import importlib.util import os +import torch +import torch.nn as nn from safetensors.torch import load_file as safetensors_load_file +from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.models import ModelConfig from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( ComponentLoader, @@ -23,6 +26,25 @@ from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) +def _convert_conv3d_weights_to_channels_last_3d(module: nn.Module) -> int: + """ + Convert Conv3d weights to channels_last_3d (NDHWC) memory format. + Returns the number of Conv3d modules converted. + """ + if not hasattr(torch, "channels_last_3d"): + return 0 + num_converted = 0 + for m in module.modules(): + if isinstance(m, nn.Conv3d): + try: + m.weight.data = m.weight.data.to(memory_format=torch.channels_last_3d) + num_converted += 1 + except Exception: + # Best-effort; skip unsupported cases. + continue + return num_converted + + class VAELoader(ComponentLoader): """Shared loader for (video/audio) VAE modules.""" @@ -85,6 +107,16 @@ class VAELoader(ComponentLoader): trust_remote_code=server_args.trust_remote_code, ) vae = vae.to(device=target_device, dtype=vae_dtype) + if ( + component_name in ("vae", "video_vae") + and torch.cuda.is_available() + and getattr(envs, "SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D", False) + ): + n = _convert_conv3d_weights_to_channels_last_3d(vae) + if n > 0: + logger.info( + "VAE: converted %d Conv3d weights to channels_last_3d", n + ) return vae # Load from ModelRegistry (standard VAE classes) @@ -111,4 +143,13 @@ class VAELoader(ComponentLoader): if unexpected_keys: logger.warning("VAE unexpected keys: %s", unexpected_keys) + if ( + component_name in ("vae", "video_vae") + and torch.cuda.is_available() + and getattr(envs, "SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D", False) + ): + n = _convert_conv3d_weights_to_channels_last_3d(vae) + if n > 0: + logger.info("VAE: converted %d Conv3d weights to channels_last_3d", n) + return vae