From 09a9147f5938b12ff1c2e0a8797be9d8b11d67dd Mon Sep 17 00:00:00 2001 From: "Zhang Yiyang (SII)" <30491374+CloudRipple@users.noreply.github.com> Date: Thu, 29 Jan 2026 09:12:08 +0800 Subject: [PATCH] [diffusion] model: support MOVA (#17704) Co-authored-by: gaoyang07 Co-authored-by: cms42 Co-authored-by: cms42 <44895820+cms42@users.noreply.github.com> Co-authored-by: Ruixiao Li Co-authored-by: Li Ruixiao(SII) <80368770+Li-dongyang@users.noreply.github.com> --- .../configs/models/bridges/__init__.py | 7 + .../configs/models/bridges/mova_dual_tower.py | 42 + .../configs/models/dits/__init__.py | 9 +- .../configs/models/dits/mova_audio.py | 66 ++ .../configs/models/dits/mova_video.py | 65 ++ .../configs/models/vaes/__init__.py | 2 + .../multimodal_gen/configs/models/vaes/dac.py | 30 + .../configs/pipeline_configs/__init__.py | 2 + .../configs/pipeline_configs/mova.py | 185 ++++ .../multimodal_gen/configs/sample/mova.py | 57 ++ python/sglang/multimodal_gen/registry.py | 23 + .../runtime/loader/component_loader.py | 152 ++- .../runtime/managers/gpu_worker.py | 5 + .../runtime/models/bridges/__init__.py | 7 + .../runtime/models/bridges/mova_dual_tower.py | 671 +++++++++++++ .../runtime/models/dits/mova_audio_dit.py | 306 ++++++ .../runtime/models/dits/mova_video_dit.py | 570 +++++++++++ .../runtime/models/model_stages/mova.py | 911 ++++++++++++++++++ .../models/schedulers/flow_match_pair.py | 501 ++++++++++ .../runtime/models/vaes/common.py | 8 +- .../multimodal_gen/runtime/models/vaes/dac.py | 627 ++++++++++++ .../runtime/pipelines/mova_pipeline.py | 119 +++ .../pipelines_core/composed_pipeline_base.py | 21 +- .../runtime/pipelines_core/schedule_batch.py | 7 +- .../pipelines_core/stages/input_validation.py | 56 +- .../multimodal_gen/runtime/server_args.py | 5 + .../runtime/utils/hf_diffusers_utils.py | 80 +- 27 files changed, 4485 insertions(+), 49 deletions(-) create mode 100644 python/sglang/multimodal_gen/configs/models/bridges/__init__.py create mode 100644 python/sglang/multimodal_gen/configs/models/bridges/mova_dual_tower.py create mode 100644 python/sglang/multimodal_gen/configs/models/dits/mova_audio.py create mode 100644 python/sglang/multimodal_gen/configs/models/dits/mova_video.py create mode 100644 python/sglang/multimodal_gen/configs/models/vaes/dac.py create mode 100644 python/sglang/multimodal_gen/configs/pipeline_configs/mova.py create mode 100644 python/sglang/multimodal_gen/configs/sample/mova.py create mode 100644 python/sglang/multimodal_gen/runtime/models/bridges/__init__.py create mode 100644 python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py create mode 100644 python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py create mode 100644 python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py create mode 100644 python/sglang/multimodal_gen/runtime/models/model_stages/mova.py create mode 100644 python/sglang/multimodal_gen/runtime/models/schedulers/flow_match_pair.py create mode 100644 python/sglang/multimodal_gen/runtime/models/vaes/dac.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines/mova_pipeline.py diff --git a/python/sglang/multimodal_gen/configs/models/bridges/__init__.py b/python/sglang/multimodal_gen/configs/models/bridges/__init__.py new file mode 100644 index 000000000..8e860b890 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/bridges/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 + +from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import ( + MOVADualTowerConfig, +) + +__all__ = ["MOVADualTowerConfig"] diff --git a/python/sglang/multimodal_gen/configs/models/bridges/mova_dual_tower.py b/python/sglang/multimodal_gen/configs/models/bridges/mova_dual_tower.py new file mode 100644 index 000000000..4bf0b83c4 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/bridges/mova_dual_tower.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Configuration for MOVA dual tower bridge model.""" + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig + + +def _is_conditioner_block(name: str, module) -> bool: + """Check if module is a ConditionalCrossAttentionBlock.""" + return "ConditionalCrossAttentionBlock" in type(module).__name__ + + +@dataclass +class MOVADualTowerArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field( + default_factory=lambda: [_is_conditioner_block] + ) + + # Model architecture parameters + visual_layers: int = 40 + audio_layers: int = 30 + visual_hidden_dim: int = 5120 + audio_hidden_dim: int = 1536 + audio_fps: float = 50.0 + head_dim: int = 128 + interaction_strategy: str = "full" + apply_cross_rope: bool = True + apply_first_frame_bias_in_rope: bool = False + trainable_condition_scale: bool = False + pooled_adaln: bool = False + eps: float = 1e-6 + + def __post_init__(self): + super().__post_init__() + self.hidden_size = self.visual_hidden_dim + self.num_attention_heads = self.visual_hidden_dim // self.head_dim + + +@dataclass +class MOVADualTowerConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=MOVADualTowerArchConfig) diff --git a/python/sglang/multimodal_gen/configs/models/dits/__init__.py b/python/sglang/multimodal_gen/configs/models/dits/__init__.py index 30b205451..1e3057352 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/dits/__init__.py @@ -1,6 +1,13 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from sglang.multimodal_gen.configs.models.dits.hunyuanvideo import HunyuanVideoConfig +from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig +from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig from sglang.multimodal_gen.configs.models.dits.wanvideo import WanVideoConfig -__all__ = ["HunyuanVideoConfig", "WanVideoConfig"] +__all__ = [ + "HunyuanVideoConfig", + "MOVAAudioConfig", + "MOVAVideoConfig", + "WanVideoConfig", +] diff --git a/python/sglang/multimodal_gen/configs/models/dits/mova_audio.py b/python/sglang/multimodal_gen/configs/models/dits/mova_audio.py new file mode 100644 index 000000000..a0c58b4ec --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/mova_audio.py @@ -0,0 +1,66 @@ +# Copied and adapted from: mossVG/mova/diffusion/models/wan_audio_dit.py +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig + + +def _is_blocks(n: str, m) -> bool: + return "blocks" in n and str.isdigit(n.split(".")[-1]) + + +@dataclass +class MOVAAudioArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_blocks]) + + param_names_mapping: dict = field( + default_factory=lambda: { + r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2", + r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", + r"^text_embedding\.0\.(.*)$": r"text_embedding.fc_in.\1", + r"^text_embedding\.2\.(.*)$": r"text_embedding.fc_out.\1", + r"^time_embedding\.0\.(.*)$": r"time_embedding.fc_in.\1", + r"^time_embedding\.2\.(.*)$": r"time_embedding.fc_out.\1", + r"^img_emb\.proj\.1\.(.*)$": r"img_emb.fc_in.\1", + r"^img_emb\.proj\.3\.(.*)$": r"img_emb.fc_out.\1", + } + ) + reverse_param_names_mapping: dict = field(default_factory=dict) + lora_param_names_mapping: dict = field(default_factory=dict) + + dim: int = 1536 + in_dim: int = 128 + ffn_dim: int = 6144 + out_dim: int = 128 + text_dim: int = 4096 + freq_dim: int = 256 + eps: float = 1e-6 + patch_size: tuple[int, int, int] = (1, 2, 2) + num_heads: int = 12 + num_layers: int = 30 + has_image_input: bool = False + has_image_pos_emb: bool = False + has_ref_conv: bool = False + add_control_adapter: bool = False + in_dim_control_adapter: int = 24 + seperated_timestep: bool = False + require_vae_embedding: bool = False + require_clip_embedding: bool = False + fuse_vae_embedding_in_latents: bool = False + vae_type: str = "dac" + + def __post_init__(self): + super().__post_init__() + self.hidden_size = self.dim + self.num_attention_heads = self.num_heads + self.num_channels_latents = self.out_dim + assert ( + not self.has_image_input + ), "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)." + + +@dataclass +class MOVAAudioConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=MOVAAudioArchConfig) + prefix: str = "mova_audio" diff --git a/python/sglang/multimodal_gen/configs/models/dits/mova_video.py b/python/sglang/multimodal_gen/configs/models/dits/mova_video.py new file mode 100644 index 000000000..0ec03932a --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/mova_video.py @@ -0,0 +1,65 @@ +# Copied and adapted from: mossVG/mova/diffusion/models/wan_video_dit.py +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig + + +def _is_blocks(n: str, m) -> bool: + return "blocks" in n and str.isdigit(n.split(".")[-1]) + + +@dataclass +class MOVAVideoArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_blocks]) + + param_names_mapping: dict = field( + default_factory=lambda: { + r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2", + r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", + r"^text_embedding\.0\.(.*)$": r"text_embedding.fc_in.\1", + r"^text_embedding\.2\.(.*)$": r"text_embedding.fc_out.\1", + r"^time_embedding\.0\.(.*)$": r"time_embedding.fc_in.\1", + r"^time_embedding\.2\.(.*)$": r"time_embedding.fc_out.\1", + r"^img_emb\.proj\.1\.(.*)$": r"img_emb.fc_in.\1", + r"^img_emb\.proj\.3\.(.*)$": r"img_emb.fc_out.\1", + } + ) + reverse_param_names_mapping: dict = field(default_factory=dict) + lora_param_names_mapping: dict = field(default_factory=dict) + + dim: int = 5120 + in_dim: int = 16 + ffn_dim: int = 13824 + out_dim: int = 16 + text_dim: int = 4096 + freq_dim: int = 256 + eps: float = 1e-6 + patch_size: tuple[int, int, int] = (1, 2, 2) + num_heads: int = 40 + num_layers: int = 40 + has_image_input: bool = False + has_image_pos_emb: bool = False + has_ref_conv: bool = False + add_control_adapter: bool = False + in_dim_control_adapter: int = 24 + seperated_timestep: bool = False + require_vae_embedding: bool = True + require_clip_embedding: bool = True + fuse_vae_embedding_in_latents: bool = False + + def __post_init__(self): + super().__post_init__() + self.hidden_size = self.dim + self.num_attention_heads = self.num_heads + self.num_channels_latents = self.out_dim + assert ( + not self.has_image_input + ), "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)." + + +@dataclass +class MOVAVideoConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=MOVAVideoArchConfig) + prefix: str = "mova_video" diff --git a/python/sglang/multimodal_gen/configs/models/vaes/__init__.py b/python/sglang/multimodal_gen/configs/models/vaes/__init__.py index 1d6e7461e..d793a6a9d 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/__init__.py @@ -1,9 +1,11 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo +from sglang.multimodal_gen.configs.models.vaes.dac import DacVAEConfig from sglang.multimodal_gen.configs.models.vaes.hunyuanvae import HunyuanVAEConfig from sglang.multimodal_gen.configs.models.vaes.wanvae import WanVAEConfig __all__ = [ + "DacVAEConfig", "HunyuanVAEConfig", "WanVAEConfig", ] diff --git a/python/sglang/multimodal_gen/configs/models/vaes/dac.py b/python/sglang/multimodal_gen/configs/models/vaes/dac.py new file mode 100644 index 000000000..63f59c6a5 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/vaes/dac.py @@ -0,0 +1,30 @@ +# Copied and adapted from: mossVG/mova/diffusion/models/dac_vae.py +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field +from typing import List + +from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig + + +@dataclass +class DacVAEArchConfig(ArchConfig): + codebook_dim: int = 8 + codebook_size: int = 1024 + continuous: bool = True + decoder_dim: int = 2048 + decoder_rates: List[int] = field(default_factory=lambda: [8, 5, 4, 3, 2]) + encoder_dim: int = 128 + encoder_rates: List[int] = field(default_factory=lambda: [2, 3, 4, 5, 8]) + hop_length: int = 3840 + latent_dim: int = 128 + n_codebooks: int = 9 + quantizer_dropout: bool = False + sample_rate: int = 48000 + + +@dataclass +class DacVAEConfig(ModelConfig): + arch_config: DacVAEArchConfig = field(default_factory=DacVAEArchConfig) + load_encoder: bool = True + load_decoder: bool = True diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py index ebb120938..534caff44 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py @@ -20,6 +20,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import ( HunyuanConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig +from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.wan import ( SelfForcingWanT2V480PConfig, WanI2V480PConfig, @@ -39,6 +40,7 @@ __all__ = [ "Flux2FinetunedPipelineConfig", "PipelineConfig", "SlidingTileAttnConfig", + "MOVAPipelineConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/mova.py b/python/sglang/multimodal_gen/configs/pipeline_configs/mova.py new file mode 100644 index 000000000..9e0131dd5 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/mova.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +MOVA pipeline configuration. +""" + +from dataclasses import dataclass, field + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +from sglang.multimodal_gen.configs.models.dits import MOVAAudioConfig, MOVAVideoConfig +from sglang.multimodal_gen.configs.models.encoders import T5Config +from sglang.multimodal_gen.configs.models.vaes import DacVAEConfig, WanVAEConfig +from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ModelTaskType, + PipelineConfig, +) +from sglang.multimodal_gen.configs.pipeline_configs.wan import t5_postprocess_text +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +@dataclass +class MOVAPipelineConfig(PipelineConfig): + """Configuration for MOVA (text+image -> video+audio) pipelines.""" + + task_type: ModelTaskType = ModelTaskType.T2V + + # Model configs + dit_config: MOVAVideoConfig = field(default_factory=MOVAVideoConfig) + audio_dit_config: MOVAAudioConfig = field(default_factory=MOVAAudioConfig) + + # Video VAE (Wan) + Audio VAE (DAC) + vae_config: WanVAEConfig = field(default_factory=WanVAEConfig) + audio_vae_config: DacVAEConfig = field(default_factory=DacVAEConfig) + audio_vae_precision: str = "fp32" + + # Text encoder (UMT5 compatible) + text_encoder_configs: tuple = field(default_factory=lambda: (T5Config(),)) + postprocess_text_funcs: tuple = field( + default_factory=lambda: (t5_postprocess_text,) + ) + + # MOVA specific + audio_vae_type: str = "dac" + boundary_ratio: float | None = 0.9 + + # temporal alignment: MOVA expects (num_frames - 1) % 4 == 0 + time_division_factor: int = 4 + time_division_remainder: int = 1 + + def _center_crop_and_resize( + self, image: torch.Tensor | Image.Image, target_height: int, target_width: int + ) -> torch.Tensor | Image.Image: + if not isinstance(image, (Image.Image, torch.Tensor)): + raise TypeError(f"Unsupported image type: {type(image)}") + if isinstance(image, Image.Image): + image = torch.from_numpy(np.array(image)) + + if image.ndim == 2: + image = image[..., None] + + if not image.dtype.is_floating_point: + image = image.to(torch.float32).div(255.0) + + if image.ndim == 3: + if image.shape[0] in (1, 3, 4) and image.shape[-1] not in (1, 3, 4): + image = image.unsqueeze(0) + else: + image = image.permute(2, 0, 1).unsqueeze(0) + elif image.ndim == 4: + if image.shape[1] not in (1, 3, 4) and image.shape[-1] in (1, 3, 4): + image = image.permute(0, 3, 1, 2) + + image_height, image_width = image.shape[-2], image.shape[-1] + if image_height == target_height and image_width == target_width: + return image + + logger.info( + "Center cropping and resizing image to %dx%d", target_width, target_height + ) + + if image_height * target_width < image_width * target_height: + cropped_width = (image_height * target_width) // target_height + left = (image_width - cropped_width) // 2 + image = image[..., :, left : left + cropped_width] + else: + cropped_height = (image_width * target_height) // target_width + top = (image_height - cropped_height) // 2 + image = image[..., top : top + cropped_height, :] + + image = F.interpolate( + image, + size=(target_height, target_width), + mode="bilinear", + align_corners=False, + antialias=True, + ) + return image + + def adjust_num_frames(self, num_frames: int) -> int: + if num_frames is None: + return num_frames + if num_frames % self.time_division_factor != self.time_division_remainder: + adjusted = ( + (num_frames + self.time_division_factor - 1) + // self.time_division_factor + * self.time_division_factor + + self.time_division_remainder + ) + logger.warning( + "`num_frames` (%s) is not compatible with MOVA temporal constraints. " + "Rounding to %s.", + num_frames, + adjusted, + ) + return adjusted + return num_frames + + def preprocess_condition_image( + self, image, target_width, target_height, _vae_image_processor + ): + image = self._center_crop_and_resize(image, target_height, target_width) + return image, (target_width, target_height) + + def prepare_latent_shape(self, batch, batch_size, num_frames): + spatial = self.vae_config.arch_config.spatial_compression_ratio + length = (num_frames - 1) // self.time_division_factor + 1 + shape = ( + batch_size, + self.dit_config.arch_config.out_dim, + length, + batch.height // spatial, + batch.width // spatial, + ) + return shape + + def prepare_audio_latent_shape(self, batch_size, num_samples, audio_vae): + latent_T = (num_samples + audio_vae.hop_length - 1) // audio_vae.hop_length + return (batch_size, audio_vae.latent_dim, latent_T) + + def normalize_video_latents(self, latents: torch.Tensor, video_vae) -> torch.Tensor: + latents_mean = getattr(video_vae.config, "latents_mean", None) + latents_std = getattr(video_vae.config, "latents_std", None) + if latents_mean is None or latents_std is None: + return latents + mean = torch.tensor( + latents_mean, device=latents.device, dtype=latents.dtype + ).view(1, video_vae.config.z_dim, 1, 1, 1) + inv_std = ( + 1.0 / torch.tensor(latents_std, device=latents.device, dtype=latents.dtype) + ).view(1, video_vae.config.z_dim, 1, 1, 1) + return (latents - mean) * inv_std + + def denormalize_video_latents( + self, latents: torch.Tensor, video_vae + ) -> torch.Tensor: + latents_mean = getattr(video_vae.config, "latents_mean", None) + latents_std = getattr(video_vae.config, "latents_std", None) + if latents_mean is None or latents_std is None: + return latents + mean = torch.tensor( + latents_mean, device=latents.device, dtype=latents.dtype + ).view(1, video_vae.config.z_dim, 1, 1, 1) + std = torch.tensor( + latents_std, device=latents.device, dtype=latents.dtype + ).view(1, video_vae.config.z_dim, 1, 1, 1) + return latents * std + mean + + +@dataclass +class MOVA360PConfig(MOVAPipelineConfig): + """Configuration for MOVA 360P (text+image -> video+audio) pipelines.""" + + max_area: int = 352 * 640 + + +@dataclass +class MOVA720PConfig(MOVAPipelineConfig): + """Configuration for MOVA 720P (text+image -> video+audio) pipelines.""" + + max_area: int = 720 * 1280 diff --git a/python/sglang/multimodal_gen/configs/sample/mova.py b/python/sglang/multimodal_gen/configs/sample/mova.py new file mode 100644 index 000000000..e27311abb --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/mova.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams + + +@dataclass +class MOVASamplingParams(SamplingParams): + # Video parameters (MOVA defaults) + height: int = 352 + width: int = 640 + num_frames: int = 193 + fps: int = 24 + + # Denoising stage + guidance_scale: float = 5.0 + num_inference_steps: int = 50 + sigma_shift: float = 5.0 + visual_shift: float = 5.0 + audio_shift: float = 5.0 + + negative_prompt: str = ( + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，" + "整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，" + "画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，" + "静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" + ) + + +@dataclass +class MOVA_360P_SamplingParams(MOVASamplingParams): + # Video parameters (MOVA 360P) + height: int = 352 + width: int = 640 + + # MOVA 360P supported resolutions + supported_resolutions: list[tuple[int, int]] = field( + default_factory=lambda: [ + (352, 640), + (640, 352), + ] + ) + + +@dataclass +class MOVA_720P_SamplingParams(MOVASamplingParams): + # Video parameters (MOVA 720P) + height: int = 720 + width: int = 1280 + + # MOVA 720P supported resolutions + supported_resolutions: list[tuple[int, int]] = field( + default_factory=lambda: [ + (720, 1280), + (1280, 720), + ] + ) diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index be9032f28..6d239c8f9 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -46,6 +46,10 @@ from sglang.multimodal_gen.configs.pipeline_configs.glm_image import ( GlmImagePipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig +from sglang.multimodal_gen.configs.pipeline_configs.mova import ( + MOVA360PConfig, + MOVA720PConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImageEditPipelineConfig, QwenImageEditPlus_2511_PipelineConfig, @@ -72,6 +76,10 @@ from sglang.multimodal_gen.configs.sample.hunyuan import ( HunyuanSamplingParams, ) from sglang.multimodal_gen.configs.sample.ltx_2 import LTX2SamplingParams +from sglang.multimodal_gen.configs.sample.mova import ( + MOVA_360P_SamplingParams, + MOVA_720P_SamplingParams, +) from sglang.multimodal_gen.configs.sample.qwenimage import ( QwenImage2512SamplingParams, QwenImageEditPlusSamplingParams, @@ -543,6 +551,21 @@ def _register_configs(): "FastVideo/FastWan2.1-T2V-1.3B-Diffusers", ], ) + # MOVA + register_configs( + sampling_param_cls=MOVA_360P_SamplingParams, + pipeline_config_cls=MOVA360PConfig, + model_detectors=[ + lambda hf_id: "mova" in hf_id.lower() and "360p" in hf_id.lower() + ], + ) + register_configs( + sampling_param_cls=MOVA_720P_SamplingParams, + pipeline_config_cls=MOVA720PConfig, + model_detectors=[ + lambda hf_id: "mova" in hf_id.lower() and "720p" in hf_id.lower() + ], + ) # FLUX register_configs( sampling_param_cls=FluxSamplingParams, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index c21197252..71a710075 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loader.py @@ -268,18 +268,31 @@ class ComponentLoader(ABC): "processor": (AutoProcessorLoader, "transformers"), "vision_language_encoder": (VisionLanguageEncoderLoader, "transformers"), } - # Loaders for audio/video specific components that might vary av_module_loaders = { + "audio_dit": (TransformerLoader, "diffusers"), "audio_vae": (VAELoader, "diffusers"), - "vocoder": (VocoderLoader, "diffusers"), "connectors": (AdapterLoader, "diffusers"), + "dual_tower_bridge": (BridgeLoader, "diffusers"), + "video_dit": (TransformerLoader, "diffusers"), + "video_vae": (VAELoader, "diffusers"), + "vocoder": (VocoderLoader, "diffusers"), } # NOTE(FlamingoPg): special for LTX-2 models if module_type == "vocoder" or module_type == "connectors": transformers_or_diffusers = "diffusers" + # NOTE(CloudRipple): special for MOVA models + # TODO(CloudRipple): remove most of these special cases after unifying the loading logic + if module_type in [ + "audio_vae", + "audio_dit", + "dual_tower_bridge", + "video_dit", + ]: + transformers_or_diffusers = "diffusers" + if module_type in module_loaders: loader_cls, expected_library = module_loaders[module_type] # Assert that the library matches what's expected for this module type @@ -616,7 +629,7 @@ class TokenizerLoader(ComponentLoader): class VAELoader(ComponentLoader): - """Loader for VAE.""" + """Shared loader for (video/audio) VAE modules.""" def should_offload( self, server_args: ServerArgs, model_config: ModelConfig | None = None @@ -636,17 +649,20 @@ class VAELoader(ComponentLoader): server_args.model_paths[module_name] = component_model_path logger.debug("HF model config: %s", config) - if module_name == "audio_vae": - vae_config = server_args.pipeline_config.audio_vae_config - vae_precision = server_args.pipeline_config.audio_vae_precision + if module_name in ("vae", "video_vae"): + pipeline_vae_config_attr = "vae_config" + pipeline_vae_precision = "vae_precision" + elif module_name in ("audio_vae",): + pipeline_vae_config_attr = "audio_vae_config" + pipeline_vae_precision = "audio_vae_precision" else: - vae_config = server_args.pipeline_config.vae_config - vae_precision = server_args.pipeline_config.vae_precision - + raise ValueError(f"Unsupported module name for VAE loader: {module_name}") + vae_config = getattr(server_args.pipeline_config, pipeline_vae_config_attr) + vae_precision = getattr(server_args.pipeline_config, pipeline_vae_precision) vae_config.update_model_arch(config) - - # NOTE: some post init logics are only available after updated with config - vae_config.post_init() + if hasattr(vae_config, "post_init"): + # NOTE: some post init logics are only available after updated with config + vae_config.post_init() should_offload = self.should_offload(server_args) target_device = self.target_device(should_offload) @@ -685,6 +701,16 @@ class VAELoader(ComponentLoader): ), 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) + + state_keys = set(vae.state_dict().keys()) + loaded_keys = set(loaded.keys()) + missing_keys = sorted(state_keys - loaded_keys) + unexpected_keys = sorted(loaded_keys - state_keys) + if missing_keys: + logger.warning("VAE missing keys: %s", missing_keys) + if unexpected_keys: + logger.warning("VAE unexpected keys: %s", unexpected_keys) + return vae.eval() @@ -753,11 +779,96 @@ class VocoderLoader(ComponentLoader): return vocoder.eval() -class TransformerLoader(ComponentLoader): - """Loader for transformer.""" +class BridgeLoader(ComponentLoader): + """Loader for MOVA dual tower bridge with FSDP support.""" + + pipeline_bridge_config_attr: str = "bridge_config" def load_customized( - self, component_model_path: str, server_args: ServerArgs, *args + self, component_model_path: str, server_args: ServerArgs, module_name: str + ): + config = get_diffusers_component_config(model_path=component_model_path) + hf_config = deepcopy(config) + class_name = config.pop("_class_name", None) + if class_name is None: + raise ValueError( + "Model config does not contain a _class_name attribute. " + "Only diffusers format is supported." + ) + server_args.model_paths[module_name] = component_model_path + + # Try to get bridge config from pipeline config, fallback to creating one + bridge_config = getattr( + server_args.pipeline_config, self.pipeline_bridge_config_attr, None + ) + if bridge_config is not None: + bridge_config.update_model_arch(config) + else: + # Create a minimal config from hf_config + from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import ( + MOVADualTowerConfig, + ) + + bridge_config = MOVADualTowerConfig() + bridge_config.update_model_arch(config) + + model_cls, _ = ModelRegistry.resolve_model_cls(class_name) + + # Find all safetensors files + safetensors_list = _list_safetensors_files(component_model_path) + if not safetensors_list: + raise ValueError(f"No safetensors files found in {component_model_path}") + + default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision] + + logger.info( + "Loading %s from %s safetensors files, default_dtype: %s", + class_name, + len(safetensors_list), + default_dtype, + ) + + # Check if FSDP loading is available + if ( + server_args.hsdp_shard_dim is not None + and hasattr(model_cls, "_fsdp_shard_conditions") + and model_cls._fsdp_shard_conditions + ): + # Load with FSDP support + model = maybe_load_fsdp_model( + model_cls=model_cls, + init_params={"config": bridge_config, "hf_config": hf_config}, + weight_dir_list=safetensors_list, + device=get_local_torch_device(), + hsdp_replicate_dim=server_args.hsdp_replicate_dim, + hsdp_shard_dim=server_args.hsdp_shard_dim, + cpu_offload=server_args.dit_cpu_offload, + pin_cpu_memory=server_args.pin_cpu_memory, + fsdp_inference=server_args.use_fsdp_inference, + default_dtype=default_dtype, + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + output_dtype=None, + strict=False, + ) + else: + # Fallback to simple loading (for non-FSDP or legacy models) + model = model_cls.from_pretrained( + component_model_path, torch_dtype=default_dtype + ) + model = model.to(device=get_local_torch_device(), dtype=default_dtype) + + total_params = sum(p.numel() for p in model.parameters()) + logger.info("Loaded bridge model with %.2fM parameters", total_params / 1e6) + + return model.eval() + + +class TransformerLoader(ComponentLoader): + """Shared loader for (video/audio) DiT transformers.""" + + def load_customized( + self, component_model_path: str, server_args: ServerArgs, module_name: str ): """Load the transformer based on the model path, and inference args.""" config = get_diffusers_component_config(model_path=component_model_path) @@ -769,10 +880,17 @@ class TransformerLoader(ComponentLoader): "Only diffusers format is supported." ) - server_args.model_paths["transformer"] = component_model_path + module_name = _normalize_module_type(module_name) + server_args.model_paths[module_name] = component_model_path + if module_name in ("transformer", "video_dit"): + pipeline_dit_config_attr = "dit_config" + elif module_name in ("audio_dit",): + pipeline_dit_config_attr = "audio_dit_config" + else: + raise ValueError(f"Invalid module name: {module_name}") # Config from Diffusers supersedes sgl_diffusion's model config - dit_config = server_args.pipeline_config.dit_config + dit_config = getattr(server_args.pipeline_config, pipeline_dit_config_attr) dit_config.update_model_arch(config) model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 3f5a15827..945134523 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -98,6 +98,9 @@ class GPUWorker: [ self.pipeline.get_module("transformer"), self.pipeline.get_module("transformer_2"), + self.pipeline.get_module("video_dit"), + self.pipeline.get_module("video_dit_2"), + self.pipeline.get_module("audio_dit"), ], ): if isinstance(dit, OffloadableDiTMixin): @@ -164,6 +167,8 @@ class GPUWorker: if isinstance(result, Req): output_batch = OutputBatch( output=result.output, + audio=getattr(result, "audio", None), + audio_sample_rate=getattr(result, "audio_sample_rate", None), timings=result.timings, trajectory_timesteps=getattr(result, "trajectory_timesteps", None), trajectory_latents=getattr(result, "trajectory_latents", None), diff --git a/python/sglang/multimodal_gen/runtime/models/bridges/__init__.py b/python/sglang/multimodal_gen/runtime/models/bridges/__init__.py new file mode 100644 index 000000000..bee1e907a --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/bridges/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 + +from sglang.multimodal_gen.runtime.models.bridges.mova_dual_tower import ( + DualTowerConditionalBridge, +) + +__all__ = ["DualTowerConditionalBridge"] diff --git a/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py b/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py new file mode 100644 index 000000000..cbc73d6e8 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py @@ -0,0 +1,671 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copied and adapted from: mossVG/mova/diffusion/models/interactionv2.py + + +from typing import Any, Dict, List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import ( + MOVADualTowerConfig, +) +from sglang.multimodal_gen.runtime.distributed import get_tp_world_size +from sglang.multimodal_gen.runtime.layers.attention import USPAttention +from sglang.multimodal_gen.runtime.layers.layernorm import ( + RMSNorm, + tensor_parallel_rms_norm, +) +from sglang.multimodal_gen.runtime.layers.linear import ( + ColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( + apply_flashinfer_rope_qk_inplace, +) +from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT +from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +@torch.no_grad() +def compute_rope_cos_sin( + position_ids: torch.Tensor, + head_dim: int, + base: float = 10000.0, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Compute RoPE cos/sin embeddings for given position IDs. + + This is a functional implementation that doesn't require storing buffers, + making it compatible with FSDP meta device initialization. + + Args: + position_ids: Position IDs tensor [B, L] or [1, L] + head_dim: Dimension of each attention head + base: RoPE base frequency (default: 10000.0) + device: Target device + dtype: Output dtype + + Returns: + (cos, sin): Each with shape [B, L, head_dim] + """ + device = device or position_ids.device + dtype = dtype or torch.float32 + + # Compute inverse frequencies + inv_freq = 1.0 / ( + base + ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim) + ) + + # Expand for batch computation: [B, L] -> [B, 1, L] @ [1, head_dim/2, 1] -> [B, head_dim/2, L] + inv_freq_expanded = inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + + # Compute frequencies: [B, head_dim/2, L] -> [B, L, head_dim/2] + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + + # Double the frequencies for full head_dim: [B, L, head_dim] + emb = torch.cat((freqs, freqs), dim=-1) + + cos = emb.cos().to(dtype=dtype) + sin = emb.sin().to(dtype=dtype) + + return cos, sin + + +class PerFrameAttentionPooling(nn.Module): + """Per-frame multi-head attention pooling. + + Flattens the input sequence [B, L, D] and grid size (T, H, W). + Performs single-query attention pooling on the H*W tokens for each time frame. + Output shape: [B, T, D]. + """ + + def __init__(self, dim: int, num_heads: int, eps: float = 1e-6): + super().__init__() + assert dim % num_heads == 0, "dim must be divisible by num_heads" + self.dim = dim + self.num_heads = num_heads + + self.probe = nn.Parameter(torch.randn(1, 1, dim)) + nn.init.normal_(self.probe, std=0.02) + + self.attention = nn.MultiheadAttention( + embed_dim=dim, num_heads=num_heads, batch_first=True + ) + self.layernorm = nn.LayerNorm(dim, eps=eps) + + def forward(self, x: torch.Tensor, grid_size: Tuple[int, int, int]) -> torch.Tensor: + """Forward pass. + + Args: + x: Input tensor of shape [B, L, D], where L = T * H * W. + grid_size: Tuple of (T, H, W). + + Returns: + Pooled tensor of shape [B, T, D]. + """ + B, L, D = x.shape + T, H, W = grid_size + assert ( + D == self.dim + ), f"Input dimension D={D} does not match module dim={self.dim}" + assert L == T * H * W, f"Flattened length L={L} does not match T*H*W={T*H*W}" + + S = H * W + x_bt_s_d = x.view(B, T, S, D).contiguous().view(B * T, S, D) # [B*T, S, D] + probe = self.probe.expand(B * T, -1, -1) # [B*T, 1, D] + + pooled_bt_1_d = self.attention(probe, x_bt_s_d, x_bt_s_d, need_weights=False)[0] + pooled_bt_d = pooled_bt_1_d.squeeze(1) # [B*T, D] + + pooled = pooled_bt_d.view(B, T, D) + pooled = self.layernorm(pooled) + return pooled + + +class CrossModalInteractionController: + """Strategy class to control dual-tower interaction. + + Manages the interaction mapping between Visual DiT (e.g., 30 layers) + and Audio DiT (e.g., 30 layers). + """ + + def __init__(self, visual_layers: int = 30, audio_layers: int = 30): + self.visual_layers = visual_layers + self.audio_layers = audio_layers + self.min_layers = min(visual_layers, audio_layers) + + def get_interaction_layers( + self, strategy: str = "shallow_focus" + ) -> Dict[str, List[Tuple[int, int]]]: + """Gets the mapping relationship of interaction layers.""" + if strategy == "shallow_focus": + num_interact = min(10, self.min_layers // 3) + interact_layers = list(range(0, num_interact)) + elif strategy == "distributed": + step = 3 + interact_layers = list(range(0, self.min_layers, step)) + elif strategy == "progressive": + shallow = list(range(0, min(8, self.min_layers))) + if self.min_layers > 8: + deep = list(range(8, self.min_layers, 3)) + interact_layers = shallow + deep + else: + interact_layers = shallow + elif strategy == "custom": + interact_layers = [0, 2, 4, 6, 8, 12, 16, 20] + interact_layers = [i for i in interact_layers if i < self.min_layers] + elif strategy == "full": + interact_layers = list(range(0, self.min_layers)) + else: + raise ValueError(f"Unknown interaction strategy: {strategy}") + + mapping = { + "v2a": [(i, i) for i in interact_layers], + "a2v": [(i, i) for i in interact_layers], + } + return mapping + + def should_interact( + self, layer_idx: int, direction: str, interaction_mapping: Dict + ) -> bool: + """Determines if the specified layer needs to interact.""" + if direction not in interaction_mapping: + return False + return any(src == layer_idx for src, _ in interaction_mapping[direction]) + + +class ConditionalCrossAttention(nn.Module): + """ + Cross-modal attention for dual-tower bridge with Tensor Parallel support. + + This module handles attention between video and audio hidden states, + which have different sequence lengths. + """ + + def __init__(self, dim: int, kv_dim: int, num_heads: int, eps: float = 1e-6): + super().__init__() + self.q_dim = dim + self.kv_dim = kv_dim + self.num_heads = num_heads + self.head_dim = self.q_dim // num_heads + + self.tp_size = get_tp_world_size() + if self.num_heads % self.tp_size != 0: + raise ValueError( + f"num_heads ({self.num_heads}) must be divisible by tp_size ({self.tp_size})." + ) + self.num_heads_per_rank = self.num_heads // self.tp_size + + # TP strategy: shard Q/K/V over heads (column-parallel), then row-parallel output. + self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False) + self.k = ColumnParallelLinear(kv_dim, dim, bias=True, gather_output=False) + self.v = ColumnParallelLinear(kv_dim, dim, bias=True, gather_output=False) + self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True) + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + + self.attn = USPAttention( + num_heads=self.num_heads_per_rank, + head_size=self.head_dim, + causal=False, + softmax_scale=None, + ) + + def forward( + self, + x: torch.Tensor, + y: torch.Tensor, + x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ): + ctx = y + q, _ = self.q(x) + k, _ = self.k(ctx) + v, _ = self.v(ctx) + + # RMSNorm over sharded hidden dimension + if self.tp_size > 1: + q = tensor_parallel_rms_norm(q, self.norm_q) + k = tensor_parallel_rms_norm(k, self.norm_k) + else: + q = self.norm_q(q) + k = self.norm_k(k) + + if x_freqs is not None: + x_cos, x_sin = x_freqs + q_view = rearrange(q, "b l (h d) -> b l h d", d=self.head_dim) + x_cos = x_cos.to(q_view.dtype).to(q_view.device).squeeze(0) + x_sin = x_sin.to(q_view.dtype).to(q_view.device).squeeze(0) + # FlashInfer expects cos_sin_cache with shape [seqlen, head_dim], + # where the first half is cos and the second half is sin, each with + # head_dim//2 elements. Since compute_rope_cos_sin duplicates the + # frequencies (cat((freqs, freqs))), we only take the first half. + half_dim = self.head_dim // 2 + cos_sin_cache = torch.cat( + [ + x_cos[:, :half_dim].to(dtype=torch.float32).contiguous(), + x_sin[:, :half_dim].to(dtype=torch.float32).contiguous(), + ], + dim=-1, + ) + q_view, _ = apply_flashinfer_rope_qk_inplace( + q_view, q_view.clone(), cos_sin_cache, is_neox=True + ) + q = rearrange(q_view, "b l h d -> b l (h d)") + + if y_freqs is not None: + y_cos, y_sin = y_freqs + k_view = rearrange(k, "b l (h d) -> b l h d", d=self.head_dim) + y_cos = y_cos.to(k_view.dtype).to(k_view.device).squeeze(0) + y_sin = y_sin.to(k_view.dtype).to(k_view.device).squeeze(0) + # FlashInfer expects cos_sin_cache with shape [seqlen, head_dim], + # where the first half is cos and the second half is sin, each with + # head_dim//2 elements. Since compute_rope_cos_sin duplicates the + # frequencies (cat((freqs, freqs))), we only take the first half. + half_dim = self.head_dim // 2 + cos_sin_cache = torch.cat( + [ + y_cos[:, :half_dim].to(dtype=torch.float32).contiguous(), + y_sin[:, :half_dim].to(dtype=torch.float32).contiguous(), + ], + dim=-1, + ) + k_view, _ = apply_flashinfer_rope_qk_inplace( + k_view, k_view.clone(), cos_sin_cache, is_neox=True + ) + k = rearrange(k_view, "b l h d -> b l (h d)") + + q = rearrange(q, "b l (h d) -> b l h d", h=self.num_heads_per_rank) + k = rearrange(k, "b l (h d) -> b l h d", h=self.num_heads_per_rank) + v = rearrange(v, "b l (h d) -> b l h d", h=self.num_heads_per_rank) + + x = self.attn(q, k, v) + x = rearrange(x, "b l h d -> b l (h d)") + x, _ = self.o(x) + return x + + +class AdaLayerNorm(nn.Module): + """ + Norm layer modified to incorporate timestep embeddings. + """ + + def __init__( + self, + embedding_dim: int, + num_embeddings: Optional[int] = None, + output_dim: Optional[int] = None, + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-5, + chunk_dim: int = 0, + ): + super().__init__() + + self.chunk_dim = chunk_dim + output_dim = output_dim or embedding_dim * 2 + + if num_embeddings is not None: + self.emb = nn.Embedding(num_embeddings, embedding_dim) + else: + self.emb = None + + self.silu = nn.SiLU() + self.linear = ReplicatedLinear(embedding_dim, output_dim) + self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine) + + def forward( + self, + x: torch.Tensor, + timestep: Optional[torch.Tensor] = None, + temb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if self.emb is not None: + temb = self.emb(timestep) + + temb, _ = self.linear(self.silu(temb)) + + if self.chunk_dim == 2: + scale, shift = temb.chunk(2, dim=2) + elif self.chunk_dim == 1: + shift, scale = temb.chunk(2, dim=1) + shift = shift[:, None, :] + scale = scale[:, None, :] + else: + scale, shift = temb.chunk(2, dim=0) + + x = self.norm(x) * (1 + scale) + shift + return x + + +class ConditionalCrossAttentionBlock(nn.Module): + """A wrapper block for ConditionalCrossAttention that applies LayerNorm to the condition input y.""" + + def __init__( + self, + dim: int, + kv_dim: int, + num_heads: int, + eps: float = 1e-6, + pooled_adaln: bool = False, + ): + super().__init__() + self.y_norm = nn.LayerNorm(kv_dim, eps=eps) + self.inner = ConditionalCrossAttention( + dim=dim, kv_dim=kv_dim, num_heads=num_heads, eps=eps + ) + self.pooled_adaln = pooled_adaln + if pooled_adaln: + self.per_frame_pooling = PerFrameAttentionPooling( + kv_dim, num_heads=num_heads, eps=eps + ) + self.adaln = AdaLayerNorm(kv_dim, output_dim=dim * 2, chunk_dim=2) + + def forward( + self, + x: torch.Tensor, + y: torch.Tensor, + x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + video_grid_size: Optional[Tuple[int, int, int]] = None, + ) -> torch.Tensor: + if self.pooled_adaln: + assert video_grid_size is not None, "video_grid_size cannot be None" + pooled_y = self.per_frame_pooling(y, video_grid_size) + if pooled_y.shape[1] != x.shape[1]: + pooled_y = F.interpolate( + pooled_y.permute(0, 2, 1), + size=x.shape[1], + mode="linear", + align_corners=False, + ).permute(0, 2, 1) + x = self.adaln(x, temb=pooled_y) + y = self.y_norm(y) + return self.inner(x=x, y=y, x_freqs=x_freqs, y_freqs=y_freqs) + + +class DualTowerConditionalBridge( + CachableDiT, + OffloadableDiTMixin, +): + """Dual-tower conditional bridge module v2 (SGLang optimized version). + + Implements the correct architecture: + 1. Audio latents -> Audio DiT -> Audio hidden states [B, L, 1536]. + 2. Visual latents -> Visual DiT -> Visual hidden states [B, L, 5120]. + 3. Cross-attention interaction between the hidden states of the two DiTs. + """ + + _fsdp_shard_conditions = MOVADualTowerConfig()._fsdp_shard_conditions + _compile_conditions = MOVADualTowerConfig()._compile_conditions + _supported_attention_backends = MOVADualTowerConfig()._supported_attention_backends + param_names_mapping = MOVADualTowerConfig().param_names_mapping + reverse_param_names_mapping = MOVADualTowerConfig().reverse_param_names_mapping + lora_param_names_mapping = MOVADualTowerConfig().lora_param_names_mapping + + def __init__( + self, + config: MOVADualTowerConfig | None = None, + hf_config: dict[str, Any] | None = None, + # Fallback parameters for from_pretrained compatibility + visual_layers: int = 40, + audio_layers: int = 30, + visual_hidden_dim: int = 5120, + audio_hidden_dim: int = 1536, + audio_fps: float = 50.0, + head_dim: int = 128, + interaction_strategy: str = "full", + apply_cross_rope: bool = True, + apply_first_frame_bias_in_rope: bool = False, + trainable_condition_scale: bool = False, + pooled_adaln: bool = False, + ): + super().__init__(config=config, hf_config=hf_config) + + # Use config if provided, otherwise use individual parameters + if config is not None: + visual_layers = config.visual_layers + audio_layers = config.audio_layers + visual_hidden_dim = config.visual_hidden_dim + audio_hidden_dim = config.audio_hidden_dim + audio_fps = config.audio_fps + head_dim = config.head_dim + interaction_strategy = config.interaction_strategy + apply_cross_rope = config.apply_cross_rope + apply_first_frame_bias_in_rope = config.apply_first_frame_bias_in_rope + trainable_condition_scale = config.trainable_condition_scale + pooled_adaln = config.pooled_adaln + + self.visual_hidden_dim = visual_hidden_dim + self.audio_hidden_dim = audio_hidden_dim + self.audio_fps = audio_fps + self.head_dim = head_dim + self.apply_cross_rope = apply_cross_rope + self.apply_first_frame_bias_in_rope = apply_first_frame_bias_in_rope + self.trainable_condition_scale = trainable_condition_scale + self.pooled_adaln = pooled_adaln + + if self.trainable_condition_scale: + self.condition_scale = nn.Parameter( + torch.tensor([1.0], dtype=torch.float32) + ) + else: + self.condition_scale = 1.0 + + self.controller = CrossModalInteractionController(visual_layers, audio_layers) + self.interaction_mapping = self.controller.get_interaction_layers( + interaction_strategy + ) + + # Cross-modal attention modules - interaction at DiT hidden states level + self.audio_to_video_conditioners = nn.ModuleDict() + self.video_to_audio_conditioners = nn.ModuleDict() + + self.rope_base = 10000.0 # RoPE base frequency hardcode. adapted from original mova implementation. + + # Audio DiT hidden states conditioning Video DiT + for v_layer, _ in self.interaction_mapping["a2v"]: + self.audio_to_video_conditioners[str(v_layer)] = ( + ConditionalCrossAttentionBlock( + dim=visual_hidden_dim, + kv_dim=audio_hidden_dim, + num_heads=visual_hidden_dim // head_dim, + pooled_adaln=False, + ) + ) + + # Visual DiT hidden states conditioning Audio DiT + for a_layer, _ in self.interaction_mapping["v2a"]: + self.video_to_audio_conditioners[str(a_layer)] = ( + ConditionalCrossAttentionBlock( + dim=audio_hidden_dim, + kv_dim=visual_hidden_dim, + num_heads=audio_hidden_dim // head_dim, + pooled_adaln=self.pooled_adaln, + ) + ) + + # Required attributes for CachableDiT/BaseDiT + self.hidden_size = visual_hidden_dim + self.num_attention_heads = visual_hidden_dim // head_dim + self.num_channels_latents = ( + visual_hidden_dim # Bridge doesn't output latents, but required by BaseDiT + ) + self.layer_names = [ + "audio_to_video_conditioners", + "video_to_audio_conditioners", + ] + self.__post_init__() + + @torch.no_grad() + def build_aligned_freqs( + self, + video_fps: float, + grid_size: Tuple[int, int, int], + audio_steps: int, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: + """Generates aligned RoPE (cos, sin) based on video FPS, grid size, and audio length. + + Uses functional RoPE computation to avoid FSDP meta device issues. + + Args: + video_fps: FPS of the video. + grid_size: Tuple of (f_v, h, w). + audio_steps: Length of the audio sequence. + device: Target device. + dtype: Output dtype. + + Returns: + A tuple of ((cos_v, sin_v), (cos_a, sin_a)). + """ + f_v, h, w = grid_size + L_v = f_v * h * w + L_a = int(audio_steps) + + device = device or next(self.parameters()).device + dtype = dtype or torch.float32 + + # Audio positions: 0, 1, 2, ..., L_a-1 + audio_pos = torch.arange(L_a, device=device, dtype=torch.float32).unsqueeze(0) + + # Video positions: Align video frames to audio step units + if self.apply_first_frame_bias_in_rope: + video_effective_fps = float(video_fps) / 4.0 + if f_v > 0: + t_starts = torch.zeros((f_v,), device=device, dtype=torch.float32) + if f_v > 1: + t_starts[1:] = (1.0 / float(video_fps)) + torch.arange( + f_v - 1, device=device, dtype=torch.float32 + ) * (1.0 / video_effective_fps) + else: + t_starts = torch.zeros((0,), device=device, dtype=torch.float32) + video_pos_per_frame = t_starts * float(self.audio_fps) + else: + scale = float(self.audio_fps) / float(video_fps / 4.0) + video_pos_per_frame = ( + torch.arange(f_v, device=device, dtype=torch.float32) * scale + ) + + video_pos = video_pos_per_frame.repeat_interleave(h * w).unsqueeze(0) + + # Use functional RoPE to compute cos/sin + cos_v, sin_v = compute_rope_cos_sin( + video_pos, self.head_dim, base=self.rope_base, device=device, dtype=dtype + ) + cos_a, sin_a = compute_rope_cos_sin( + audio_pos, self.head_dim, base=self.rope_base, device=device, dtype=dtype + ) + + return (cos_v, sin_v), (cos_a, sin_a) + + def should_interact(self, layer_idx: int, direction: str) -> bool: + return self.controller.should_interact( + layer_idx, direction, self.interaction_mapping + ) + + def apply_conditional_control( + self, + layer_idx: int, + direction: str, + primary_hidden_states: torch.Tensor, + condition_hidden_states: torch.Tensor, + x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + condition_scale: Optional[float] = None, + video_grid_size: Optional[Tuple[int, int, int]] = None, + ) -> torch.Tensor: + """Applies conditional control at the DiT hidden states level.""" + if not self.controller.should_interact( + layer_idx, direction, self.interaction_mapping + ): + return primary_hidden_states + + if direction == "a2v": + conditioner = self.audio_to_video_conditioners[str(layer_idx)] + elif direction == "v2a": + conditioner = self.video_to_audio_conditioners[str(layer_idx)] + else: + raise ValueError(f"Invalid direction: {direction}") + + conditioned_features = conditioner( + x=primary_hidden_states, + y=condition_hidden_states, + x_freqs=x_freqs, + y_freqs=y_freqs, + video_grid_size=video_grid_size, + ) + + if self.trainable_condition_scale and condition_scale is not None: + logger.warning( + "The current model has a trainable condition_scale, but condition_scale " + "was passed externally. Ignoring the trainable condition_scale and " + "using the external condition_scale=%s.", + condition_scale, + ) + + scale = condition_scale if condition_scale is not None else self.condition_scale + + primary_hidden_states = primary_hidden_states + conditioned_features * scale + + return primary_hidden_states + + def forward( + self, + layer_idx: int, + visual_hidden_states: torch.Tensor, + audio_hidden_states: torch.Tensor, + *, + x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + a2v_condition_scale: Optional[float] = None, + v2a_condition_scale: Optional[float] = None, + condition_scale: Optional[float] = None, + video_grid_size: Optional[Tuple[int, int, int]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Performs bidirectional conditional control for both visual and audio towers.""" + visual_conditioned = self.apply_conditional_control( + layer_idx=layer_idx, + direction="a2v", + primary_hidden_states=visual_hidden_states, + condition_hidden_states=audio_hidden_states, + x_freqs=x_freqs, + y_freqs=y_freqs, + condition_scale=( + a2v_condition_scale + if a2v_condition_scale is not None + else condition_scale + ), + video_grid_size=video_grid_size, + ) + + audio_conditioned = self.apply_conditional_control( + layer_idx=layer_idx, + direction="v2a", + primary_hidden_states=audio_hidden_states, + condition_hidden_states=visual_hidden_states, + x_freqs=y_freqs, + y_freqs=x_freqs, + condition_scale=( + v2a_condition_scale + if v2a_condition_scale is not None + else condition_scale + ), + video_grid_size=video_grid_size, + ) + + return visual_conditioned, audio_conditioned + + +EntryClass = DualTowerConditionalBridge diff --git a/python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py b/python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py new file mode 100644 index 000000000..a659a73d1 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py @@ -0,0 +1,306 @@ +# Copied and adapted from: mossVG/mova/diffusion/models/wan_audio_dit.py +# SPDX-License-Identifier: Apache-2.0 +# +# NOTE: This module reuses common functions from mova_video_dit.py to reduce code duplication. +# Audio-specific functions (precompute_freqs_cis_1d, legacy_precompute_freqs_cis_1d) are kept here. + +import math +from typing import Any, Optional, Tuple + +import torch +import torch.nn as nn +from einops import rearrange +from torch.distributed.tensor import DTensor + +from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig +from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear +from sglang.multimodal_gen.runtime.layers.mlp import MLP +from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT +from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin + +# Reuse common functions and classes from mova_video_dit +from .mova_video_dit import DiTBlock, precompute_freqs_cis, sinusoidal_embedding_1d + + +# Audio-specific positional encoding functions +def legacy_precompute_freqs_cis_1d( + dim: int, + end: int = 16384, + theta: float = 10000.0, + base_tps=4.0, + target_tps=44100 / 2048, +): + s = float(base_tps) / float(target_tps) + # 1d rope precompute + f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta, s) + # No positional encoding is applied to the remaining dimensions + no_freqs_cis = precompute_freqs_cis(dim // 3, end, theta, s) + no_freqs_cis = torch.ones_like(no_freqs_cis) + return f_freqs_cis, no_freqs_cis, no_freqs_cis + + +def precompute_freqs_cis_1d(dim: int, end: int = 16384, theta: float = 10000.0): + f_freqs_cis = precompute_freqs_cis(dim, end, theta) + return f_freqs_cis.chunk(3, dim=-1) + + +class Head(nn.Module): + def __init__( + self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float + ): + super().__init__() + self.dim = dim + self.patch_size = patch_size + self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.head = ReplicatedLinear(dim, out_dim * math.prod(patch_size)) + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, t_mod): + if len(t_mod.shape) == 3: + shift, scale = ( + self.modulation.unsqueeze(0).to(dtype=t_mod.dtype, device=t_mod.device) + + t_mod.unsqueeze(2) + ).chunk(2, dim=2) + x, _ = self.head(self.norm(x) * (1 + scale.squeeze(2)) + shift.squeeze(2)) + else: + # NOTE: t_mod was originally [B, C]. This works correctly with broadcasting when B=1, but it won't match [1, 2, C] when B > 1. + shift, scale = ( + self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + + t_mod.unsqueeze(1) + ).chunk(2, dim=1) + x, _ = self.head(self.norm(x) * (1 + scale) + shift) + return x + + +class Conv1dLocalIsland(nn.Conv1d): + """Inherits from Conv1d and overrides forward. + + - Parameters remain as DTensors (optimizer consistency is maintained). + - In the forward pass, x, weight, and bias are aggregated as Replicate, + and then local convolution is performed via to_local. + - The output is then redistributed as a DTensor (default is Replicate, + placements can be customized). + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, input): + if isinstance(input, DTensor): + x_local = input.to_local() # type: ignore[attr-defined] + w_local = self.weight.to_local() # type: ignore[attr-defined] + b_local = ( + self.bias.to_local() if self.bias is not None else None # type: ignore[attr-defined] + ) + + return self._conv_forward(x_local, w_local, b_local) + else: + return super().forward(input) + + +class WanAudioModel(CachableDiT, OffloadableDiTMixin): + _fsdp_shard_conditions = MOVAAudioConfig()._fsdp_shard_conditions + _compile_conditions = MOVAAudioConfig()._compile_conditions + _supported_attention_backends = MOVAAudioConfig()._supported_attention_backends + param_names_mapping = MOVAAudioConfig().param_names_mapping + reverse_param_names_mapping = MOVAAudioConfig().reverse_param_names_mapping + lora_param_names_mapping = MOVAAudioConfig().lora_param_names_mapping + + def __init__(self, config: MOVAAudioConfig, hf_config: dict[str, Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + + # Extract parameters from config + dim = config.dim + in_dim = config.in_dim + ffn_dim = config.ffn_dim + out_dim = config.out_dim + text_dim = config.text_dim + freq_dim = config.freq_dim + eps = config.eps + patch_size = config.patch_size + num_heads = config.num_heads + num_layers = config.num_layers + has_image_pos_emb = config.has_image_pos_emb + has_ref_conv = config.has_ref_conv + add_control_adapter = config.add_control_adapter + in_dim_control_adapter = config.in_dim_control_adapter + seperated_timestep = config.seperated_timestep + require_vae_embedding = config.require_vae_embedding + require_clip_embedding = config.require_clip_embedding + fuse_vae_embedding_in_latents = config.fuse_vae_embedding_in_latents + vae_type = config.vae_type + + self.dim = dim + self.freq_dim = freq_dim + self.patch_size = patch_size + self.seperated_timestep = seperated_timestep + self.require_vae_embedding = require_vae_embedding + self.require_clip_embedding = require_clip_embedding + self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents + self.vae_type = vae_type + # self.patch_embedding = nn.Conv3d( + # in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.patch_embedding = Conv1dLocalIsland( + in_dim, dim, kernel_size=patch_size, stride=patch_size + ) + self.text_embedding = MLP( + text_dim, dim, output_dim=dim, act_type="gelu_pytorch_tanh" + ) + self.time_embedding = MLP(freq_dim, dim, output_dim=dim, act_type="silu") + # Preserve state_dict keys (time_projection.1.weight/bias). + self.time_projection = nn.Sequential(nn.SiLU(), ReplicatedLinear(dim, dim * 6)) + self.blocks = nn.ModuleList( + [DiTBlock(dim, num_heads, ffn_dim, eps) for _ in range(num_layers)] + ) + self.head = Head(dim, out_dim, patch_size, eps) + self.num_heads = num_heads + self.freqs = None + self.img_pos_emb = None + if has_ref_conv: + self.ref_conv = nn.Conv2d(16, dim, kernel_size=(2, 2), stride=(2, 2)) + self.has_image_pos_emb = has_image_pos_emb + self.has_ref_conv = has_ref_conv + self.hidden_size = dim + self.num_attention_heads = num_heads + self.num_channels_latents = out_dim + self.layer_names = ["blocks"] + self.cnt = 0 + self.teacache_thresh = 0 + self.coefficients = [] + self.accumulated_rel_l1_distance = 0 + self.previous_modulated_input = None + self.previous_resiual = None + self.previous_e0_even = None + self.previous_e0_odd = None + self.previous_residual_even = None + self.previous_residual_odd = None + self.is_even = False + self.should_calc_even = True + self.should_calc_odd = True + self.accumulated_rel_l1_distance_even = 0 + self.accumulated_rel_l1_distance_odd = 0 + self.__post_init__() + if add_control_adapter: + from .wan_video_camera_controller import SimpleAdapter + + self.control_adapter = SimpleAdapter( + in_dim_control_adapter, + dim, + kernel_size=patch_size[1:], + stride=patch_size[1:], + ) + else: + self.control_adapter = None + + def _init_freqs(self): + if self.freqs is not None: + return + head_dim = self.dim // self.num_heads + if self.vae_type == "dac": + self.freqs = precompute_freqs_cis_1d(head_dim) + else: + raise ValueError(f"Invalid VAE type: {self.vae_type}") + + def patchify( + self, + x: torch.Tensor, + control_camera_latents_input: Optional[torch.Tensor] = None, + ): + x = self.patch_embedding(x) + if ( + self.control_adapter is not None + and control_camera_latents_input is not None + ): + y_camera = self.control_adapter(control_camera_latents_input) + if isinstance(x, list): + x = [u + v for u, v in zip(x, y_camera)] + x = x[0].unsqueeze(0) + else: + if isinstance(y_camera, list): + x = x + y_camera[0] + else: + x = x + y_camera + grid_size = x.shape[2:] + x = rearrange(x, "b c f -> b f c").contiguous() + return x, grid_size # x, grid_size: (f) + + def unpatchify(self, x: torch.Tensor, grid_size: tuple[int]): + return rearrange( + x, "b f (p c) -> b c (f p)", f=grid_size[0], p=self.patch_size[0] + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + guidance=None, + use_gradient_checkpointing: bool = False, + use_gradient_checkpointing_offload: bool = False, + **kwargs, + ) -> torch.Tensor: + # MOVA audio uses x/context naming historically. + x = hidden_states + context = ( + encoder_hidden_states[0] + if isinstance(encoder_hidden_states, list) + else encoder_hidden_states + ) + + t = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep)) + t_proj, _ = self.time_projection(t) + t_mod = t_proj.unflatten(1, (6, self.dim)) + context = self.text_embedding(context) + + x, (f,) = self.patchify(x) + + freqs = ( + torch.cat( + [ + self.freqs[0][:f].view(f, -1).expand(f, -1), + self.freqs[1][:f].view(f, -1).expand(f, -1), + self.freqs[2][:f].view(f, -1).expand(f, -1), + ], + dim=-1, + ) + .reshape(f, 1, -1) + .to(x.device) + ) + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + for block in self.blocks: + if self.training and use_gradient_checkpointing: + if use_gradient_checkpointing_offload: + with torch.autograd.graph.save_on_cpu(): + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, + context, + t_mod, + freqs, + use_reentrant=False, + ) + else: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, + context, + t_mod, + freqs, + use_reentrant=False, + ) + else: + x = block(x, context, t_mod, freqs) + + x = self.head(x, t) + x = self.unpatchify(x, (f,)) + return x + + +EntryClass = WanAudioModel diff --git a/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py b/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py new file mode 100644 index 000000000..9fd7babcc --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py @@ -0,0 +1,570 @@ +# Copied and adapted from: mossVG/mova/diffusion/models/wan_video_dit.py +# SPDX-License-Identifier: Apache-2.0 +# +# NOTE: This module shares common functions (sinusoidal_embedding_1d, precompute_freqs_cis, etc.) +# with wanvideo.py. These functions are kept here for MOVA-specific model architecture, +# but could be refactored to a common module in the future. + +import math +from typing import Any, Tuple + +import torch +import torch.nn as nn +from einops import rearrange +from torch.distributed.tensor import DTensor + +from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig +from sglang.multimodal_gen.runtime.distributed import get_tp_world_size +from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAttention + +# Reuse SGLang's optimized RMSNorm instead of torch.nn.RMSNorm or custom SlowRMSNorm +from sglang.multimodal_gen.runtime.layers.layernorm import ( + RMSNorm, + tensor_parallel_rms_norm, +) +from sglang.multimodal_gen.runtime.layers.linear import ( + ColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.multimodal_gen.runtime.layers.mlp import MLP +from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT +from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +# @torch.compile(fullgraph=True) +def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor): + return x * (1 + scale) + shift + + +def sinusoidal_embedding_1d(dim, position): + sinusoid = torch.outer( + position.type(torch.float64), + torch.pow( + 10000, + -torch.arange(dim // 2, dtype=torch.float64, device=position.device).div( + dim // 2 + ), + ), + ) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x.to(position.dtype) + + +def precompute_freqs_cis_3d(dim: int, end: int = 1024, theta: float = 10000.0): + # 3d rope precompute + f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta) + h_freqs_cis = precompute_freqs_cis(dim // 3, end, theta) + w_freqs_cis = precompute_freqs_cis(dim // 3, end, theta) + return f_freqs_cis, h_freqs_cis, w_freqs_cis + + +def precompute_freqs_cis( + dim: int, end: int = 1024, theta: float = 10000.0, s: float = 1.0 +): + # 1d rope precompute + # Note: s parameter is used for audio-specific scaling (e.g., tps adjustment) + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].double() / dim)) + pos = torch.arange(end, dtype=torch.float64, device=freqs.device) * s + freqs = torch.outer(pos, freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 + return freqs_cis + + +def rope_apply(x, freqs, num_heads): + x = rearrange(x, "b s (n d) -> b s n d", n=num_heads) + x_out = torch.view_as_complex( + x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2) + ) + x_out = torch.view_as_real(x_out * freqs).flatten(2) + return x_out.to(x.dtype) + + +def rope_apply_head_dim(x, freqs, head_dim): + x = rearrange(x, "b s (n d) -> b s n d", d=head_dim) + x_out = torch.view_as_complex( + x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2) + ) + # print(f"{x_out.shape = }, {freqs.shape = }") + x_out = torch.view_as_real(x_out * freqs).flatten(2) + return x_out.to(x.dtype) + + +class SelfAttention(nn.Module): + """ + Self-Attention module for MOVA DiT with Sequence Parallelism support. + + SP is handled at the pipeline level (latents are pre-sharded before DiT forward). + USPAttention internally handles the all-to-all communication for distributed attention. + Input x should already be the local shard [B, S_local, D] when SP is enabled. + """ + + def __init__(self, dim: int, num_heads: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + + self.tp_size = get_tp_world_size() + if self.num_heads % self.tp_size != 0: + raise ValueError( + f"num_heads ({self.num_heads}) must be divisible by tp_size ({self.tp_size})." + ) + self.num_heads_per_rank = self.num_heads // self.tp_size + + # TP strategy: shard Q/K/V over heads (column-parallel), then row-parallel output. + self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False) + self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False) + self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False) + self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True) + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + + self.attn = USPAttention( + # Local heads per TP rank. + num_heads=self.num_heads_per_rank, + head_size=self.head_dim, + causal=False, + softmax_scale=None, + ) + + def forward(self, x, freqs): + """ + Forward pass for self-attention. + + Args: + x: Input tensor [B, S_local, D] - already sharded by SP when SP > 1 + freqs: RoPE frequencies [S_local, 1, head_dim] - should match x's sequence length + + Returns: + Output tensor [B, S_local, D] + """ + if isinstance(freqs, DTensor): + freqs = freqs.to_local() + + # Compute Q, K, V on local sequence + q, _ = self.q(x) + k, _ = self.k(x) + v, _ = self.v(x) + + # RMSNorm over sharded hidden dimension. + if self.tp_size > 1: + q = tensor_parallel_rms_norm(q, self.norm_q) + k = tensor_parallel_rms_norm(k, self.norm_k) + else: + q = self.norm_q(q) + k = self.norm_k(k) + + # Apply RoPE + q = rope_apply_head_dim(q, freqs, self.head_dim) + k = rope_apply_head_dim(k, freqs, self.head_dim) + + # USPAttention expects [B, S_local, H, D] format + q = rearrange(q, "b s (n d) -> b s n d", n=self.num_heads_per_rank) + k = rearrange(k, "b s (n d) -> b s n d", n=self.num_heads_per_rank) + v = rearrange(v, "b s (n d) -> b s n d", n=self.num_heads_per_rank) + + # USPAttention handles SP communication internally + out = self.attn(q, k, v) + out = rearrange(out, "b s n d -> b s (n d)") + + out, _ = self.o(out) + return out + + +class CrossAttention(nn.Module): + """ + Cross-Attention module for MOVA DiT. + + Cross-attention does NOT require SP communication because: + - Query comes from the main sequence (already sharded by SP) + - Key/Value come from context (text embeddings, which are replicated across all ranks) + + Uses LocalAttention instead of USPAttention for efficiency. + """ + + def __init__(self, dim: int, num_heads: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + + self.tp_size = get_tp_world_size() + if self.num_heads % self.tp_size != 0: + raise ValueError( + f"num_heads ({self.num_heads}) must be divisible by tp_size ({self.tp_size})." + ) + self.num_heads_per_rank = self.num_heads // self.tp_size + + self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False) + self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False) + self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False) + self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True) + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + + # Use LocalAttention for cross-attention (no SP communication needed) + self.attn = LocalAttention( + num_heads=self.num_heads_per_rank, + head_size=self.head_dim, + causal=False, + softmax_scale=None, + ) + + def forward(self, x: torch.Tensor, y: torch.Tensor): + """ + Forward pass for cross-attention. + + Args: + x: Query tensor [B, S_local, D] - the main sequence (sharded by SP) + y: Context tensor [B, S_ctx, D] - text/image embeddings (replicated) + + Returns: + Output tensor [B, S_local, D] + """ + ctx = y + + q, _ = self.q(x) + k, _ = self.k(ctx) + v, _ = self.v(ctx) + + if self.tp_size > 1: + q = tensor_parallel_rms_norm(q, self.norm_q) + k = tensor_parallel_rms_norm(k, self.norm_k) + else: + q = self.norm_q(q) + k = self.norm_k(k) + + q = rearrange(q, "b s (n d) -> b s n d", n=self.num_heads_per_rank) + k = rearrange(k, "b s (n d) -> b s n d", n=self.num_heads_per_rank) + v = rearrange(v, "b s (n d) -> b s n d", n=self.num_heads_per_rank) + x = self.attn(q, k, v) + x = rearrange(x, "b s n d -> b s (n d)") + x, _ = self.o(x) + return x + + +class GateModule(nn.Module): + def __init__( + self, + ): + super().__init__() + + def forward(self, x, gate, residual): + return x + gate * residual + + +class DiTBlock(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + ffn_dim: int, + eps: float = 1e-6, + ): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.ffn_dim = ffn_dim + + self.self_attn = SelfAttention(dim, num_heads, eps) + self.cross_attn = CrossAttention(dim, num_heads, eps) + self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm3 = nn.LayerNorm(dim, eps=eps) + self.ffn = MLP(dim, ffn_dim, output_dim=dim, act_type="gelu_pytorch_tanh") + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + self.gate = GateModule() + + def forward(self, x, context, t_mod, freqs): + has_seq = len(t_mod.shape) == 4 + chunk_dim = 2 if has_seq else 1 + # msa: multi-head self-attention mlp: multi-layer perceptron + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod + ).chunk(6, dim=chunk_dim) + if has_seq: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + shift_msa.squeeze(2), + scale_msa.squeeze(2), + gate_msa.squeeze(2), + shift_mlp.squeeze(2), + scale_mlp.squeeze(2), + gate_mlp.squeeze(2), + ) + input_x = modulate(self.norm1(x), shift_msa, scale_msa) + x = self.gate(x, gate_msa, self.self_attn(input_x, freqs)) + x = x + self.cross_attn(self.norm3(x), context) + input_x = modulate(self.norm2(x), shift_mlp, scale_mlp) + x = self.gate(x, gate_mlp, self.ffn(input_x)) + return x + + +class Head(nn.Module): + def __init__( + self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float + ): + super().__init__() + self.dim = dim + self.patch_size = patch_size + self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + # Output dim is small for MOVA; replicate to avoid TP shape coupling. + self.head = ReplicatedLinear(dim, out_dim * math.prod(patch_size)) + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, t_mod): + if len(t_mod.shape) == 3: + shift, scale = ( + self.modulation.unsqueeze(0).to(dtype=t_mod.dtype, device=t_mod.device) + + t_mod.unsqueeze(2) + ).chunk(2, dim=2) + x, _ = self.head(self.norm(x) * (1 + scale.squeeze(2)) + shift.squeeze(2)) + else: + shift, scale = ( + self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod + ).chunk(2, dim=1) + x, _ = self.head(self.norm(x) * (1 + scale) + shift) + return x + + +class Conv3dLocalIsland(nn.Conv3d): + """ + Inherits from Conv3d and overrides the forward method. + + Key behaviors: + - Parameters are kept as DTensor to maintain optimizer consistency. + - The forward pass aggregates input, weight, and bias into a Replicate state, + then performs the convolution locally using to_local(). + - The output is then redistributed as a DTensor (defaults to Replicate, + but placements can be customized). + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, input): + if isinstance(input, DTensor): + # NOTE: DTensor typing stubs are incomplete; at runtime DTensor has + # to_local() and parameters may also be DTensor. + x_local = input.to_local() # type: ignore[attr-defined] + w_local = self.weight.to_local() # type: ignore[attr-defined] + b_local = ( + self.bias.to_local() if self.bias is not None else None # type: ignore[attr-defined] + ) + + return self._conv_forward(x_local, w_local, b_local) + else: + return super().forward(input) + + +class WanModel(CachableDiT, OffloadableDiTMixin): + _fsdp_shard_conditions = MOVAVideoConfig()._fsdp_shard_conditions + _compile_conditions = MOVAVideoConfig()._compile_conditions + _supported_attention_backends = MOVAVideoConfig()._supported_attention_backends + param_names_mapping = MOVAVideoConfig().param_names_mapping + reverse_param_names_mapping = MOVAVideoConfig().reverse_param_names_mapping + lora_param_names_mapping = MOVAVideoConfig().lora_param_names_mapping + + def __init__(self, config: MOVAVideoConfig, hf_config: dict[str, Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + + # Extract parameters from config + dim = config.dim + in_dim = config.in_dim + ffn_dim = config.ffn_dim + out_dim = config.out_dim + text_dim = config.text_dim + freq_dim = config.freq_dim + eps = config.eps + patch_size = config.patch_size + num_heads = config.num_heads + num_layers = config.num_layers + has_image_pos_emb = config.has_image_pos_emb + has_ref_conv = config.has_ref_conv + add_control_adapter = config.add_control_adapter + in_dim_control_adapter = config.in_dim_control_adapter + seperated_timestep = config.seperated_timestep + require_vae_embedding = config.require_vae_embedding + require_clip_embedding = config.require_clip_embedding + fuse_vae_embedding_in_latents = config.fuse_vae_embedding_in_latents + + self.dim = dim + self.freq_dim = freq_dim + self.patch_size = patch_size + self.seperated_timestep = seperated_timestep + self.require_vae_embedding = require_vae_embedding + self.require_clip_embedding = require_clip_embedding + self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents + + self.patch_embedding = Conv3dLocalIsland( + in_dim, dim, kernel_size=patch_size, stride=patch_size + ) + self.text_embedding = MLP( + text_dim, dim, output_dim=dim, act_type="gelu_pytorch_tanh" + ) + self.time_embedding = MLP(freq_dim, dim, output_dim=dim, act_type="silu") + # Preserve state_dict keys (time_projection.1.weight/bias). + self.time_projection = nn.Sequential(nn.SiLU(), ReplicatedLinear(dim, dim * 6)) + self.blocks = nn.ModuleList( + [DiTBlock(dim, num_heads, ffn_dim, eps) for _ in range(num_layers)] + ) + self.head = Head(dim, out_dim, patch_size, eps) + self.num_heads = num_heads + self.freqs = None + + if has_ref_conv: + self.ref_conv = nn.Conv2d(16, dim, kernel_size=(2, 2), stride=(2, 2)) + self.has_image_pos_emb = has_image_pos_emb + self.has_ref_conv = has_ref_conv + self.hidden_size = dim + self.num_attention_heads = num_heads + self.num_channels_latents = out_dim + self.layer_names = ["blocks"] + self.cnt = 0 + self.teacache_thresh = 0 + self.coefficients = [] + self.accumulated_rel_l1_distance = 0 + self.previous_modulated_input = None + self.previous_resiual = None + self.previous_e0_even = None + self.previous_e0_odd = None + self.previous_residual_even = None + self.previous_residual_odd = None + self.is_even = False + self.should_calc_even = True + self.should_calc_odd = True + self.accumulated_rel_l1_distance_even = 0 + self.accumulated_rel_l1_distance_odd = 0 + self.__post_init__() + if add_control_adapter: + from .wan_video_camera_controller import SimpleAdapter + + self.control_adapter = SimpleAdapter( + in_dim_control_adapter, + dim, + kernel_size=patch_size[1:], + stride=patch_size[1:], + ) + else: + self.control_adapter = None + + def _init_freqs(self): + if self.freqs is not None: + return + head_dim = self.dim // self.num_heads + self.freqs = precompute_freqs_cis_3d(head_dim) + + def patchify( + self, x: torch.Tensor, control_camera_latents_input: torch.Tensor | None = None + ): + # NOTE(dhyu): avoid slow_conv + x = x.contiguous(memory_format=torch.channels_last_3d) + x = self.patch_embedding(x) + if ( + self.control_adapter is not None + and control_camera_latents_input is not None + ): + y_camera = self.control_adapter(control_camera_latents_input) + if isinstance(x, list): + x = [u + v for u, v in zip(x, y_camera)] + x = x[0].unsqueeze(0) + else: + # Some adapters may return a list even when x is a Tensor. + if isinstance(y_camera, list): + x = x + y_camera[0] + else: + x = x + y_camera + grid_size = x.shape[2:] + x = rearrange(x, "b c f h w -> b (f h w) c").contiguous() + return x, grid_size # x, grid_size: (f, h, w) + + def unpatchify(self, x: torch.Tensor, grid_size: tuple[int, int, int]): + return rearrange( + x, + "b (f h w) (x y z c) -> b c (f x) (h y) (w z)", + f=grid_size[0], + h=grid_size[1], + w=grid_size[2], + x=self.patch_size[0], + y=self.patch_size[1], + z=self.patch_size[2], + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor], + y: torch.Tensor, + guidance=None, + use_gradient_checkpointing: bool = False, + use_gradient_checkpointing_offload: bool = False, + **kwargs, + ) -> torch.Tensor: + # MOVA code historically uses x/context/y/clip_feature naming. + x = hidden_states + context = ( + encoder_hidden_states[0] + if isinstance(encoder_hidden_states, list) + else encoder_hidden_states + ) + t = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep)) + t_proj, _ = self.time_projection(t) + t_mod = t_proj.unflatten(1, (6, self.dim)) + context = self.text_embedding(context) + + x, (f, h, w) = self.patchify(x) + + freqs = ( + torch.cat( + [ + self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ) + .reshape(f * h * w, 1, -1) + .to(x.device) + ) + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + for block in self.blocks: + if self.training and use_gradient_checkpointing: + if use_gradient_checkpointing_offload: + with torch.autograd.graph.save_on_cpu(): + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, + context, + t_mod, + freqs, + use_reentrant=False, + ) + else: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, + context, + t_mod, + freqs, + use_reentrant=False, + ) + else: + x = block(x, context, t_mod, freqs) + + x = self.head(x, t) + x = self.unpatchify(x, (f, h, w)) + return x + + +EntryClass = WanModel diff --git a/python/sglang/multimodal_gen/runtime/models/model_stages/mova.py b/python/sglang/multimodal_gen/runtime/models/model_stages/mova.py new file mode 100644 index 000000000..afbe4a914 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/model_stages/mova.py @@ -0,0 +1,911 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +MOVA-specific pipeline stages. + +Sequence Parallelism (SP) Support: +- Video latents are sharded along the sequence dimension (T*H*W) after patchify +- Audio latents are sharded along the sequence dimension (L) after patchify +- USPAttention handles all-to-all communication internally +- Latents are gathered before unpatchify to restore full sequence +""" + +from __future__ import annotations + +import functools +import inspect +import os +from collections.abc import Iterable + +import torch +from diffusers.utils.torch_utils import randn_tensor +from tqdm.auto import tqdm + +from sglang.multimodal_gen.runtime.distributed import ( + get_local_torch_device, + get_world_group, +) +from sglang.multimodal_gen.runtime.distributed.communication_op import ( + cfg_model_parallel_all_reduce, + sequence_model_parallel_all_gather, +) +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_cfg_group, + get_classifier_free_guidance_rank, + get_sp_parallel_rank, + get_sp_world_size, +) +from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context + +# Both audio and video DiT use the same sinusoidal_embedding_1d function +# Import from mova_video_dit where it's defined (mova_audio_dit re-exports it) +from sglang.multimodal_gen.runtime.models.dits.mova_video_dit import ( + sinusoidal_embedding_1d, +) + +# Create aliases for backward compatibility +video_sinusoidal_embedding_1d = sinusoidal_embedding_1d +audio_sinusoidal_embedding_1d = sinusoidal_embedding_1d +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( + PipelineStage, + StageParallelismType, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import ( + _ensure_tensor_decode_output, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( + StageValidators as V, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( + VerificationResult, +) +from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args +from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler +from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE + +logger = init_logger(__name__) + + +class MOVALatentPreparationStage(PipelineStage): + """Prepare video/audio noise latents for MOVA.""" + + def __init__(self, audio_vae, require_vae_embedding: bool = True) -> None: + super().__init__() + self.audio_vae = audio_vae + self.require_vae_embedding = require_vae_embedding + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + batch_size = batch.batch_size + num_frames = batch.num_frames + if num_frames is None: + raise ValueError("num_frames is required for MOVA") + + audio_num_samples = int(self.audio_vae.sample_rate * num_frames / batch.fps) + + video_shape = server_args.pipeline_config.prepare_latent_shape( + batch, batch_size, num_frames + ) + audio_shape = server_args.pipeline_config.prepare_audio_latent_shape( + batch_size, audio_num_samples, self.audio_vae + ) + + device = get_local_torch_device() + generator = batch.generator + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + dit_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision] + batch.latents = randn_tensor( + video_shape, generator=generator, device=device, dtype=dit_dtype + ) + batch.audio_latents = randn_tensor( + audio_shape, generator=generator, device=device, dtype=dit_dtype + ) + + if batch.image_latent is not None: + batch.y = batch.image_latent.to(device=device, dtype=dit_dtype) + elif self.require_vae_embedding: + raise ValueError("MOVA requires reference image latents for denoising") + return batch + + +class MOVATimestepPreparationStage(PipelineStage): + """Prepare paired timesteps for MOVA.""" + + def __init__(self, scheduler) -> None: + super().__init__() + self.scheduler = scheduler + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + self.scheduler.set_timesteps( + batch.num_inference_steps, + denoising_strength=1.0, + shift=getattr(batch, "sigma_shift", self.scheduler.shift), + ) + self.scheduler.set_pair_postprocess_by_name( + "dual_sigma_shift", + visual_shift=getattr(batch, "visual_shift", 5.0), + audio_shift=getattr(batch, "audio_shift", 5.0), + ) + paired = self.scheduler.get_pairs() + batch.paired_timesteps = paired + batch.timesteps = paired + return batch + + +class MOVADenoisingStage(PipelineStage): + """Run MOVA dual-tower denoising loop.""" + + def __init__(self, video_dit, video_dit_2, audio_dit, dual_tower_bridge, scheduler): + super().__init__() + self.video_dit = video_dit + self.video_dit_2 = video_dit_2 + self.audio_dit = audio_dit + self.dual_tower_bridge = dual_tower_bridge + self.scheduler = scheduler + self._cache_dit_enabled = False + self._cached_num_steps = None + self._torch_compiled = False + + @property + def parallelism_type(self) -> StageParallelismType: + if get_global_server_args().enable_cfg_parallel: + return StageParallelismType.CFG_PARALLEL + return StageParallelismType.REPLICATED + + def _predict( + self, + visual_dit, + visual_latents, + audio_latents, + y, + context, + timestep, + audio_timestep, + video_fps, + timestep_index: int, + attn_metadata, + forward_batch: Req | None = None, + ): + # Set forward context for distributed attention (USPAttention) + with set_forward_context( + current_timestep=timestep_index, + attn_metadata=attn_metadata, + forward_batch=forward_batch, + ): + return self.inference_single_step( + visual_dit=visual_dit, + visual_latents=visual_latents, + audio_latents=audio_latents, + y=y, + context=context, + timestep=timestep, + audio_timestep=audio_timestep, + video_fps=video_fps, + ) + + def _cfg_combine(self, pos, neg, guidance_scale, cfg_rank, enable_cfg_parallel): + if not enable_cfg_parallel: + return neg + guidance_scale * (pos - neg) + if cfg_rank == 0: + partial = guidance_scale * pos + else: + partial = (1 - guidance_scale) * neg + return cfg_model_parallel_all_reduce(partial) + + def compile_module_with_torch_compile(self, module, server_args: ServerArgs): + if not server_args.enable_torch_compile or module is None: + return module + if not hasattr(module, "forward"): + return module + try: + import torch._inductor.config as _inductor_cfg + + _inductor_cfg.reorder_for_compute_comm_overlap = True + except ImportError: + pass + mode = os.environ.get("SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs") + logger.info("Compiling %s with mode: %s", module.__class__.__name__, mode) + compiled_forward = torch.compile(getattr(module, "forward"), mode=mode) + setattr(module, "forward", compiled_forward) + return module + + def _maybe_compile_dits(self, server_args: ServerArgs): + if self._torch_compiled or not server_args.enable_torch_compile: + return + for module in filter(None, [self.video_dit, self.video_dit_2, self.audio_dit]): + self.compile_module_with_torch_compile(module, server_args) + self._torch_compiled = True + + def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + """Verify denoising stage inputs.""" + result = VerificationResult() + result.add_check("y", batch.y, V.is_tensor) + result.add_check("paired_timesteps", batch.paired_timesteps, V.is_tensor) + result.add_check("latents", batch.latents, V.is_tensor) + result.add_check("audio_latents", batch.audio_latents, V.is_tensor) + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + result.add_check( + "negative_prompt_embeds", + batch.negative_prompt_embeds, + lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x), + ) + result.add_check( + "num_inference_steps", batch.num_inference_steps, V.positive_int + ) + result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float) + result.add_check( + "guidance_rescale", batch.guidance_rescale, V.non_negative_float + ) + result.add_check( + "do_classifier_free_guidance", + batch.do_classifier_free_guidance, + V.bool_value, + ) + return result + + def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + """Verify denoising stage outputs.""" + result = VerificationResult() + result.add_check("latents", batch.latents, V.is_tensor) + result.add_check("audio_latents", batch.audio_latents, V.is_tensor) + return result + + def progress_bar( + self, iterable: Iterable | None = None, total: int | None = None + ) -> tqdm: + """ + Create a progress bar for the denoising process. + """ + local_rank = get_world_group().local_rank + disable = local_rank != 0 + return tqdm(iterable=iterable, total=total, disable=disable) + + def step_profile(self): + profiler = SGLDiffusionProfiler.get_instance() + if profiler: + profiler.step_denoising_step() + + def rescale_noise_cfg( + self, noise_cfg, noise_pred_text, guidance_rescale=0.0 + ) -> torch.Tensor: + """ + Rescale noise prediction according to guidance_rescale. + + Based on findings of "Common Diffusion Noise Schedules and Sample Steps are Flawed" + (https://arxiv.org/pdf/2305.08891.pdf), Section 3.4. + """ + std_text = noise_pred_text.std( + dim=list(range(1, noise_pred_text.ndim)), keepdim=True + ) + std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) + noise_pred_rescaled = noise_cfg * (std_text / std_cfg) + noise_cfg = ( + guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg + ) + return noise_cfg + + def prepare_extra_func_kwargs(self, func, kwargs) -> dict[str, object]: + if not kwargs: + return {} + + if isinstance(func, functools.partial) and func.args: + func = getattr(func.args[0], "_original_forward", func) + + target_func = inspect.unwrap(func) + params = inspect.signature(target_func).parameters + return {k: v for k, v in kwargs.items() if k in params} + + def _build_attn_metadata( + self, i: int, batch: Req, server_args: ServerArgs + ) -> object | None: + return None + + def _manage_device_placement( + self, + model_to_use: torch.nn.Module | None, + model_to_offload: torch.nn.Module | None, + server_args: ServerArgs, + ): + if not server_args.dit_cpu_offload: + return + + if ( + model_to_offload is not None + and next(model_to_offload.parameters()).device.type == "cuda" + ): + model_to_offload.to("cpu") + + if ( + model_to_use is not None + and next(model_to_use.parameters()).device.type == "cpu" + ): + model_to_use.to(get_local_torch_device()) + + def _select_visual_dit( + self, timestep: float, boundary_ratio: float | None, server_args: ServerArgs + ): + if boundary_ratio is None or self.video_dit_2 is None: + self._manage_device_placement(self.video_dit, None, server_args) + return self.video_dit + + boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps + if timestep >= boundary_timestep: + current_model = self.video_dit + model_to_offload = self.video_dit_2 + else: + current_model = self.video_dit_2 + model_to_offload = self.video_dit + + self._manage_device_placement(current_model, model_to_offload, server_args) + return current_model + + def _apply_guidance_rescale( + self, + noise_pred, + noise_pred_text, + guidance_rescale, + cfg_rank, + enable_cfg_parallel, + ): + if guidance_rescale <= 0.0: + return noise_pred + if enable_cfg_parallel: + std_cfg = noise_pred.std(dim=list(range(1, noise_pred.ndim)), keepdim=True) + if cfg_rank == 0: + assert noise_pred_text is not None + std_text = noise_pred_text.std( + dim=list(range(1, noise_pred_text.ndim)), keepdim=True + ) + else: + std_text = torch.empty_like(std_cfg) + std_text = get_cfg_group().broadcast(std_text, src=0) + noise_pred_rescaled = noise_pred * (std_text / std_cfg) + return guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * ( + noise_pred + ) + return self.rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale) + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + self._maybe_compile_dits(server_args) + self._manage_device_placement(self.audio_dit, None, server_args) + + paired_timesteps = batch.paired_timesteps + if paired_timesteps is None: + raise ValueError("paired_timesteps must be set for MOVA") + + y = batch.y if batch.y is not None else batch.image_latent + if getattr(self.video_dit, "require_vae_embedding", False) and y is None: + raise ValueError("MOVA requires reference image latents for denoising") + + boundary_ratio = server_args.pipeline_config.boundary_ratio + total_steps = paired_timesteps.shape[0] + cfg_rank = get_classifier_free_guidance_rank() + enable_cfg_parallel = server_args.enable_cfg_parallel + + is_warmup = batch.is_warmup + extra_step_kwargs = self.prepare_extra_func_kwargs( + self.scheduler.step_from_to, + getattr(batch, "extra_step_kwargs", None) or {}, + ) + + timings = getattr(batch, "timings", None) + perf_dump_path_provided = getattr(batch, "perf_dump_path", None) is not None + + with self.progress_bar(total=total_steps) as progress_bar: + for idx_step in range(total_steps): + with StageProfiler( + f"denoising_step_{idx_step}", + logger=logger, + timings=timings, + perf_dump_path_provided=perf_dump_path_provided, + ): + pair_t = paired_timesteps[idx_step] + if getattr(pair_t, "shape", None) == (2,): + timestep, audio_timestep = pair_t + else: + timestep = pair_t + audio_timestep = pair_t + + cur_visual_dit = self._select_visual_dit( + timestep.item(), boundary_ratio, server_args + ) + + timestep = timestep.unsqueeze(0).to(device=get_local_torch_device()) + audio_timestep = audio_timestep.unsqueeze(0).to( + device=get_local_torch_device() + ) + + attn_metadata = self._build_attn_metadata( + idx_step, batch, server_args + ) + + if not batch.do_classifier_free_guidance: + visual_noise_pred, audio_noise_pred = self._predict( + cur_visual_dit, + batch.latents, + batch.audio_latents, + y, + batch.prompt_embeds[0], + timestep, + audio_timestep, + batch.fps, + idx_step, + attn_metadata, + batch, + ) + else: + if enable_cfg_parallel: + if cfg_rank == 0: + pos = self._predict( + cur_visual_dit, + batch.latents, + batch.audio_latents, + y, + batch.prompt_embeds[0], + timestep, + audio_timestep, + batch.fps, + idx_step, + attn_metadata, + batch, + ) + neg = (None, None) + else: + pos = (None, None) + neg = self._predict( + cur_visual_dit, + batch.latents, + batch.audio_latents, + y, + batch.negative_prompt_embeds[0], + timestep, + audio_timestep, + batch.fps, + idx_step, + attn_metadata, + batch, + ) + else: + pos = self._predict( + cur_visual_dit, + batch.latents, + batch.audio_latents, + y, + batch.prompt_embeds[0], + timestep, + audio_timestep, + batch.fps, + idx_step, + attn_metadata, + batch, + ) + neg = self._predict( + cur_visual_dit, + batch.latents, + batch.audio_latents, + y, + batch.negative_prompt_embeds[0], + timestep, + audio_timestep, + batch.fps, + idx_step, + attn_metadata, + batch, + ) + + visual_noise_pred = self._cfg_combine( + pos[0] if pos[0] is not None else neg[0], + neg[0] if neg[0] is not None else pos[0], + batch.guidance_scale, + cfg_rank, + enable_cfg_parallel, + ) + audio_noise_pred = self._cfg_combine( + pos[1] if pos[1] is not None else neg[1], + neg[1] if neg[1] is not None else pos[1], + batch.guidance_scale, + cfg_rank, + enable_cfg_parallel, + ) + + if batch.guidance_rescale > 0.0: + visual_noise_pred = self._apply_guidance_rescale( + visual_noise_pred, + pos[0] if pos[0] is not None else None, + batch.guidance_rescale, + cfg_rank, + enable_cfg_parallel, + ) + audio_noise_pred = self._apply_guidance_rescale( + audio_noise_pred, + pos[1] if pos[1] is not None else None, + batch.guidance_rescale, + cfg_rank, + enable_cfg_parallel, + ) + + if idx_step + 1 < total_steps: + next_pair_t = paired_timesteps[idx_step + 1] + if getattr(next_pair_t, "shape", None) == (2,): + next_timestep, next_audio_timestep = next_pair_t + else: + next_timestep = next_pair_t + next_audio_timestep = next_pair_t + else: + next_timestep = None + next_audio_timestep = None + + batch.latents = self.scheduler.step_from_to( + visual_noise_pred, + timestep, + next_timestep, + batch.latents, + **extra_step_kwargs, + ) + batch.audio_latents = self.scheduler.step_from_to( + audio_noise_pred, + audio_timestep, + next_audio_timestep, + batch.audio_latents, + **extra_step_kwargs, + ) + + if progress_bar is not None: + progress_bar.update() + if not is_warmup and hasattr(self, "step_profile"): + self.step_profile() + + for dit in filter(None, [self.video_dit, self.video_dit_2, self.audio_dit]): + if isinstance(dit, OffloadableDiTMixin): + dit.prepare_for_next_denoise() + + return batch + + def _shard_sequence_for_sp( + self, x: torch.Tensor, dim: int = 1 + ) -> tuple[torch.Tensor, int]: + """ + Shard tensor along sequence dimension for Sequence Parallelism. + + Args: + x: Input tensor + dim: Dimension to shard along + + Returns: + (sharded_tensor, pad_len) + """ + sp_size = get_sp_world_size() + if sp_size <= 1: + return x, 0 + + sp_rank = get_sp_parallel_rank() + seq_len = x.shape[dim] + + # Pad if needed + pad_len = (sp_size - (seq_len % sp_size)) % sp_size + if pad_len > 0: + pad_shape = list(x.shape) + pad_shape[dim] = pad_len + pad = torch.zeros(pad_shape, dtype=x.dtype, device=x.device) + x = torch.cat([x, pad], dim=dim) + + # Shard + chunk_size = x.shape[dim] // sp_size + start = sp_rank * chunk_size + end = start + chunk_size + idx = [slice(None)] * x.dim() + idx[dim] = slice(start, end) + return x[tuple(idx)], pad_len + + def _gather_sequence_from_sp( + self, x: torch.Tensor, pad_len: int, dim: int = 1 + ) -> torch.Tensor: + """ + Gather tensor along sequence dimension after Sequence Parallelism. + + Args: + x: Sharded tensor + pad_len: Padding length that was added during sharding + dim: Dimension to gather along + + Returns: + Gathered tensor with padding removed + """ + sp_size = get_sp_world_size() + if sp_size <= 1: + return x + + gathered = sequence_model_parallel_all_gather(x, dim=dim) + if pad_len > 0: + idx = [slice(None)] * gathered.dim() + idx[dim] = slice(0, gathered.shape[dim] - pad_len) + gathered = gathered[tuple(idx)] + return gathered + + def inference_single_step( + self, + visual_dit, + visual_latents: torch.Tensor, + audio_latents: torch.Tensor, + y, + context: torch.Tensor, + timestep: torch.Tensor, + audio_timestep: torch.Tensor, + video_fps: float, + ): + """ + Single inference step for MOVA dual-tower denoising. + + Supports Sequence Parallelism (SP): + - After patchify, sequences are sharded across SP ranks + - USPAttention handles distributed attention communication + - Before unpatchify, sequences are gathered back + """ + model_dtype = visual_dit.time_embedding.fc_in.weight.dtype + device = visual_latents.device + + visual_context = context.to(device=device, dtype=model_dtype) + audio_context = context.to(device=device, dtype=model_dtype) + with torch.autocast( + device_type=current_platform.device_type, dtype=torch.float32 + ): + visual_t = visual_dit.time_embedding( + video_sinusoidal_embedding_1d(visual_dit.freq_dim, timestep) + ) + visual_t_mod, _ = visual_dit.time_projection(visual_t) + visual_t_mod = visual_t_mod.unflatten(1, (6, visual_dit.dim)) + + audio_t = self.audio_dit.time_embedding( + audio_sinusoidal_embedding_1d(self.audio_dit.freq_dim, audio_timestep) + ) + audio_t_mod, _ = self.audio_dit.time_projection(audio_t) + audio_t_mod = audio_t_mod.unflatten(1, (6, self.audio_dit.dim)) + + visual_t = visual_t.to(model_dtype) + visual_t_mod = visual_t_mod.to(model_dtype) + audio_t = audio_t.to(model_dtype) + audio_t_mod = audio_t_mod.to(model_dtype) + + visual_context_emb = visual_dit.text_embedding(visual_context) + audio_context_emb = self.audio_dit.text_embedding(audio_context) + + visual_x = visual_latents.to(model_dtype) + audio_x = audio_latents.to(model_dtype) + + if getattr(visual_dit, "require_vae_embedding", False): + visual_x = torch.cat([visual_x, y], dim=1) + + # Patchify visual latents + visual_x, (t, h, w) = visual_dit.patchify(visual_x) + grid_size = (t, h, w) + full_visual_seq_len = t * h * w + + # Build visual freqs for full sequence + visual_dit._init_freqs() + visual_freqs = tuple(freq.to(visual_x.device) for freq in visual_dit.freqs) + visual_freqs = ( + torch.cat( + [ + visual_freqs[0][:t].view(t, 1, 1, -1).expand(t, h, w, -1), + visual_freqs[1][:h].view(1, h, 1, -1).expand(t, h, w, -1), + visual_freqs[2][:w].view(1, 1, w, -1).expand(t, h, w, -1), + ], + dim=-1, + ) + .reshape(full_visual_seq_len, 1, -1) + .to(visual_x.device) + ) + + # Patchify audio latents + audio_x, (f,) = self.audio_dit.patchify(audio_x, None) + full_audio_seq_len = f + + # Build audio freqs for full sequence + self.audio_dit._init_freqs() + audio_freqs = ( + torch.cat( + [ + self.audio_dit.freqs[0][:f].view(f, -1).expand(f, -1), + self.audio_dit.freqs[1][:f].view(f, -1).expand(f, -1), + self.audio_dit.freqs[2][:f].view(f, -1).expand(f, -1), + ], + dim=-1, + ) + .reshape(full_audio_seq_len, 1, -1) + .to(audio_x.device) + ) + + # Shard sequences for SP + visual_x, visual_pad_len = self._shard_sequence_for_sp(visual_x, dim=1) + audio_x, audio_pad_len = self._shard_sequence_for_sp(audio_x, dim=1) + + # Shard freqs to match local sequence length + visual_freqs, _ = self._shard_sequence_for_sp(visual_freqs, dim=0) + audio_freqs, _ = self._shard_sequence_for_sp(audio_freqs, dim=0) + + # Forward through dual-tower DiT + visual_x, audio_x = self.forward_dual_tower_dit( + visual_dit=visual_dit, + visual_x=visual_x, + audio_x=audio_x, + visual_context=visual_context_emb, + audio_context=audio_context_emb, + visual_t_mod=visual_t_mod, + audio_t_mod=audio_t_mod, + visual_freqs=visual_freqs, + audio_freqs=audio_freqs, + grid_size=grid_size, + video_fps=video_fps, + full_visual_seq_len=full_visual_seq_len, + full_audio_seq_len=full_audio_seq_len, + ) + + # Gather sequences back from SP before head/unpatchify + visual_x = self._gather_sequence_from_sp(visual_x, visual_pad_len, dim=1) + audio_x = self._gather_sequence_from_sp(audio_x, audio_pad_len, dim=1) + + visual_output = visual_dit.head(visual_x, visual_t) + visual_output = visual_dit.unpatchify(visual_output, grid_size) + + audio_output = self.audio_dit.head(audio_x, audio_t) + audio_output = self.audio_dit.unpatchify(audio_output, (f,)) + + return visual_output.float(), audio_output.float() + + def forward_dual_tower_dit( + self, + visual_dit, + visual_x: torch.Tensor, + audio_x: torch.Tensor, + visual_context: torch.Tensor, + audio_context: torch.Tensor, + visual_t_mod: torch.Tensor, + audio_t_mod: torch.Tensor, + visual_freqs: torch.Tensor, + audio_freqs: torch.Tensor, + grid_size: tuple[int, int, int], + video_fps: float, + full_visual_seq_len: int, + full_audio_seq_len: int, + condition_scale: float | None = 1.0, + a2v_condition_scale: float | None = None, + v2a_condition_scale: float | None = None, + ): + """ + Forward pass through dual-tower DiT with cross-modal interaction. + + Sequence Parallelism (SP) Support: + - visual_x and audio_x are already sharded along sequence dimension + - visual_freqs and audio_freqs match the local sequence length + - USPAttention in self-attention handles distributed communication + - LocalAttention in cross-attention operates on local sequence vs replicated context + - Cross-modal attention (dual_tower_bridge) uses LocalAttention (no SP communication) + + Args: + full_visual_seq_len: Full visual sequence length before SP sharding + full_audio_seq_len: Full audio sequence length before SP sharding + """ + min_layers = min(len(visual_dit.blocks), len(self.audio_dit.blocks)) + visual_layers = len(visual_dit.blocks) + sp_size = get_sp_world_size() + + # Build RoPE frequencies for cross-attention if needed (only used when SP == 1) + # When SP > 1, we rebuild freqs inside the loop after gathering full sequences + visual_rope_cos_sin, audio_rope_cos_sin = ( + self.dual_tower_bridge.build_aligned_freqs( + video_fps=video_fps, + grid_size=grid_size, + audio_steps=full_audio_seq_len, + device=visual_x.device, + dtype=visual_x.dtype, + ) + ) + if visual_rope_cos_sin is not None: + visual_rope_cos_sin = [ + self._shard_sequence_for_sp(rope_cos_sin, dim=1)[0] + for rope_cos_sin in visual_rope_cos_sin + ] + if audio_rope_cos_sin is not None: + audio_rope_cos_sin = [ + self._shard_sequence_for_sp(rope_cos_sin, dim=1)[0] + for rope_cos_sin in audio_rope_cos_sin + ] + + for layer_idx in range(min_layers): + visual_block = visual_dit.blocks[layer_idx] + audio_block = self.audio_dit.blocks[layer_idx] + + # Cross-modal interaction via dual tower bridge + # Bridge operations (PerFrameAttentionPooling, RoPE) expect full sequences + # When SP is enabled, we need to gather before bridge and shard after + if self.dual_tower_bridge.should_interact(layer_idx, "a2v"): + visual_x, audio_x = self.dual_tower_bridge( + layer_idx, + visual_x, + audio_x, + x_freqs=visual_rope_cos_sin, + y_freqs=audio_rope_cos_sin, + a2v_condition_scale=a2v_condition_scale, + v2a_condition_scale=v2a_condition_scale, + condition_scale=condition_scale, + video_grid_size=grid_size, + ) + + # Self-attention and FFN in DiT blocks + visual_x = visual_block( + visual_x, visual_context, visual_t_mod, visual_freqs + ) + audio_x = audio_block(audio_x, audio_context, audio_t_mod, audio_freqs) + + # Process remaining visual layers (if visual has more layers than audio) + for layer_idx in range(min_layers, visual_layers): + visual_block = visual_dit.blocks[layer_idx] + visual_x = visual_block( + visual_x, visual_context, visual_t_mod, visual_freqs + ) + + return visual_x, audio_x + + +class MOVADecodingStage(PipelineStage): + """Decode video and audio outputs for MOVA.""" + + def __init__(self, video_vae, audio_vae) -> None: + super().__init__() + self.video_vae = video_vae + self.audio_vae = audio_vae + + @property + def parallelism_type(self) -> StageParallelismType: + if get_global_server_args().enable_cfg_parallel: + return StageParallelismType.MAIN_RANK_ONLY + return StageParallelismType.REPLICATED + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: + self.video_vae = self.video_vae.to(get_local_torch_device()) + self.audio_vae = self.audio_vae.to(get_local_torch_device()) + + video_latents = server_args.pipeline_config.denormalize_video_latents( + batch.latents, self.video_vae + ) + + vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + vae_autocast_enabled = ( + vae_dtype != torch.float32 + ) and not server_args.disable_autocast + + with torch.autocast( + device_type=current_platform.device_type, + dtype=vae_dtype, + enabled=vae_autocast_enabled, + ): + if server_args.pipeline_config.vae_tiling: + self.video_vae.enable_tiling() + if not vae_autocast_enabled: + video_latents = video_latents.to(vae_dtype) + decode_output = self.video_vae.decode(video_latents) + video = _ensure_tensor_decode_output(decode_output) + + video = (video / 2 + 0.5).clamp(0, 1) + + with torch.autocast( + device_type=current_platform.device_type, dtype=torch.float32 + ): + audio = self.audio_vae.decode(batch.audio_latents) + output_batch = OutputBatch( + output=video, + audio=audio, + audio_sample_rate=getattr(self.audio_vae, "sample_rate", None), + timings=batch.timings, + ) + return output_batch diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/flow_match_pair.py b/python/sglang/multimodal_gen/runtime/models/schedulers/flow_match_pair.py new file mode 100644 index 000000000..679399ab7 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/flow_match_pair.py @@ -0,0 +1,501 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copied and adapted from: mossVG/mova/diffusion/schedulers/flow_match.py and flow_match_pair.py + +from __future__ import annotations + +import math + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin + +from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler + + +class FlowMatchScheduler(BaseScheduler): + def __init__( + self, + num_inference_steps=100, + num_train_timesteps=1000, + shift=3.0, + sigma_max=1.0, + sigma_min=0.003 / 1.002, + inverse_timesteps=False, + extra_one_step=False, + reverse_sigmas=False, + exponential_shift=False, + exponential_shift_mu=None, + shift_terminal=None, + ): + self.order = 1 + self.num_train_timesteps = num_train_timesteps + self.shift = shift + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self.inverse_timesteps = inverse_timesteps + self.extra_one_step = extra_one_step + self.reverse_sigmas = reverse_sigmas + self.exponential_shift = exponential_shift + self.exponential_shift_mu = exponential_shift_mu + self.shift_terminal = shift_terminal + self.train_timesteps = None + self.train_sigmas = None + self.set_timesteps(num_train_timesteps) + self.set_timesteps(num_inference_steps) + BaseScheduler.__init__(self) + + def set_shift(self, shift: float) -> None: + self.shift = shift + + def set_timesteps( + self, + num_inference_steps=100, + denoising_strength=1.0, + training=False, + shift=None, + dynamic_shift_len=None, + ): + if shift is not None: + self.shift = shift + sigma_start = ( + self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength + ) + if self.extra_one_step: + self.sigmas = torch.linspace( + sigma_start, self.sigma_min, num_inference_steps + 1 + )[:-1] + else: + self.sigmas = torch.linspace( + sigma_start, self.sigma_min, num_inference_steps + ) + if self.inverse_timesteps: + self.sigmas = torch.flip(self.sigmas, dims=[0]) + if self.exponential_shift: + mu = ( + self.calculate_shift(dynamic_shift_len) + if dynamic_shift_len is not None + else self.exponential_shift_mu + ) + self.sigmas = math.exp(mu) / (math.exp(mu) + (1 / self.sigmas - 1)) + else: + self.sigmas = ( + self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) + ) + if self.shift_terminal is not None: + one_minus_z = 1 - self.sigmas + scale_factor = one_minus_z[-1] / (1 - self.shift_terminal) + self.sigmas = 1 - (one_minus_z / scale_factor) + if self.reverse_sigmas: + self.sigmas = 1 - self.sigmas + self.timesteps = self.sigmas * self.num_train_timesteps + # 第一次设置 train_timesteps + if self.train_timesteps is None: + self.train_timesteps = self.timesteps + self.train_sigmas = self.sigmas + if training: + x = self.timesteps + y = torch.exp( + -2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2 + ) + y_shifted = y - y.min() + bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum()) + self.linear_timesteps_weights = bsmntw_weighing + self.training = True + else: + self.training = False + + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None): + return sample + + def step(self, model_output, timestep, sample, to_final=False, **kwargs): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + if to_final or timestep_id + 1 >= len(self.timesteps): + sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0 + else: + sigma_ = self.sigmas[timestep_id + 1] + prev_sample = sample + model_output * (sigma_ - sigma) + return prev_sample + + def return_to_timestep(self, timestep, sample, sample_stablized): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + model_output = (sample - sample_stablized) / sigma + return model_output + + def add_noise(self, original_samples, noise, timestep): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + sample = (1 - sigma) * original_samples + sigma * noise + return sample + + def training_target(self, sample, noise, timestep): + target = noise - sample + return target + + def training_weight(self, timestep): + timestep_id = torch.argmin( + (self.timesteps - timestep.to(self.timesteps.device)).abs() + ) + weights = self.linear_timesteps_weights[timestep_id] + return weights + + def calculate_shift( + self, + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 8192, + base_shift: float = 0.5, + max_shift: float = 0.9, + ): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +class FlowMatchPairScheduler(FlowMatchScheduler, SchedulerMixin, ConfigMixin): + """ + 在 FlowMatchScheduler 的基础上，提供便捷的配对接口： + - 默认返回形状为 [num_timesteps, 2] 的张量，每一行为 (t, t) + - 允许通过 set_pair_postprocess(fn) 设置一个后处理函数以修改配对行为 + """ + + @register_to_config + def __init__( + self, + num_inference_steps=100, + num_train_timesteps=1000, + shift=3.0, + sigma_max=1.0, + sigma_min=0.003 / 1.002, + inverse_timesteps=False, + extra_one_step=False, + reverse_sigmas=False, + exponential_shift=False, + exponential_shift_mu=None, + shift_terminal=None, + ): + self._pair_postprocess_fn = None + self._pair_postprocess_requires_source = False + self.pair_timesteps: torch.Tensor | None = None + self.pair_sigmas: torch.Tensor | None = None + self.timesteps: torch.Tensor | None = None + self.sigmas: torch.Tensor | None = None + super().__init__( + num_inference_steps=num_inference_steps, + num_train_timesteps=num_train_timesteps, + shift=shift, + sigma_max=sigma_max, + sigma_min=sigma_min, + inverse_timesteps=inverse_timesteps, + extra_one_step=extra_one_step, + reverse_sigmas=reverse_sigmas, + exponential_shift=exponential_shift, + exponential_shift_mu=exponential_shift_mu, + shift_terminal=shift_terminal, + ) + + def set_pair_postprocess(self, fn): + """ + 设置一个后处理函数，用于在默认配对生成后进行自定义修改。 + 要求： + - fn(pairs: torch.Tensor) -> torch.Tensor + - 返回的张量必须与输入 pairs 形状一致，否则直接 raise。 + """ + if fn is not None and not callable(fn): + raise TypeError("pair_postprocess 必须是可调用对象或 None") + self._pair_postprocess_fn = fn + self._pair_postprocess_requires_source = ( + False if fn is None else bool(getattr(fn, "_requires_source", False)) + ) + if self.timesteps is None or self.sigmas is None: + raise RuntimeError("调度器未初始化，请先调用 set_timesteps()") + self._refresh_pair_cache() + + def set_pair_postprocess_by_name(self, name: str | None, **kwargs): + """ + 通过名称快速配置后处理函数。 + 支持： + - None/"none"/"off"/"false"/"no": 关闭 + - "quadratic_perp_bulge_swap": x2=x+d, y2=x-d，其中 d=4*amp*s*(1-s), s=t/T + - "v2a_sequential": 假设原 pairs 为 (t,t)，用列0步长2采样一半序列，先让列0按该序列变化，再让列1按该序列变化 + - "a2v_sequential": 同上，先列1后列0 + - "dual_sigma_shift": 仅使用 timestep 数量，重新按照 FlowMatchScheduler 的 sigma 变换逻辑为两列独立构建调度，可配置 visual_shift/audio_shift + + 额外参数： + - amp: 浮点幅度，默认 150.0 + """ + if name is None or str(name).lower() in ("none", "off", "false", "no"): + self.set_pair_postprocess(None) + return + if name == "quadratic_perp_bulge_swap": + amp = float(kwargs.get("amp", 150.0)) + + def _quadratic_perp_bulge_swap(pairs: torch.Tensor): + if ( + not isinstance(pairs, torch.Tensor) + or pairs.ndim != 2 + or pairs.shape[1] != 2 + ): + raise ValueError("pairs 必须是形状 [N, 2] 的 torch.Tensor") + x = pairs[:, 0] + T = float(self.num_train_timesteps) + s = x / T + d = 4.0 * amp * s * (1.0 - s) + x2 = x + d + y2 = x - d + return torch.stack([x2, y2], dim=1) + + self.set_pair_postprocess(_quadratic_perp_bulge_swap) + return + if name == "v2a_sequential": + + def _v2a(pairs: torch.Tensor): + if ( + not isinstance(pairs, torch.Tensor) + or pairs.ndim != 2 + or pairs.shape[1] != 2 + ): + raise ValueError("pairs 必须是形状 [N, 2] 的 torch.Tensor") + N = pairs.shape[0] + base = pairs[:, 0] + seq_half = base[::2] + m = int(seq_half.shape[0]) + col0 = torch.cat([seq_half, seq_half[-1:].repeat(m)], dim=0)[:N] + col1 = torch.cat([seq_half[0:1].repeat(m), seq_half], dim=0)[:N] + return torch.stack( + [ + col0.to(dtype=pairs.dtype, device=pairs.device), + col1.to(dtype=pairs.dtype, device=pairs.device), + ], + dim=1, + ) + + self.set_pair_postprocess(_v2a) + return + if name == "a2v_sequential": + + def _a2v(pairs: torch.Tensor): + if ( + not isinstance(pairs, torch.Tensor) + or pairs.ndim != 2 + or pairs.shape[1] != 2 + ): + raise ValueError("pairs 必须是形状 [N, 2] 的 torch.Tensor") + N = pairs.shape[0] + base = pairs[:, 0] + seq_half = base[::2] + m = int(seq_half.shape[0]) + col0 = torch.cat([seq_half[0:1].repeat(m), seq_half], dim=0)[:N] + col1 = torch.cat([seq_half, seq_half[-1:].repeat(m)], dim=0)[:N] + return torch.stack( + [ + col0.to(dtype=pairs.dtype, device=pairs.device), + col1.to(dtype=pairs.dtype, device=pairs.device), + ], + dim=1, + ) + + self.set_pair_postprocess(_a2v) + return + if name == "v2a": + + def _v2a_classic(pairs: torch.Tensor): + if ( + not isinstance(pairs, torch.Tensor) + or pairs.ndim != 2 + or pairs.shape[1] != 2 + ): + raise ValueError("pairs 必须是形状 [N, 2] 的 torch.Tensor") + zeros = torch.zeros_like(pairs[:, 0]) + return torch.stack([zeros, pairs[:, 1]], dim=1) + + self.set_pair_postprocess(_v2a_classic) + return + if name == "a2v": + + def _a2v_classic(pairs: torch.Tensor): + if ( + not isinstance(pairs, torch.Tensor) + or pairs.ndim != 2 + or pairs.shape[1] != 2 + ): + raise ValueError("pairs 必须是形状 [N, 2] 的 torch.Tensor") + zeros = torch.zeros_like(pairs[:, 1]) + return torch.stack([pairs[:, 0], zeros], dim=1) + + self.set_pair_postprocess(_a2v_classic) + return + if name == "dual_sigma_shift": + visual_shift = float(kwargs.get("visual_shift", self.shift)) + audio_shift = float(kwargs.get("audio_shift", self.shift)) + visual_denoising_strength = float( + kwargs.get("visual_denoising_strength", 1.0) + ) + audio_denoising_strength = float( + kwargs.get("audio_denoising_strength", 1.0) + ) + visual_mu = kwargs.get( + "visual_exponential_shift_mu", self.exponential_shift_mu + ) + audio_mu = kwargs.get( + "audio_exponential_shift_mu", self.exponential_shift_mu + ) + + def _dual_sigma_shift(pairs: torch.Tensor, *, source: str): + if not isinstance(pairs, torch.Tensor): + raise TypeError("pairs 必须是 torch.Tensor") + if pairs.ndim != 2 or pairs.shape[1] != 2: + raise ValueError("pairs 必须是形状 [N, 2] 的 torch.Tensor") + if pairs.shape[0] == 0: + raise ValueError("pairs 的长度必须大于 0") + if source not in ("timesteps", "sigmas"): + raise ValueError("source 仅支持 'timesteps' 或 'sigmas'") + + num_steps = pairs.shape[0] + device = pairs.device + dtype = pairs.dtype + + def _build_column( + shift_value: float, denoising_strength: float, mu_override + ): + if shift_value <= 0: + raise ValueError("shift 必须为正数") + if denoising_strength <= 0: + raise ValueError("denoising_strength 必须为正数") + + sigma_start = ( + self.sigma_min + + (self.sigma_max - self.sigma_min) * denoising_strength + ) + if self.extra_one_step: + base = torch.linspace( + sigma_start, + self.sigma_min, + num_steps + 1, + device=device, + dtype=dtype, + )[:-1] + else: + base = torch.linspace( + sigma_start, + self.sigma_min, + num_steps, + device=device, + dtype=dtype, + ) + + if self.inverse_timesteps: + base = torch.flip(base, dims=[0]) + + if self.exponential_shift: + mu_value = mu_override + if mu_value is None: + raise RuntimeError( + "启用了 exponential_shift 但未提供 exponential_shift_mu" + ) + exp_mu = math.exp(float(mu_value)) + base = exp_mu / (exp_mu + (1 / base - 1)) + else: + base = shift_value * base / (1 + (shift_value - 1) * base) + + if self.shift_terminal is not None: + one_minus_z = 1 - base + scale_factor = one_minus_z[-1] / (1 - self.shift_terminal) + base = 1 - (one_minus_z / scale_factor) + + if self.reverse_sigmas: + base = 1 - base + + if source == "timesteps": + return base * self.num_train_timesteps + return base + + col0 = _build_column(visual_shift, visual_denoising_strength, visual_mu) + col1 = _build_column(audio_shift, audio_denoising_strength, audio_mu) + return torch.stack([col0, col1], dim=1) + + _dual_sigma_shift._requires_source = True + self.set_pair_postprocess(_dual_sigma_shift) + return + raise ValueError(f"Unknown pair_postprocess name: {name}") + + def _make_pairs_from_vector(self, vec: torch.Tensor) -> torch.Tensor: + if vec.ndim != 1: + raise ValueError("vec must be 1D") + return torch.stack([vec, vec], dim=1) + + def get_pairs(self, source: str = "timesteps") -> torch.Tensor: + if source == "timesteps": + if self.pair_timesteps is None: + self._refresh_pair_cache() + return self.pair_timesteps + if source == "sigmas": + if self.pair_sigmas is None: + self._refresh_pair_cache() + return self.pair_sigmas + raise ValueError("source must be 'timesteps' or 'sigmas'") + + def timestep_to_sigma(self, timestep: torch.Tensor | float) -> torch.Tensor: + """根据给定的 timestep（标量）返回对应的 sigma（按最近邻在 self.timesteps 中查找）。""" + t_value = float(timestep) + t_cpu = torch.tensor(t_value) + idx = torch.argmin((self.train_timesteps - t_cpu).abs()) + return self.train_sigmas[idx] + + def step_from_to( + self, + model_output: torch.Tensor, + timestep_from: torch.Tensor, + timestep_to: torch.Tensor | None, + sample: torch.Tensor, + ) -> torch.Tensor: + """ + 使用显式给定的 (from, to) timestep 对，按照对应的 sigma 差进行一步更新： + x_{to} = x_{from} + model_output * (sigma(to) - sigma(from)) + 该方法可用于两个模态分别沿着它们各自列序列推进。 + """ + sigma_from = self.timestep_to_sigma(timestep_from) + if timestep_to is None: + sigma_to = torch.tensor( + 1.0 if (self.inverse_timesteps or self.reverse_sigmas) else 0.0, + device=sigma_from.device, + dtype=sigma_from.dtype, + ) + else: + sigma_to = self.timestep_to_sigma(timestep_to) + prev_sample = sample + model_output * (sigma_to - sigma_from) + return prev_sample + + def _refresh_pair_cache(self) -> None: + if self.timesteps is None or self.sigmas is None: + raise RuntimeError("调度器未初始化，请先调用 set_timesteps()") + + def _apply_postprocess(pairs: torch.Tensor, source: str) -> torch.Tensor: + if self._pair_postprocess_fn is None: + return pairs + if self._pair_postprocess_requires_source: + modified = self._pair_postprocess_fn(pairs, source=source) + else: + modified = self._pair_postprocess_fn(pairs) + if not isinstance(modified, torch.Tensor): + raise TypeError("pair_postprocess 返回值必须是 torch.Tensor") + if modified.shape != pairs.shape: + raise ValueError("pair_postprocess 返回的张量形状必须与输入一致") + return modified + + base_pairs_timesteps = self._make_pairs_from_vector(self.timesteps) + base_pairs_sigmas = self._make_pairs_from_vector(self.sigmas) + + self.pair_timesteps = _apply_postprocess(base_pairs_timesteps, "timesteps") + self.pair_sigmas = _apply_postprocess(base_pairs_sigmas, "sigmas") + + +EntryClass = FlowMatchPairScheduler diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/common.py b/python/sglang/multimodal_gen/runtime/models/vaes/common.py index d3011b623..095ce4957 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/common.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/common.py @@ -611,7 +611,9 @@ class DiagonalGaussianDistribution: return x def kl( - self, other: Optional["DiagonalGaussianDistribution"] = None + self, + other: Optional["DiagonalGaussianDistribution"] = None, + dims: tuple[int, ...] = (1, 2, 3), ) -> torch.Tensor: if self.deterministic: return torch.Tensor([0.0]) @@ -619,7 +621,7 @@ class DiagonalGaussianDistribution: if other is None: return 0.5 * torch.sum( torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, - dim=[1, 2, 3], + dim=dims, ) else: return 0.5 * torch.sum( @@ -628,7 +630,7 @@ class DiagonalGaussianDistribution: - 1.0 - self.logvar + other.logvar, - dim=[1, 2, 3], + dim=dims, ) def nll( diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/dac.py b/python/sglang/multimodal_gen/runtime/models/vaes/dac.py new file mode 100644 index 000000000..3d6d821ab --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/vaes/dac.py @@ -0,0 +1,627 @@ +# Copied and adapted from: https://github.com/descriptinc/descript-audio-codec + +# SPDX-License-Identifier: MIT + +import math +from bisect import bisect_right +from typing import Union + +import torch +import torch.nn.functional as F +from einops import rearrange +from torch import nn + +from sglang.multimodal_gen.configs.models.vaes.dac import DacVAEConfig +from sglang.multimodal_gen.runtime.models.vaes.common import ( + DiagonalGaussianDistribution, +) + + +# Scripting this brings model speed up 1.4x +@torch.jit.script +def snake(x, alpha): + shape = x.shape + x = x.reshape(shape[0], shape[1], -1) + x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2) + x = x.reshape(shape) + return x + + +class Snake1d(nn.Module): + def __init__(self, channels): + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, x): + return snake(x, self.alpha) + + +class VectorQuantize(nn.Module): + """ + Implementation of VQ similar to Karpathy's repo: + https://github.com/karpathy/deep-vector-quantization + Additionally uses following tricks from Improved VQGAN + (https://arxiv.org/pdf/2110.04627.pdf): + 1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space + for improved codebook usage + 2. l2-normalized codes: Converts euclidean distance to cosine similarity which + improves training stability + """ + + def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int): + super().__init__() + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + + self.in_proj = nn.Conv1d(input_dim, codebook_dim, kernel_size=1) + self.out_proj = nn.Conv1d(codebook_dim, input_dim, kernel_size=1) + self.codebook = nn.Embedding(codebook_size, codebook_dim) + + def forward(self, z): + """Quantize the input tensor using a fixed codebook and return the corresponding codebook vectors. + + Args: + z (torch.Tensor): Input tensor with shape ``[B, D, T]``. + + Returns: + tuple: A tuple containing: + - z_q (torch.Tensor): Quantized continuous representation with shape ``[B, D, T]``. + - commitment_loss (torch.Tensor): Commitment loss scalar to train encoder to predict + vectors closer to codebook entries. + - codebook_loss (torch.Tensor): Codebook loss scalar to update the codebook. + - indices (torch.Tensor): Codebook indices (quantized discrete representation) with shape ``[B, T]``. + - z_e (torch.Tensor): Projected latents (continuous representation before quantization) with shape ``[B, D, T]``. + """ + + # Factorized codes (ViT-VQGAN) Project input into low-dimensional space + z_e = self.in_proj(z) # z_e : (B x D x T) + z_q, indices = self.decode_latents(z_e) + + commitment_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2]) + codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2]) + + z_q = ( + z_e + (z_q - z_e).detach() + ) # noop in forward pass, straight-through gradient estimator in backward pass + + z_q = self.out_proj(z_q) + + return z_q, commitment_loss, codebook_loss, indices, z_e + + def embed_code(self, embed_id): + return F.embedding(embed_id, self.codebook.weight) + + def decode_code(self, embed_id): + return self.embed_code(embed_id).transpose(1, 2) + + def decode_latents(self, latents): + encodings = rearrange(latents, "b d t -> (b t) d") + codebook = self.codebook.weight # codebook: (N x D) + + # L2 normalize encodings and codebook (ViT-VQGAN) + encodings = F.normalize(encodings) + codebook = F.normalize(codebook) + + # Compute euclidean distance with codebook + dist = ( + encodings.pow(2).sum(1, keepdim=True) + - 2 * encodings @ codebook.t() + + codebook.pow(2).sum(1, keepdim=True).t() + ) + indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0)) + z_q = self.decode_code(indices) + return z_q, indices + + +class ResidualVectorQuantize(nn.Module): + """ + Introduced in SoundStream: An end2end neural audio codec + https://arxiv.org/abs/2107.03312 + """ + + def __init__( + self, + input_dim: int = 512, + n_codebooks: int = 9, + codebook_size: int = 1024, + codebook_dim: Union[int, list] = 8, + quantizer_dropout: float = 0.0, + ): + super().__init__() + if isinstance(codebook_dim, int): + codebook_dim = [codebook_dim for _ in range(n_codebooks)] + + self.n_codebooks = n_codebooks + self.codebook_dim = codebook_dim + self.codebook_size = codebook_size + dim_offsets = [0] + for dim in self.codebook_dim: + dim_offsets.append(dim_offsets[-1] + dim) + self._codebook_dim_offsets = tuple(dim_offsets) + + self.quantizers = nn.ModuleList( + [ + VectorQuantize(input_dim, codebook_size, codebook_dim[i]) + for i in range(n_codebooks) + ] + ) + self.quantizer_dropout = quantizer_dropout + + def forward(self, z, n_quantizers: int = None): + """Quantize the input tensor using a fixed set of codebooks and return the corresponding codebook vectors. + + Args: + z (torch.Tensor): Input tensor with shape ``[B, D, T]``. + n_quantizers (int, optional): Number of quantizers to use. If ``None``, + all quantizers are used. When ``n_quantizers`` < ``self.n_codebooks``, + quantizer dropout is applied. Note: if ``self.quantizer_dropout`` > 0 + and in training mode, this argument is ignored and a random number of + quantizers is used. + + Returns: + tuple: A tuple containing: + - z_q (torch.Tensor): Quantized continuous representation with shape ``[B, D, T]``. + - codes (torch.Tensor): Codebook indices for each codebook with shape ``[B, N, T]`` + (quantized discrete representation of input). + - latents (torch.Tensor): Projected latents with shape ``[B, N*D, T]`` + (continuous representation before quantization). + - commitment_loss (torch.Tensor): Commitment loss scalar to train encoder to predict + vectors closer to codebook entries. + - codebook_loss (torch.Tensor): Codebook loss scalar to update the codebook. + """ + z_q = 0 + residual = z + commitment_loss = 0 + codebook_loss = 0 + + codebook_indices = [] + latents = [] + + if n_quantizers is None: + n_quantizers = self.n_codebooks + quantizers = self.quantizers + if self.training: + batch_size = z.shape[0] + device = z.device + n_quantizers = torch.full( + (batch_size,), + self.n_codebooks + 1, + device=device, + dtype=torch.long, + ) + if self.quantizer_dropout > 0: + dropout = torch.randint( + 1, + self.n_codebooks + 1, + (batch_size,), + device=device, + ) + n_dropout = int(batch_size * self.quantizer_dropout) + if n_dropout > 0: + n_quantizers[:n_dropout] = dropout[:n_dropout] + + for i, quantizer in enumerate(quantizers): + z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer( + residual + ) + + # Create mask to apply quantizer dropout + mask = i < n_quantizers + z_q = z_q + z_q_i * mask[:, None, None] + residual = residual - z_q_i + + # Sum losses + commitment_loss += (commitment_loss_i * mask).mean() + codebook_loss += (codebook_loss_i * mask).mean() + + codebook_indices.append(indices_i) + latents.append(z_e_i) + else: + for i, quantizer in enumerate(quantizers): + if i >= n_quantizers: + break + z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer( + residual + ) + z_q = z_q + z_q_i + residual = residual - z_q_i + + commitment_loss += commitment_loss_i.mean() + codebook_loss += codebook_loss_i.mean() + + codebook_indices.append(indices_i) + latents.append(z_e_i) + + codes = torch.stack(codebook_indices, dim=1) + latents = torch.cat(latents, dim=1) + + return z_q, codes, latents, commitment_loss, codebook_loss + + def from_codes(self, codes: torch.Tensor): + """Reconstruct the continuous representation from quantized codes. + + Args: + codes (torch.Tensor): Quantized discrete representation with shape ``[B, N, T]``. + + Returns: + tuple: A tuple containing: + - z_q (torch.Tensor): Quantized continuous representation with shape ``[B, D, T]``. + - z_p (torch.Tensor): Concatenated latent space representation with shape ``[B, N*D, T]``. + - codes (torch.Tensor): Original input codebook indices with shape ``[B, N, T]``. + """ + z_q = 0.0 + z_p = [] + n_codebooks = codes.shape[1] + for i in range(n_codebooks): + z_p_i = self.quantizers[i].decode_code(codes[:, i, :]) + z_p.append(z_p_i) + + z_q_i = self.quantizers[i].out_proj(z_p_i) + z_q = z_q + z_q_i + return z_q, torch.cat(z_p, dim=1), codes + + def from_latents(self, latents: torch.Tensor): + """Reconstruct the continuous representation from unquantized latents. + + Args: + latents (torch.Tensor): Continuous representation after projection with shape ``[B, N*D, T]``. + + Returns: + tuple: A tuple containing: + - z_q (torch.Tensor): Quantized representation of full-projected space with shape ``[B, D, T]``. + - z_p (torch.Tensor): Quantized representation of latent space with shape ``[B, N*D, T]``. + - codes (torch.Tensor): Codebook indices with shape ``[B, N, T]``. + """ + z_q = 0 + z_p = [] + codes = [] + dims = self._codebook_dim_offsets + n_codebooks = bisect_right(dims, latents.shape[1]) - 1 + for i in range(n_codebooks): + j, k = dims[i], dims[i + 1] + z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :]) + z_p.append(z_p_i) + codes.append(codes_i) + + z_q_i = self.quantizers[i].out_proj(z_p_i) + z_q = z_q + z_q_i + + return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1) + + +class ResidualUnit(nn.Module): + def __init__(self, dim: int = 16, dilation: int = 1): + super().__init__() + pad = ((7 - 1) * dilation) // 2 + self.block = nn.Sequential( + Snake1d(dim), + nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad), + Snake1d(dim), + nn.Conv1d(dim, dim, kernel_size=1), + ) + + def forward(self, x): + y = self.block(x) + pad = (x.shape[-1] - y.shape[-1]) // 2 + if pad > 0: + x = x[..., pad:-pad] + return x + y + + +class EncoderBlock(nn.Module): + def __init__(self, dim: int = 16, stride: int = 1): + super().__init__() + self.block = nn.Sequential( + ResidualUnit(dim // 2, dilation=1), + ResidualUnit(dim // 2, dilation=3), + ResidualUnit(dim // 2, dilation=9), + Snake1d(dim // 2), + nn.Conv1d( + dim // 2, + dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ), + ) + + def forward(self, x): + return self.block(x) + + +class Encoder(nn.Module): + def __init__( + self, + d_model: int = 64, + strides: list = [2, 4, 8, 8], + d_latent: int = 64, + ): + super().__init__() + # Create first convolution + self.block = [nn.Conv1d(1, d_model, kernel_size=7, padding=3)] + + # Create EncoderBlocks that double channels as they downsample by `stride` + for stride in strides: + d_model *= 2 + self.block += [EncoderBlock(d_model, stride=stride)] + + # Create last convolution + self.block += [ + Snake1d(d_model), + nn.Conv1d(d_model, d_latent, kernel_size=3, padding=1), + ] + + # Wrap black into nn.Sequential + self.block = nn.Sequential(*self.block) + self.enc_dim = d_model + + def forward(self, x): + return self.block(x) + + +class DecoderBlock(nn.Module): + def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1): + super().__init__() + self.block = nn.Sequential( + Snake1d(input_dim), + nn.ConvTranspose1d( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + output_padding=stride % 2, + ), + ResidualUnit(output_dim, dilation=1), + ResidualUnit(output_dim, dilation=3), + ResidualUnit(output_dim, dilation=9), + ) + + def forward(self, x): + return self.block(x) + + +class Decoder(nn.Module): + def __init__( + self, + input_channel, + channels, + rates, + d_out: int = 1, + ): + super().__init__() + + # Add first conv layer + layers = [nn.Conv1d(input_channel, channels, kernel_size=7, padding=3)] + + # Add upsampling + MRF blocks + for i, stride in enumerate(rates): + input_dim = channels // 2**i + output_dim = channels // 2 ** (i + 1) + layers += [DecoderBlock(input_dim, output_dim, stride)] + + # Add final conv layer + layers += [ + Snake1d(output_dim), + nn.Conv1d(output_dim, d_out, kernel_size=7, padding=3), + nn.Tanh(), + ] + + self.model = nn.Sequential(*layers) + + def forward(self, x): + return self.model(x) + + +class DAC(nn.Module): + def __init__( + self, + config: DacVAEConfig, + ): + super().__init__() + + self.continuous = config.continuous + self.decoder_dim = config.decoder_dim + self.decoder_rates = config.decoder_rates + self.encoder_dim = config.encoder_dim + self.encoder_rates = config.encoder_rates + self.hop_length = math.prod(config.encoder_rates) + self.sample_rate = config.sample_rate + + if config.latent_dim is None: + latent_dim = config.encoder_dim * (2 ** len(config.encoder_rates)) + else: + latent_dim = config.latent_dim + + self.latent_dim = latent_dim + + if config.load_encoder: + self.encoder = Encoder(config.encoder_dim, config.encoder_rates, latent_dim) + + if not config.continuous: + self.n_codebooks = config.n_codebooks + self.codebook_size = config.codebook_size + self.codebook_dim = config.codebook_dim + self.quantizer = ResidualVectorQuantize( + input_dim=latent_dim, + n_codebooks=config.n_codebooks, + codebook_size=config.codebook_size, + codebook_dim=config.codebook_dim, + quantizer_dropout=config.quantizer_dropout, + ) + else: + self.quant_conv = torch.nn.Conv1d(latent_dim, 2 * latent_dim, 1) + self.post_quant_conv = torch.nn.Conv1d(latent_dim, latent_dim, 1) + + if config.load_decoder: + self.decoder = Decoder( + latent_dim, + config.decoder_dim, + config.decoder_rates, + ) + + self.apply(self.init_weights) + + @staticmethod + def init_weights(m): + if isinstance(m, nn.Conv1d): + nn.init.trunc_normal_(m.weight, std=0.02) + nn.init.constant_(m.bias, 0) + + @property + def dtype(self): + return next(self.parameters()).dtype + + @property + def device(self): + return next(self.parameters()).device + + def preprocess(self, audio_data, sample_rate): + if sample_rate is None: + sample_rate = self.sample_rate + assert sample_rate == self.sample_rate + + length = audio_data.shape[-1] + right_pad = math.ceil(length / self.hop_length) * self.hop_length - length + audio_data = nn.functional.pad(audio_data, (0, right_pad)) + + return audio_data + + def encode( + self, + audio_data: torch.Tensor, + n_quantizers: int = None, + ): + """Encode audio data into latent representations. + + This method processes audio through the encoder network and optionally applies + vector quantization (in VQ mode) or projects to a Gaussian distribution (in + continuous mode) to produce latent representations. + + Args: + audio_data (torch.Tensor): Audio data to encode, with shape ``[B, 1, T]``. + n_quantizers (int, optional): Number of quantizers to use. If ``None``, + all quantizers are used. Only applicable in VQ mode (``continuous=False``). + + Returns: + tuple: A tuple containing: + - z (torch.Tensor): Encoded representation. In VQ mode, this is the + quantized continuous representation with shape ``[B, D, T]``. In + continuous mode, this is a ``DiagonalGaussianDistribution`` object. + - codes (torch.Tensor or None): Codebook indices with shape ``[B, N, T]`` + in VQ mode, ``None`` in continuous mode. + - latents (torch.Tensor or None): Projected latents with shape ``[B, N*D, T]`` + in VQ mode, ``None`` in continuous mode. + - commitment_loss (torch.Tensor): Commitment loss scalar. + - codebook_loss (torch.Tensor): Codebook loss scalar. + + Note: + In continuous mode, the encoded representation is projected through a + quantization convolution layer and wrapped in a ``DiagonalGaussianDistribution`` + for VAE training. + """ + z = self.encoder(audio_data) # [B x D x T] + if not self.continuous: + z, codes, latents, commitment_loss, codebook_loss = self.quantizer( + z, n_quantizers + ) + else: + z = self.quant_conv(z) # [B x 2D x T] + z = DiagonalGaussianDistribution(z) + codes, latents, commitment_loss, codebook_loss = None, None, 0, 0 + + return z, codes, latents, commitment_loss, codebook_loss + + def decode(self, z: torch.Tensor): + """Decode latent representations back to audio waveforms. + + This method takes latent representations (either quantized from VQ mode or sampled + from the posterior in continuous mode) and reconstructs the corresponding audio + through the decoder network. + + Args: + z (torch.Tensor): Latent representation to decode, with shape ``[B, D, T]``. + In VQ mode (``continuous=False``), this is the quantized continuous + representation. In continuous mode (``continuous=True``), this is sampled + from the posterior distribution. + + Returns: + torch.Tensor: Decoded audio data with shape ``[B, 1, T']``. The output length + T' is determined by the decoder's upsampling rates and may differ from the + input temporal dimension T. + + Note: + In continuous mode (``continuous=True``), the input is first passed through + a post-quantization convolution layer before being fed to the decoder. + """ + if not self.continuous: + audio = self.decoder(z) + else: + z = self.post_quant_conv(z) + audio = self.decoder(z) + + return audio + + def forward( + self, + audio_data: torch.Tensor, + sample_rate: int = None, + n_quantizers: int = None, + ): + """Model forward pass. + + Args: + audio_data (torch.Tensor): Audio to encode, shape [B, 1, T]. + sample_rate (int, optional): Sample rate in Hz. Defaults to + ``self.sample_rate`` when ``None``. + n_quantizers (int, optional): Number of quantizers to use. When ``None``, + all quantizers are used. Only used in VQ mode (``continuous=False``). + + Returns: + dict: A dictionary containing different keys depending on the mode: + + **VQ Mode (``continuous=False``):** + - "audio" (torch.Tensor): Decoded audio, shape [B, 1, length]. + - "z" (torch.Tensor): Quantized continuous representation, shape [B, D, T]. + - "codes" (torch.Tensor): Codebook indices, shape [B, N, T]. + - "latents" (torch.Tensor): Projected latents, shape [B, N*D, T]. + - "vq/commitment_loss" (torch.Tensor): Commitment loss. + - "vq/codebook_loss" (torch.Tensor): Codebook loss. + + **Continuous Mode (``continuous=True``):** + - "audio" (torch.Tensor): Decoded audio, shape [B, 1, length]. + - "z" (torch.Tensor): Latent representation, shape [B, D, T]. + - "kl_loss" (torch.Tensor): KL divergence loss (for VAE training). + """ + length = audio_data.shape[-1] + audio_data = self.preprocess(audio_data, sample_rate) + if not self.continuous: + z, codes, latents, commitment_loss, codebook_loss = self.encode( + audio_data, n_quantizers + ) + + x = self.decode(z) + return { + "audio": x[..., :length], + "z": z, + "codes": codes, + "latents": latents, + "vq/commitment_loss": commitment_loss, + "vq/codebook_loss": codebook_loss, + } + else: + posterior, _, _, _, _ = self.encode(audio_data, n_quantizers) + z = posterior.sample() + x = self.decode(z) + + kl_loss = posterior.kl(dims=(1, 2)) + kl_loss = kl_loss.mean() + + return { + "audio": x[..., :length], + "z": z, + "kl_loss": kl_loss, + } + + +EntryClass = DAC diff --git a/python/sglang/multimodal_gen/runtime/pipelines/mova_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/mova_pipeline.py new file mode 100644 index 000000000..2a8ed13b3 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/mova_pipeline.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +MOVA pipeline integration (native SGLang pipeline). +""" + +from __future__ import annotations + +from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig +from sglang.multimodal_gen.configs.sample.mova import MOVASamplingParams +from sglang.multimodal_gen.runtime.models.model_stages.mova import ( + MOVADecodingStage, + MOVADenoisingStage, + MOVALatentPreparationStage, + MOVATimestepPreparationStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages import ( + ConditioningStage, + ImageVAEEncodingStage, + InputValidationStage, + TextEncodingStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +class MOVAPipeline(ComposedPipelineBase): + """MOVA pipeline with SGLang stage orchestration.""" + + pipeline_name = "MOVA" + is_video_pipeline = True + _required_config_modules = [ + "video_vae", + "audio_vae", + "text_encoder", + "tokenizer", + "scheduler", + "video_dit", + "video_dit_2", + "audio_dit", + "dual_tower_bridge", + ] + pipeline_config_cls = MOVAPipelineConfig + sampling_params_cls = MOVASamplingParams + + def initialize_pipeline(self, server_args: ServerArgs) -> None: + """ + Initialize the pipeline. + + MOVA supports Context Parallel (sequence parallel) through USPAttention, + which uses Ulysses-style all-to-all communication for distributed attention. + """ + if server_args.sp_degree > 1: + logger.info( + "MOVA Context Parallel enabled with sp_degree=%d. " + "Using USPAttention for distributed self-attention.", + server_args.sp_degree, + ) + + def create_pipeline_stages(self, server_args: ServerArgs) -> None: + self.add_stage( + stage_name="input_validation_stage", stage=InputValidationStage() + ) + self.add_stage( + stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ), + ) + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + if getattr(self.get_module("video_dit"), "require_vae_embedding", True): + self.add_stage( + stage_name="image_latent_preparation_stage", + stage=ImageVAEEncodingStage(vae=self.get_module("video_vae")), + ) + self.add_stage( + stage_name="mova_latent_preparation_stage", + stage=MOVALatentPreparationStage( + audio_vae=self.get_module("audio_vae"), + require_vae_embedding=getattr( + self.get_module("video_dit"), "require_vae_embedding", True + ), + ), + ) + self.add_stage( + stage_name="mova_timestep_preparation_stage", + stage=MOVATimestepPreparationStage( + scheduler=self.get_module("scheduler"), + ), + ) + self.add_stage( + stage_name="mova_denoising_stage", + stage=MOVADenoisingStage( + video_dit=self.get_module("video_dit"), + video_dit_2=self.get_module("video_dit_2"), + audio_dit=self.get_module("audio_dit"), + dual_tower_bridge=self.get_module("dual_tower_bridge"), + scheduler=self.get_module("scheduler"), + ), + ) + self.add_stage( + stage_name="mova_decoding_stage", + stage=MOVADecodingStage( + video_vae=self.get_module("video_vae"), + audio_vae=self.get_module("audio_vae"), + ), + ) + + +class MOVAPipelineAlias(MOVAPipeline): + pipeline_name = "MOVAPipeline" + + +EntryClass = [MOVAPipeline, MOVAPipelineAlias] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index 5ca4e8092..ffd61b60e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -182,12 +182,25 @@ class ComposedPipelineBase(ABC): "boundary_ratio" in model_index and model_index["boundary_ratio"] is not None ): - logger.info( - "MoE pipeline detected. Adding transformer_2 to self.required_config_modules..." + has_transformer = ( + "transformer" in model_index + or "transformer_2" in model_index + or "transformer" in self.required_config_modules + or "transformer_2" in self.required_config_modules ) - self.required_config_modules.append("transformer_2") + if has_transformer: + logger.info( + "MoE pipeline detected. Adding transformer_2 to self.required_config_modules..." + ) + if "transformer_2" not in self.required_config_modules: + self.required_config_modules.append("transformer_2") + else: + logger.info( + "Boundary ratio found in model_index.json without transformers; " + "using it for pipeline config only." + ) logger.info( - "MoE pipeline detected. Setting boundary ratio to %s", + "Setting boundary ratio to %s", model_index["boundary_ratio"], ) server_args.pipeline_config.dit_config.boundary_ratio = model_index[ diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index 963959bc8..fe038960b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -90,11 +90,13 @@ class Req: # Latent tensors latents: torch.Tensor | None = None + y: torch.Tensor | None = None # Flux-2 latent_ids: torch.Tensor | None = None - # Audio Latents (LTX-2) + # Audio Latents audio_latents: torch.Tensor | None = None + audio_noise: torch.Tensor | None = None raw_audio_latent_shape: tuple[int, ...] | None = None # Audio Parameters @@ -114,6 +116,7 @@ class Req: # Timesteps timesteps: torch.Tensor | None = None + paired_timesteps: torch.Tensor | None = None timestep: torch.Tensor | float | int | None = None step_index: int | None = None @@ -154,6 +157,8 @@ class Req: # results output: torch.Tensor | None = None + audio: torch.Tensor | None = None + audio_sample_rate: int | None = None def __init__(self, **kwargs): # Initialize dataclass fields diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py index d4b1c71d6..f4b85c260 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py @@ -11,6 +11,7 @@ from PIL import Image from sglang.multimodal_gen.configs.pipeline_configs import WanI2V480PConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType +from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig from sglang.multimodal_gen.runtime.models.vision_utils import load_image, load_video from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage @@ -46,6 +47,25 @@ class InputValidationStage(PipelineStage): super().__init__() self.vae_image_processor = vae_image_processor + @staticmethod + def _calculate_dimensions_from_area( + max_area: float, aspect_ratio: float, mod_value: int + ) -> tuple[int, int]: + """ + Calculate output dimensions based on maximum area and aspect ratio. + + Args: + max_area: Maximum area constraint for the output + aspect_ratio: Target aspect ratio (height/width) + mod_value: Value to round dimensions to (typically vae_scale * patch_size) + + Returns: + Tuple of (width, height) rounded to multiples of mod_value + """ + height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value + width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value + return width, height + def _generate_seeds(self, batch: Req, server_args: ServerArgs): """Generate seeds for the inference""" seed = batch.seed @@ -168,13 +188,45 @@ class InputValidationStage(PipelineStage): server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial * server_args.pipeline_config.dit_config.arch_config.patch_size[1] ) - height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value - width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value + width, height = self._calculate_dimensions_from_area( + max_area, aspect_ratio, mod_value + ) batch.condition_image = batch.condition_image.resize((width, height)) batch.height = height batch.width = width + elif issubclass(type(server_args.pipeline_config), MOVAPipelineConfig): + # resize image only, MOVA + image = batch.condition_image + if isinstance(image, list): + image = image[0] # not support multi image input yet. + + max_area = server_args.pipeline_config.max_area + if hasattr(batch, "height") and hasattr(batch, "width"): + aspect_ratio = batch.height / batch.width + else: + aspect_ratio = ( + batch.sampling_params.height / batch.sampling_params.width + ) + mod_value = ( + server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial + * server_args.pipeline_config.dit_config.arch_config.patch_size[1] + ) + width, height = self._calculate_dimensions_from_area( + max_area, aspect_ratio, mod_value + ) + + config = server_args.pipeline_config + image, (final_w, final_h) = ( + server_args.pipeline_config.preprocess_condition_image( + image, width, height, self.vae_image_processor + ) + ) + batch.condition_image = image + batch.width = final_w + batch.height = final_h + def forward( self, batch: Req, diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 6a6be2064..f53ee67f5 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -342,6 +342,11 @@ class ServerArgs: default_factory=lambda: { "transformer": True, "vae": True, + "video_vae": True, + "audio_vae": True, + "video_dit": True, + "audio_dit": True, + "dual_tower_bridge": True, } ) diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py index 0bde44405..55ca4f349 100644 --- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py @@ -412,20 +412,6 @@ def verify_model_config_and_directory(model_path: str) -> dict[str, Any]: "Only HuggingFace diffusers format is supported." ) - # Check for transformer and vae directories - transformer_dir = os.path.join(model_path, "transformer") - vae_dir = os.path.join(model_path, "vae") - - if not os.path.exists(transformer_dir): - raise ValueError( - f"Model directory {model_path} does not contain a transformer/ directory." - ) - - if not os.path.exists(vae_dir): - raise ValueError( - f"Model directory {model_path} does not contain a vae/ directory." - ) - # Load the config with open(config_path) as f: config = json.load(f) @@ -435,6 +421,37 @@ def verify_model_config_and_directory(model_path: str) -> dict[str, Any]: raise ValueError("model_index.json does not contain _diffusers_version") logger.info("Diffusers version: %s", config["_diffusers_version"]) + + component_keys = [ + key + for key, value in config.items() + if isinstance(value, (list, tuple)) + and len(value) == 2 + and all(isinstance(item, str) for item in value) + ] + if component_keys: + missing_components = [ + component_key + for component_key in component_keys + if not os.path.exists(os.path.join(model_path, component_key)) + ] + if missing_components: + missing_str = ", ".join(missing_components) + raise ValueError( + f"Model directory {model_path} is missing required component " + f"directories: {missing_str}." + ) + else: + transformer_dir = os.path.join(model_path, "transformer") + vae_dir = os.path.join(model_path, "vae") + if not os.path.exists(transformer_dir): + raise ValueError( + f"Model directory {model_path} does not contain a transformer/ directory." + ) + if not os.path.exists(vae_dir): + raise ValueError( + f"Model directory {model_path} does not contain a vae/ directory." + ) return cast(dict[str, Any], config) @@ -543,14 +560,35 @@ def maybe_download_model( def _verify_model_complete(path: str) -> bool: """Check if model directory has required subdirectories.""" + config_path = os.path.join(path, "model_index.json") + if not os.path.exists(config_path): + return False + + try: + with open(config_path) as config_file: + model_index = json.load(config_file) + except Exception as exc: + logger.warning( + "Failed to read model_index.json at %s: %s", config_path, exc + ) + return False + + component_keys = [ + key + for key, value in model_index.items() + if isinstance(value, (list, tuple)) + and len(value) == 2 + and all(isinstance(item, str) for item in value) + ] + if component_keys: + return all( + os.path.exists(os.path.join(path, component_key)) + for component_key in component_keys + ) + transformer_dir = os.path.join(path, "transformer") vae_dir = os.path.join(path, "vae") - config_path = os.path.join(path, "model_index.json") - return ( - os.path.exists(config_path) - and os.path.exists(transformer_dir) - and os.path.exists(vae_dir) - ) + return os.path.exists(transformer_dir) and os.path.exists(vae_dir) # 1. Local path check: if path exists locally, verify it's complete (skip for LoRA) if os.path.exists(model_name_or_path): @@ -584,7 +622,7 @@ def maybe_download_model( return model_name_or_path else: logger.warning( - "Local model at %s appears incomplete (missing transformer/ or vae/), " + "Local model at %s appears incomplete (missing required components), " "will attempt re-download", model_name_or_path, )