From 29ce7b3612da63ab8ed6ca2860a1ee351fbed2e1 Mon Sep 17 00:00:00 2001 From: Yuhao Yang <47235274+yhyang201@users.noreply.github.com> Date: Sat, 27 Dec 2025 13:25:05 +0800 Subject: [PATCH] [diffusion] chore: remove stepvideo code (#15918) --- .../configs/models/dits/__init__.py | 3 +- .../configs/models/dits/stepvideo.py | 64 - .../configs/models/vaes/__init__.py | 2 - .../configs/models/vaes/stepvideovae.py | 31 - .../configs/pipeline_configs/__init__.py | 2 - .../configs/pipeline_configs/base.py | 26 - .../configs/pipeline_configs/stepvideo.py | 36 - .../configs/sample/sampling_params.py | 11 - .../configs/sample/stepvideo.py | 22 - python/sglang/multimodal_gen/registry.py | 12 - .../runtime/loader/component_loader.py | 3 +- .../runtime/models/dits/stepvideo.py | 731 ---------- .../runtime/models/encoders/stepllm.py | 617 --------- .../runtime/models/vaes/stepvideovae.py | 1184 ----------------- .../runtime/pipelines/stepvideo_pipeline.py | 182 --- .../runtime/pipelines_core/stages/__init__.py | 4 - .../stages/latent_preparation.py | 2 - .../stages/stepvideo_encoding.py | 99 -- .../multimodal_gen/runtime/server_args.py | 5 - 19 files changed, 2 insertions(+), 3034 deletions(-) delete mode 100644 python/sglang/multimodal_gen/configs/models/dits/stepvideo.py delete mode 100644 python/sglang/multimodal_gen/configs/models/vaes/stepvideovae.py delete mode 100644 python/sglang/multimodal_gen/configs/pipeline_configs/stepvideo.py delete mode 100644 python/sglang/multimodal_gen/configs/sample/stepvideo.py delete mode 100644 python/sglang/multimodal_gen/runtime/models/dits/stepvideo.py delete mode 100644 python/sglang/multimodal_gen/runtime/models/encoders/stepllm.py delete mode 100644 python/sglang/multimodal_gen/runtime/models/vaes/stepvideovae.py delete mode 100644 python/sglang/multimodal_gen/runtime/pipelines/stepvideo_pipeline.py delete mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/stepvideo_encoding.py diff --git a/python/sglang/multimodal_gen/configs/models/dits/__init__.py b/python/sglang/multimodal_gen/configs/models/dits/__init__.py index 67e6d97b4..30b205451 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/dits/__init__.py @@ -1,7 +1,6 @@ # 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.stepvideo import StepVideoConfig from sglang.multimodal_gen.configs.models.dits.wanvideo import WanVideoConfig -__all__ = ["HunyuanVideoConfig", "WanVideoConfig", "StepVideoConfig"] +__all__ = ["HunyuanVideoConfig", "WanVideoConfig"] diff --git a/python/sglang/multimodal_gen/configs/models/dits/stepvideo.py b/python/sglang/multimodal_gen/configs/models/dits/stepvideo.py deleted file mode 100644 index 1d7fe21a6..000000000 --- a/python/sglang/multimodal_gen/configs/models/dits/stepvideo.py +++ /dev/null @@ -1,64 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -# SPDX-License-Identifier: Apache-2.0 -from dataclasses import dataclass, field - -from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig - - -def is_transformer_blocks(n, m): - return "transformer_blocks" in n and n.split(".")[-1].isdigit() - - -@dataclass -class StepVideoArchConfig(DiTArchConfig): - _fsdp_shard_conditions: list = field( - default_factory=lambda: [is_transformer_blocks] - ) - - param_names_mapping: dict = field( - default_factory=lambda: { - # transformer block - r"^transformer_blocks\.(\d+)\.norm1\.(weight|bias)$": r"transformer_blocks.\1.norm1.norm.\2", - r"^transformer_blocks\.(\d+)\.norm2\.(weight|bias)$": r"transformer_blocks.\1.norm2.norm.\2", - r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.weight$": r"transformer_blocks.\1.ff.fc_in.weight", - r"^transformer_blocks\.(\d+)\.ff\.net\.2\.weight$": r"transformer_blocks.\1.ff.fc_out.weight", - # adanorm block - r"^adaln_single\.emb\.timestep_embedder\.linear_1\.(weight|bias)$": r"adaln_single.emb.mlp.fc_in.\1", - r"^adaln_single\.emb\.timestep_embedder\.linear_2\.(weight|bias)$": r"adaln_single.emb.mlp.fc_out.\1", - # caption projection - r"^caption_projection\.linear_1\.(weight|bias)$": r"caption_projection.fc_in.\1", - r"^caption_projection\.linear_2\.(weight|bias)$": r"caption_projection.fc_out.\1", - } - ) - - num_attention_heads: int = 48 - attention_head_dim: int = 128 - in_channels: int = 64 - out_channels: int | None = 64 - num_layers: int = 48 - dropout: float = 0.0 - patch_size: int = 1 - norm_type: str = "ada_norm_single" - norm_elementwise_affine: bool = False - norm_eps: float = 1e-6 - caption_channels: int | list[int] | tuple[int, ...] | None = field( - default_factory=lambda: [6144, 1024] - ) - attention_type: str | None = "torch" - use_additional_conditions: bool | None = False - exclude_lora_layers: list[str] = field(default_factory=lambda: []) - - def __post_init__(self): - self.hidden_size = self.num_attention_heads * self.attention_head_dim - self.out_channels = ( - self.in_channels if self.out_channels is None else self.out_channels - ) - self.num_channels_latents = self.out_channels - - -@dataclass -class StepVideoConfig(DiTConfig): - arch_config: DiTArchConfig = field(default_factory=StepVideoArchConfig) - - prefix: str = "StepVideo" diff --git a/python/sglang/multimodal_gen/configs/models/vaes/__init__.py b/python/sglang/multimodal_gen/configs/models/vaes/__init__.py index e9b478618..1d6e7461e 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/__init__.py @@ -1,11 +1,9 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo from sglang.multimodal_gen.configs.models.vaes.hunyuanvae import HunyuanVAEConfig -from sglang.multimodal_gen.configs.models.vaes.stepvideovae import StepVideoVAEConfig from sglang.multimodal_gen.configs.models.vaes.wanvae import WanVAEConfig __all__ = [ "HunyuanVAEConfig", "WanVAEConfig", - "StepVideoVAEConfig", ] diff --git a/python/sglang/multimodal_gen/configs/models/vaes/stepvideovae.py b/python/sglang/multimodal_gen/configs/models/vaes/stepvideovae.py deleted file mode 100644 index 6794e9792..000000000 --- a/python/sglang/multimodal_gen/configs/models/vaes/stepvideovae.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -# SPDX-License-Identifier: Apache-2.0 -from dataclasses import dataclass, field - -from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig - - -@dataclass -class StepVideoVAEArchConfig(VAEArchConfig): - in_channels: int = 3 - out_channels: int = 3 - z_channels: int = 64 - num_res_blocks: int = 2 - version: int = 2 - frame_len: int = 17 - world_size: int = 1 - - spatial_compression_ratio: int = 16 - temporal_compression_ratio: int = 8 - - scaling_factor: float = 1.0 - - -@dataclass -class StepVideoVAEConfig(VAEConfig): - arch_config: VAEArchConfig = field(default_factory=StepVideoVAEArchConfig) - use_tiling: bool = False - use_temporal_tiling: bool = False - use_parallel_tiling: bool = False - use_temporal_scaling_frames: bool = False diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py index c482cc4a7..dbb394ba8 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py @@ -12,7 +12,6 @@ from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import ( FastHunyuanConfig, HunyuanConfig, ) -from sglang.multimodal_gen.configs.pipeline_configs.stepvideo import StepVideoT2VConfig from sglang.multimodal_gen.configs.pipeline_configs.wan import ( SelfForcingWanT2V480PConfig, WanI2V480PConfig, @@ -33,7 +32,6 @@ __all__ = [ "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig", - "StepVideoT2VConfig", "SelfForcingWanT2V480PConfig", "ZImagePipelineConfig", ] diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 570da92c8..ab09b8112 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -195,11 +195,6 @@ class PipelineConfig: field(default_factory=lambda: (postprocess_text,)) ) - # StepVideo specific parameters - pos_magic: str | None = None - neg_magic: str | None = None - timesteps_scale: bool | None = None - # STA (Sliding Tile Attention) parameters mask_strategy_file_path: str | None = None STA_mode: STA_Mode = STA_Mode.STA_INFERENCE @@ -468,27 +463,6 @@ class PipelineConfig: choices=["fp32", "fp16", "bf16"], help="Precision for image encoder", ) - parser.add_argument( - f"--{prefix_with_dot}pos_magic", - type=str, - dest=f"{prefix_with_dot.replace('-', '_')}pos_magic", - default=PipelineConfig.pos_magic, - help="Positive magic prompt for sampling, used in stepvideo", - ) - parser.add_argument( - f"--{prefix_with_dot}neg_magic", - type=str, - dest=f"{prefix_with_dot.replace('-', '_')}neg_magic", - default=PipelineConfig.neg_magic, - help="Negative magic prompt for sampling, used in stepvideo", - ) - parser.add_argument( - f"--{prefix_with_dot}timesteps_scale", - type=bool, - dest=f"{prefix_with_dot.replace('-', '_')}timesteps_scale", - default=PipelineConfig.timesteps_scale, - help="Bool for applying scheduler scale in set_timesteps, used in stepvideo", - ) # DMD parameters parser.add_argument( diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/stepvideo.py b/python/sglang/multimodal_gen/configs/pipeline_configs/stepvideo.py deleted file mode 100644 index aff18e5cf..000000000 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/stepvideo.py +++ /dev/null @@ -1,36 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -# SPDX-License-Identifier: Apache-2.0 -from dataclasses import dataclass, field - -from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig -from sglang.multimodal_gen.configs.models.dits import StepVideoConfig -from sglang.multimodal_gen.configs.models.vaes import StepVideoVAEConfig -from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig - - -@dataclass -class StepVideoT2VConfig(PipelineConfig): - """Base configuration for StepVideo pipeline architecture.""" - - # WanConfig-specific parameters with defaults - # DiT - dit_config: DiTConfig = field(default_factory=StepVideoConfig) - # VAE - vae_config: VAEConfig = field(default_factory=StepVideoVAEConfig) - vae_tiling: bool = False - vae_sp: bool = False - - # Denoising stage - flow_shift: int = 13 - timesteps_scale: bool = False - pos_magic: str = ( - "超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。" - ) - neg_magic: str = ( - "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。" - ) - - # Precision for each component - precision: str = "bf16" - vae_precision: str = "bf16" diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 4d538a93f..7afa6fc77 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -308,8 +308,6 @@ class SamplingParams: if use_temporal_scaling_frames: orig_latent_num_frames = (num_frames - 1) // temporal_scale_factor + 1 - else: # stepvideo only - orig_latent_num_frames = self.num_frames // 17 * 3 if orig_latent_num_frames % server_args.num_gpus != 0: # Adjust latent frames to be divisible by number of GPUs @@ -328,15 +326,6 @@ class SamplingParams: new_num_frames = ( new_latent_num_frames - 1 ) * temporal_scale_factor + 1 - else: # stepvideo only - # Find the least common multiple of 3 and num_gpus - divisor = math.lcm(3, num_gpus) - # Round up to the nearest multiple of this LCM - new_latent_num_frames = ( - (new_latent_num_frames + divisor - 1) // divisor - ) * divisor - # Convert back to actual frames using the StepVideo formula - new_num_frames = new_latent_num_frames // 3 * 17 logger.info( "Adjusting number of frames from %s to %s based on number of GPUs (%s)", diff --git a/python/sglang/multimodal_gen/configs/sample/stepvideo.py b/python/sglang/multimodal_gen/configs/sample/stepvideo.py deleted file mode 100644 index 4fff150b0..000000000 --- a/python/sglang/multimodal_gen/configs/sample/stepvideo.py +++ /dev/null @@ -1,22 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -# SPDX-License-Identifier: Apache-2.0 -from dataclasses import dataclass - -from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams - - -@dataclass -class StepVideoT2VSamplingParams(SamplingParams): - # Video parameters - height: int = 720 - width: int = 1280 - num_frames: int = 81 - - # Denoising stage - guidance_scale: float = 9.0 - num_inference_steps: int = 50 - - # neg magic and pos magic - # pos_magic: str = "超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。" - # neg_magic: str = "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。" diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 7d256c0b9..42cb20259 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -18,7 +18,6 @@ from sglang.multimodal_gen.configs.pipeline_configs import ( FastHunyuanConfig, FluxPipelineConfig, HunyuanConfig, - StepVideoT2VConfig, WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, @@ -50,7 +49,6 @@ from sglang.multimodal_gen.configs.sample.qwenimage import ( QwenImageLayeredSamplingParams, QwenImageSamplingParams, ) -from sglang.multimodal_gen.configs.sample.stepvideo import StepVideoT2VSamplingParams from sglang.multimodal_gen.configs.sample.wan import ( FastWanT2V480PConfig, Wan2_1_Fun_1_3B_InP_SamplingParams, @@ -312,16 +310,6 @@ def _register_configs(): ], ) - # StepVideo - register_configs( - sampling_param_cls=StepVideoT2VSamplingParams, - pipeline_config_cls=StepVideoT2VConfig, - hf_model_paths=[ - "FastVideo/stepvideo-t2v-diffusers", - ], - model_detectors=[lambda hf_id: "stepvideo" in hf_id.lower()], - ) - # Wan register_configs( sampling_param_cls=WanT2V_1_3B_SamplingParams, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index 29ba60bc9..324ee2b9a 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loader.py @@ -752,8 +752,7 @@ class SchedulerLoader(ComponentLoader): scheduler = scheduler_cls(**config) if server_args.pipeline_config.flow_shift is not None: scheduler.set_shift(server_args.pipeline_config.flow_shift) - if server_args.pipeline_config.timesteps_scale is not None: - scheduler.set_timesteps_scale(server_args.pipeline_config.timesteps_scale) + return scheduler diff --git a/python/sglang/multimodal_gen/runtime/models/dits/stepvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/stepvideo.py deleted file mode 100644 index 6dca091ac..000000000 --- a/python/sglang/multimodal_gen/runtime/models/dits/stepvideo.py +++ /dev/null @@ -1,731 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -# Copyright 2025 StepFun Inc. All Rights Reserved. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# ============================================================================== -from typing import Any - -import torch -from einops import rearrange, repeat -from torch import nn - -from sglang.multimodal_gen.configs.models.dits import StepVideoConfig -from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size -from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAttention -from sglang.multimodal_gen.runtime.layers.layernorm import LayerNormScaleShift -from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear -from sglang.multimodal_gen.runtime.layers.mlp import MLP -from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( - _apply_rotary_emb, - get_rotary_pos_embed, -) -from sglang.multimodal_gen.runtime.layers.visual_embedding import TimestepEmbedder -from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT -from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum - - -class PatchEmbed2D(nn.Module): - """2D Image to Patch Embedding - - Image to Patch Embedding using Conv2d - - A convolution based approach to patchifying a 2D image w/ embedding projection. - - Based on the impl in https://github.com/google-research/vision_transformer - - Hacked together by / Copyright 2020 Ross Wightman - - Remove the _assert function in forward function to be compatible with multi-resolution images. - """ - - def __init__( - self, - patch_size=16, - in_chans=3, - embed_dim=768, - norm_layer=None, - flatten=True, - bias=True, - dtype=None, - prefix: str = "", - ): - super().__init__() - # Convert patch_size to 2-tuple - if isinstance(patch_size, list | tuple): - if len(patch_size) == 1: - patch_size = (patch_size[0], patch_size[0]) - else: - patch_size = (patch_size, patch_size) - - self.patch_size = patch_size - self.flatten = flatten - - self.proj = nn.Conv2d( - in_chans, - embed_dim, - kernel_size=patch_size, - stride=patch_size, - bias=bias, - dtype=dtype, - ) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x): - x = self.proj(x) - if self.flatten: - x = x.flatten(2).transpose(1, 2) # BCHW -> BNC - x = self.norm(x) - return x - - -class StepVideoRMSNorm(nn.Module): - - def __init__( - self, - dim: int, - elementwise_affine=True, - eps: float = 1e-6, - device=None, - dtype=None, - ): - """ - Initialize the RMSNorm normalization layer. - - Args: - dim (int): The dimension of the input tensor. - eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. - - Attributes: - eps (float): A small value added to the denominator for numerical stability. - weight (nn.Parameter): Learnable scaling parameter. - - """ - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - if elementwise_affine: - self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) - - def _norm(self, x) -> torch.Tensor: - """ - Apply the RMSNorm normalization to the input tensor. - - Args: - x (torch.Tensor): The input tensor. - - Returns: - torch.Tensor: The normalized tensor. - - """ - return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) - - def forward(self, x): - """ - Forward pass through the RMSNorm layer. - - Args: - x (torch.Tensor): The input tensor. - - Returns: - torch.Tensor: The output tensor after applying RMSNorm. - - """ - output = self._norm(x.float()).type_as(x) - if hasattr(self, "weight"): - output = output * self.weight - return output - - -class SelfAttention(nn.Module): - - def __init__( - self, - hidden_dim, - head_dim, - rope_split: tuple[int, int, int] = (64, 32, 32), - bias: bool = False, - with_rope: bool = True, - with_qk_norm: bool = True, - attn_type: str = "torch", - supported_attention_backends=( - AttentionBackendEnum.FA, - AttentionBackendEnum.AITER, - AttentionBackendEnum.TORCH_SDPA, - ), - ): - super().__init__() - self.head_dim = head_dim - self.hidden_dim = hidden_dim - self.rope_split = list(rope_split) - self.n_heads = hidden_dim // head_dim - - self.wqkv = ReplicatedLinear(hidden_dim, hidden_dim * 3, bias=bias) - self.wo = ReplicatedLinear(hidden_dim, hidden_dim, bias=bias) - - self.with_rope = with_rope - self.with_qk_norm = with_qk_norm - if self.with_qk_norm: - self.q_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True) - self.k_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True) - - # self.core_attention = self.attn_processor(attn_type=attn_type) - self.parallel = attn_type == "parallel" - self.attn = USPAttention( - num_heads=self.n_heads, - head_size=head_dim, - causal=False, - supported_attention_backends=supported_attention_backends, - ) - - def _apply_rope(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): - """ - x: [B, S, H, D] - cos: [S, D/2] where D = head_dim = sum(self.rope_split) - sin: [S, D/2] - returns x with rotary applied exactly as v0 did - """ - B, S, H, D = x.shape - # 1) split cos/sin per chunk - half_splits = [c // 2 for c in self.rope_split] # [32,16,16] for [64,32,32] - cos_splits = cos.split(half_splits, dim=1) - sin_splits = sin.split(half_splits, dim=1) - - outs = [] - idx = 0 - for chunk_size, cos_i, sin_i in zip( - self.rope_split, cos_splits, sin_splits, strict=True - ): - # slice the corresponding channels - x_chunk = x[..., idx : idx + chunk_size] # [B,S,H,chunk_size] - idx += chunk_size - - # flatten to [S, B*H, chunk_size] - x_flat = rearrange(x_chunk, "b s h d -> s (b h) d") - - # apply rotary on *that* chunk - out_flat = _apply_rotary_emb(x_flat, cos_i, sin_i, is_neox_style=True) - - # restore [B,S,H,chunk_size] - out = rearrange(out_flat, "s (b h) d -> b s h d", b=B, h=H) - outs.append(out) - - # concatenate back to [B,S,H,D] - return torch.cat(outs, dim=-1) - - def forward( - self, - x, - cu_seqlens=None, - max_seqlen=None, - rope_positions=None, - cos_sin=None, - attn_mask=None, - mask_strategy=None, - ): - - B, S, _ = x.shape - xqkv, _ = self.wqkv(x) - xqkv = xqkv.view(*x.shape[:-1], self.n_heads, 3 * self.head_dim) - q, k, v = torch.split(xqkv, [self.head_dim] * 3, dim=-1) # [B,S,H,D] - - if self.with_qk_norm: - q = self.q_norm(q) - k = self.k_norm(k) - - if self.with_rope: - if rope_positions is not None: - F, Ht, W = rope_positions - assert F * Ht * W == S, "rope_positions mismatches sequence length" - - cos, sin = cos_sin - cos = cos.to(x.device, dtype=x.dtype) - sin = sin.to(x.device, dtype=x.dtype) - - q = self._apply_rope(q, cos, sin) - k = self._apply_rope(k, cos, sin) - - output = self.attn(q, k, v) # [B,heads,S,D] - - output = rearrange(output, "b s h d -> b s (h d)") - output, _ = self.wo(output) - - return output - - -class CrossAttention(nn.Module): - - def __init__( - self, - hidden_dim, - head_dim, - bias=False, - with_qk_norm=True, - supported_attention_backends=( - AttentionBackendEnum.FA, - AttentionBackendEnum.AITER, - AttentionBackendEnum.TORCH_SDPA, - ), - ) -> None: - super().__init__() - self.head_dim = head_dim - self.n_heads = hidden_dim // head_dim - - self.wq = ReplicatedLinear(hidden_dim, hidden_dim, bias=bias) - self.wkv = ReplicatedLinear(hidden_dim, hidden_dim * 2, bias=bias) - self.wo = ReplicatedLinear(hidden_dim, hidden_dim, bias=bias) - - self.with_qk_norm = with_qk_norm - if self.with_qk_norm: - self.q_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True) - self.k_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True) - - self.attn = LocalAttention( - num_heads=self.n_heads, - head_size=head_dim, - causal=False, - supported_attention_backends=supported_attention_backends, - ) - - def forward( - self, x: torch.Tensor, encoder_hidden_states: torch.Tensor, attn_mask=None - ) -> torch.Tensor: - - xq, _ = self.wq(x) - xq = xq.view(*xq.shape[:-1], self.n_heads, self.head_dim) - - xkv, _ = self.wkv(encoder_hidden_states) - xkv = xkv.view(*xkv.shape[:-1], self.n_heads, 2 * self.head_dim) - - xk, xv = torch.split(xkv, [self.head_dim] * 2, dim=-1) ## seq_len, n, dim - - if self.with_qk_norm: - xq = self.q_norm(xq) - xk = self.k_norm(xk) - - output = self.attn(xq, xk, xv) - - output = rearrange(output, "b s h d -> b s (h d)") - output, _ = self.wo(output) - - return output - - -class AdaLayerNormSingle(nn.Module): - r""" - Norm layer adaptive layer norm single (adaLN-single). - - As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3). - - Parameters: - embedding_dim (`int`): The size of each embedding vector. - use_additional_conditions (`bool`): To use additional conditions for normalization or not. - """ - - def __init__(self, embedding_dim: int, time_step_rescale=1000): - super().__init__() - - self.emb = TimestepEmbedder(embedding_dim) - - self.silu = nn.SiLU() - self.linear = ReplicatedLinear(embedding_dim, 6 * embedding_dim, bias=True) - - self.time_step_rescale = time_step_rescale ## timestep usually in [0, 1], we rescale it to [0,1000] for stability - - def forward( - self, - timestep: torch.Tensor, - added_cond_kwargs: dict[str, torch.Tensor] | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - embedded_timestep = self.emb(timestep * self.time_step_rescale) - - out, _ = self.linear(self.silu(embedded_timestep)) - - return out, embedded_timestep - - -class StepVideoTransformerBlock(nn.Module): - r""" - A basic Transformer block. - - Parameters: - dim (`int`): The number of channels in the input and output. - num_attention_heads (`int`): The number of heads to use for multi-head attention. - attention_head_dim (`int`): The number of channels in each head. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. - activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. - num_embeds_ada_norm (: - obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. - attention_bias (: - obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. - only_cross_attention (`bool`, *optional*): - Whether to use only cross-attention layers. In this case two cross attention layers are used. - double_self_attention (`bool`, *optional*): - Whether to use two self-attention layers. In this case no cross attention layers are used. - upcast_attention (`bool`, *optional*): - Whether to upcast the attention computation to float32. This is useful for mixed precision training. - norm_elementwise_affine (`bool`, *optional*, defaults to `True`): - Whether to use learnable elementwise affine parameters for normalization. - norm_type (`str`, *optional*, defaults to `"layer_norm"`): - The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`. - final_dropout (`bool` *optional*, defaults to False): - Whether to apply a final dropout after the last feed-forward layer. - positional_embeddings (`str`, *optional*, defaults to `None`): - The type of positional embeddings to apply to. - num_positional_embeddings (`int`, *optional*, defaults to `None`): - The maximum number of positional embeddings to apply. - """ - - def __init__( - self, - dim: int, - attention_head_dim: int, - norm_eps: float = 1e-5, - ff_inner_dim: int | None = None, - ff_bias: bool = False, - attention_type: str = "torch", - ): - super().__init__() - self.dim = dim - self.norm1 = LayerNormScaleShift( - dim, norm_type="layer", elementwise_affine=True, eps=norm_eps - ) - self.attn1 = SelfAttention( - dim, - attention_head_dim, - bias=False, - with_rope=True, - with_qk_norm=True, - ) - - self.norm2 = LayerNormScaleShift( - dim, norm_type="layer", elementwise_affine=True, eps=norm_eps - ) - self.attn2 = CrossAttention( - dim, attention_head_dim, bias=False, with_qk_norm=True - ) - - self.ff = MLP( - input_dim=dim, - mlp_hidden_dim=dim * 4 if ff_inner_dim is None else ff_inner_dim, - act_type="gelu_pytorch_tanh", - bias=ff_bias, - ) - - self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) - - @torch.no_grad() - def forward( - self, - q: torch.Tensor, - kv: torch.Tensor, - t_expand: torch.LongTensor, - attn_mask=None, - rope_positions: list | None = None, - cos_sin=None, - mask_strategy=None, - ) -> torch.Tensor: - - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( - torch.clone(chunk) - for chunk in ( - self.scale_shift_table[None] + t_expand.reshape(-1, 6, self.dim) - ).chunk(6, dim=1) - ) - - scale_shift_q = self.norm1( - q, scale=scale_msa.squeeze(1), shift=shift_msa.squeeze(1) - ) - - attn_q = self.attn1( - scale_shift_q, - rope_positions=rope_positions, - cos_sin=cos_sin, - mask_strategy=mask_strategy, - ) - - q = attn_q * gate_msa + q - - attn_q = self.attn2(q, kv, attn_mask) - - q = attn_q + q - - scale_shift_q = self.norm2( - q, scale=scale_mlp.squeeze(1), shift=shift_mlp.squeeze(1) - ) - - ff_output = self.ff(scale_shift_q) - - q = ff_output * gate_mlp + q - - return q - - -class StepVideoModel(BaseDiT): - # (Optional) Keep the same attribute for compatibility with splitting, etc. - _fsdp_shard_conditions = [ - lambda n, m: "transformer_blocks" in n and n.split(".")[-1].isdigit(), - # lambda n, m: "pos_embed" in n # If needed for the patch embedding. - ] - param_names_mapping = StepVideoConfig().param_names_mapping - reverse_param_names_mapping = StepVideoConfig().reverse_param_names_mapping - lora_param_names_mapping = StepVideoConfig().lora_param_names_mapping - _supported_attention_backends = StepVideoConfig()._supported_attention_backends - - def __init__(self, config: StepVideoConfig, hf_config: dict[str, Any]) -> None: - super().__init__(config=config, hf_config=hf_config) - self.num_attention_heads = config.num_attention_heads - self.attention_head_dim = config.attention_head_dim - self.in_channels = config.in_channels - self.out_channels = config.out_channels - self.num_layers = config.num_layers - self.dropout = config.dropout - self.patch_size = config.patch_size - self.norm_type = config.norm_type - self.norm_elementwise_affine = config.norm_elementwise_affine - self.norm_eps = config.norm_eps - self.use_additional_conditions = config.use_additional_conditions - self.caption_channels = config.caption_channels - self.attention_type = config.attention_type - self.num_channels_latents = config.num_channels_latents - # Compute inner dimension. - self.hidden_size = config.hidden_size - - # Image/video patch embedding. - self.pos_embed = PatchEmbed2D( - patch_size=self.patch_size, - in_chans=self.in_channels, - embed_dim=self.hidden_size, - ) - - self._rope_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {} - # Transformer blocks. - self.transformer_blocks = nn.ModuleList( - [ - StepVideoTransformerBlock( - dim=self.hidden_size, - attention_head_dim=self.attention_head_dim, - attention_type=self.attention_type, - ) - for _ in range(self.num_layers) - ] - ) - - # Output blocks. - self.norm_out = LayerNormScaleShift( - self.hidden_size, - norm_type="layer", - eps=self.norm_eps, - elementwise_affine=self.norm_elementwise_affine, - ) - self.scale_shift_table = nn.Parameter( - torch.randn(2, self.hidden_size) / (self.hidden_size**0.5) - ) - self.proj_out = ReplicatedLinear( - self.hidden_size, self.patch_size * self.patch_size * self.out_channels - ) - # Time modulation via adaptive layer norm. - self.adaln_single = AdaLayerNormSingle(self.hidden_size) - - # Set up caption conditioning. - if isinstance(self.caption_channels, int): - caption_channel = self.caption_channels - else: - caption_channel, clip_channel = self.caption_channels - self.clip_projection = ReplicatedLinear(clip_channel, self.hidden_size) - self.caption_norm = nn.LayerNorm( - caption_channel, - eps=self.norm_eps, - elementwise_affine=self.norm_elementwise_affine, - ) - self.caption_projection = MLP( - input_dim=caption_channel, - mlp_hidden_dim=self.hidden_size, - act_type="gelu_pytorch_tanh", - ) - - # Flag to indicate if using parallel attention. - self.parallel = self.attention_type == "parallel" - - self.__post_init__() - - def patchfy(self, hidden_states) -> torch.Tensor: - hidden_states = rearrange(hidden_states, "b f c h w -> (b f) c h w") - hidden_states = self.pos_embed(hidden_states) - return hidden_states - - def prepare_attn_mask( - self, encoder_attention_mask, encoder_hidden_states, q_seqlen - ) -> tuple[torch.Tensor, torch.Tensor]: - kv_seqlens = encoder_attention_mask.sum(dim=1).int() - mask = torch.zeros( - [len(kv_seqlens), q_seqlen, max(kv_seqlens)], - dtype=torch.bool, - device=encoder_attention_mask.device, - ) - encoder_hidden_states = encoder_hidden_states[:, : max(kv_seqlens)] - for i, kv_len in enumerate(kv_seqlens): - mask[i, :, :kv_len] = 1 - return encoder_hidden_states, mask - - def block_forward( - self, - hidden_states, - encoder_hidden_states=None, - t_expand=None, - rope_positions=None, - cos_sin=None, - attn_mask=None, - parallel=True, - mask_strategy=None, - ) -> torch.Tensor: - - for i, block in enumerate(self.transformer_blocks): - hidden_states = block( - hidden_states, - encoder_hidden_states, - t_expand=t_expand, - attn_mask=attn_mask, - rope_positions=rope_positions, - cos_sin=cos_sin, - mask_strategy=mask_strategy[i], - ) - - return hidden_states - - def _get_rope( - self, - rope_positions: tuple[int, int, int], - dtype: torch.dtype, - device: torch.device, - ): - F, Ht, W = rope_positions - key = (F, Ht, W, dtype) - if key not in self._rope_cache: - cos, sin = get_rotary_pos_embed( - rope_sizes=(F * get_sp_world_size(), Ht, W), - hidden_size=self.hidden_size, - heads_num=self.hidden_size // self.attention_head_dim, - rope_dim_list=(64, 32, 32), # same split you used - rope_theta=1.0e4, - dtype=torch.float32, # build once in fp32 - ) - # move & cast once - self._rope_cache[key] = ( - cos.to(device, dtype=dtype), - sin.to(device, dtype=dtype), - ) - return self._rope_cache[key] - - @torch.inference_mode() - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor | None = None, - t_expand: torch.LongTensor | None = None, - encoder_hidden_states_2: torch.Tensor | None = None, - added_cond_kwargs: dict[str, torch.Tensor] | None = None, - encoder_attention_mask: torch.Tensor | None = None, - fps: torch.Tensor | None = None, - return_dict: bool = True, - mask_strategy=None, - guidance=None, - ): - assert hidden_states.ndim == 5 - "hidden_states's shape should be (bsz, f, ch, h ,w)" - frame = hidden_states.shape[2] - hidden_states = rearrange(hidden_states, "b c f h w -> b f c h w", f=frame) - if mask_strategy is None: - mask_strategy = [None, None] - bsz, frame, _, height, width = hidden_states.shape - height, width = height // self.patch_size, width // self.patch_size - - hidden_states = self.patchfy(hidden_states) - len_frame = hidden_states.shape[1] - - t_expand, embedded_timestep = self.adaln_single(t_expand) - encoder_hidden_states = self.caption_projection( - self.caption_norm(encoder_hidden_states) - ) - - if encoder_hidden_states_2 is not None and hasattr(self, "clip_projection"): - clip_embedding, _ = self.clip_projection(encoder_hidden_states_2) - encoder_hidden_states = torch.cat( - [clip_embedding, encoder_hidden_states], dim=1 - ) - - hidden_states = rearrange( - hidden_states, "(b f) l d-> b (f l) d", b=bsz, f=frame, l=len_frame - ).contiguous() - encoder_hidden_states, attn_mask = self.prepare_attn_mask( - encoder_attention_mask, encoder_hidden_states, q_seqlen=frame * len_frame - ) - - cos_sin = self._get_rope( - (frame, height, width), hidden_states.dtype, hidden_states.device - ) - - hidden_states = self.block_forward( - hidden_states, - encoder_hidden_states, - t_expand=t_expand, - rope_positions=[frame, height, width], - cos_sin=cos_sin, - attn_mask=attn_mask, - parallel=self.parallel, - mask_strategy=mask_strategy, - ) - - hidden_states = rearrange( - hidden_states, "b (f l) d -> (b f) l d", b=bsz, f=frame, l=len_frame - ) - - embedded_timestep = repeat( - embedded_timestep, "b d -> (b f) d", f=frame - ).contiguous() - - shift, scale = ( - self.scale_shift_table[None] + embedded_timestep[:, None] - ).chunk(2, dim=1) - hidden_states = self.norm_out( - hidden_states, shift=shift.squeeze(1), scale=scale.squeeze(1) - ) - # Modulation - hidden_states, _ = self.proj_out(hidden_states) - - # unpatchify - hidden_states = hidden_states.reshape( - shape=( - -1, - height, - width, - self.patch_size, - self.patch_size, - self.out_channels, - ) - ) - - hidden_states = rearrange(hidden_states, "n h w p q c -> n c h p w q") - output = hidden_states.reshape( - shape=( - -1, - self.out_channels, - height * self.patch_size, - width * self.patch_size, - ) - ) - - output = rearrange(output, "(b f) c h w -> b c f h w", f=frame) - return output - - -EntryClass = StepVideoModel diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/stepllm.py b/python/sglang/multimodal_gen/runtime/models/encoders/stepllm.py deleted file mode 100644 index 26e3a8670..000000000 --- a/python/sglang/multimodal_gen/runtime/models/encoders/stepllm.py +++ /dev/null @@ -1,617 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -# SPDX-License-Identifier: Apache-2.0 -# type: ignore -# Copyright 2025 StepFun Inc. All Rights Reserved. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# ============================================================================== -import os -from functools import wraps - -import torch -import torch.nn as nn -import torch.nn.functional as F -from einops import rearrange -from transformers import PretrainedConfig, PreTrainedModel - -from sglang.multimodal_gen.runtime.models.dits.stepvideo import StepVideoRMSNorm -from sglang.multimodal_gen.runtime.platforms import current_platform - - -class EmptyInitOnDevice(torch.overrides.TorchFunctionMode): - - def __init__(self, device=None): - self.device = device - - def __torch_function__(self, func, types, args=(), kwargs=None): - kwargs = kwargs or {} - if getattr(func, "__module__", None) == "torch.nn.init": - if "tensor" in kwargs: - return kwargs["tensor"] - else: - return args[0] - if ( - self.device is not None - and func in torch.utils._device._device_constructors() - and kwargs.get("device") is None - ): - kwargs["device"] = self.device - return func(*args, **kwargs) - - -def with_empty_init(func): - - @wraps(func) - def wrapper(*args, **kwargs): - with EmptyInitOnDevice("cpu"): - return func(*args, **kwargs) - - return wrapper - - -class LLaMaEmbedding(nn.Module): - """Language model embeddings. - - Arguments: - hidden_size: hidden size - vocab_size: vocabulary size - max_sequence_length: maximum size of sequence. This - is used for positional embedding - embedding_dropout_prob: dropout probability for embeddings - init_method: weight initialization method - num_tokentypes: size of the token-type embeddings. 0 value - will ignore this embedding - """ - - def __init__( - self, - cfg, - ): - super().__init__() - self.hidden_size = cfg.hidden_size - self.params_dtype = cfg.params_dtype - self.fp32_residual_connection = cfg.fp32_residual_connection - self.embedding_weights_in_fp32 = cfg.embedding_weights_in_fp32 - self.word_embeddings = torch.nn.Embedding( - cfg.padded_vocab_size, - self.hidden_size, - ) - self.embedding_dropout = torch.nn.Dropout(cfg.hidden_dropout) - - def forward(self, input_ids): - # Embeddings. - if self.embedding_weights_in_fp32: - self.word_embeddings = self.word_embeddings.to(torch.float32) - embeddings = self.word_embeddings(input_ids) - if self.embedding_weights_in_fp32: - embeddings = embeddings.to(self.params_dtype) - self.word_embeddings = self.word_embeddings.to(self.params_dtype) - - # Data format change to avoid explicit transposes : [b s h] --> [s b h]. - embeddings = embeddings.transpose(0, 1).contiguous() - - # If the input flag for fp32 residual connection is set, convert for float. - if self.fp32_residual_connection: - embeddings = embeddings.float() - - # Dropout. - embeddings = self.embedding_dropout(embeddings) - - return embeddings - - -class StepChatTokenizer: - """Step Chat Tokenizer""" - - def __init__( - self, - model_file, - name="StepChatTokenizer", - bot_token="<|BOT|>", # Begin of Turn - eot_token="<|EOT|>", # End of Turn - call_start_token="<|CALL_START|>", # Call Start - call_end_token="<|CALL_END|>", # Call End - think_start_token="<|THINK_START|>", # Think Start - think_end_token="<|THINK_END|>", # Think End - mask_start_token="<|MASK_1e69f|>", # Mask start - mask_end_token="<|UNMASK_1e69f|>", # Mask end - ): - import sentencepiece - - self._tokenizer = sentencepiece.SentencePieceProcessor(model_file=model_file) - - self._vocab = {} - self._inv_vocab = {} - - self._special_tokens = {} - self._inv_special_tokens = {} - - self._t5_tokens = [] - - for idx in range(self._tokenizer.get_piece_size()): - text = self._tokenizer.id_to_piece(idx) - self._inv_vocab[idx] = text - self._vocab[text] = idx - - if self._tokenizer.is_control(idx) or self._tokenizer.is_unknown(idx): - self._special_tokens[text] = idx - self._inv_special_tokens[idx] = text - - self._unk_id = self._tokenizer.unk_id() - self._bos_id = self._tokenizer.bos_id() - self._eos_id = self._tokenizer.eos_id() - - for token in [ - bot_token, - eot_token, - call_start_token, - call_end_token, - think_start_token, - think_end_token, - ]: - assert token in self._vocab, f"Token '{token}' not found in tokenizer" - assert ( - token in self._special_tokens - ), f"Token '{token}' is not a special token" - - for token in [mask_start_token, mask_end_token]: - assert token in self._vocab, f"Token '{token}' not found in tokenizer" - - self._bot_id = self._tokenizer.piece_to_id(bot_token) - self._eot_id = self._tokenizer.piece_to_id(eot_token) - self._call_start_id = self._tokenizer.piece_to_id(call_start_token) - self._call_end_id = self._tokenizer.piece_to_id(call_end_token) - self._think_start_id = self._tokenizer.piece_to_id(think_start_token) - self._think_end_id = self._tokenizer.piece_to_id(think_end_token) - self._mask_start_id = self._tokenizer.piece_to_id(mask_start_token) - self._mask_end_id = self._tokenizer.piece_to_id(mask_end_token) - - self._underline_id = self._tokenizer.piece_to_id("\u2581") - - @property - def vocab(self): - return self._vocab - - @property - def inv_vocab(self): - return self._inv_vocab - - @property - def vocab_size(self): - return self._tokenizer.vocab_size() - - def tokenize(self, text: str) -> list[int]: - return self._tokenizer.encode_as_ids(text) - - def detokenize(self, token_ids: list[int]) -> str: - return self._tokenizer.decode_ids(token_ids) - - -class Tokens: - - def __init__( - self, input_ids, cu_input_ids, attention_mask, cu_seqlens, max_seq_len - ) -> None: - self.input_ids = input_ids - self.attention_mask = attention_mask - self.cu_input_ids = cu_input_ids - self.cu_seqlens = cu_seqlens - self.max_seq_len = max_seq_len - - def to(self, device): - self.input_ids = self.input_ids.to(device) - self.attention_mask = self.attention_mask.to(device) - self.cu_input_ids = self.cu_input_ids.to(device) - self.cu_seqlens = self.cu_seqlens.to(device) - return self - - -class Wrapped_StepChatTokenizer(StepChatTokenizer): - - def __call__( - self, - text, - max_length=320, - padding="max_length", - truncation=True, - return_tensors="pt", - ): - # [bos, ..., eos, pad, pad, ..., pad] - self.BOS = 1 - self.EOS = 2 - self.PAD = 2 - out_tokens = [] - attn_mask = [] - if len(text) == 0: - part_tokens = [self.BOS] + [self.EOS] - valid_size = len(part_tokens) - if len(part_tokens) < max_length: - part_tokens += [self.PAD] * (max_length - valid_size) - out_tokens.append(part_tokens) - attn_mask.append([1] * valid_size + [0] * (max_length - valid_size)) - else: - for part in text: - part_tokens = self.tokenize(part) - part_tokens = part_tokens[ - : (max_length - 2) - ] # leave 2 space for bos and eos - part_tokens = [self.BOS] + part_tokens + [self.EOS] - valid_size = len(part_tokens) - if len(part_tokens) < max_length: - part_tokens += [self.PAD] * (max_length - valid_size) - out_tokens.append(part_tokens) - attn_mask.append([1] * valid_size + [0] * (max_length - valid_size)) - - out_tokens = torch.tensor(out_tokens, dtype=torch.long) - attn_mask = torch.tensor(attn_mask, dtype=torch.long) - - # padding y based on tp size - padded_len = 0 - padded_flag = False - if padded_len > 0: - padded_flag = True - if padded_flag: - pad_tokens = torch.tensor( - [[self.PAD] * max_length], device=out_tokens.device - ) - pad_attn_mask = torch.tensor( - [[1] * padded_len + [0] * (max_length - padded_len)], - device=attn_mask.device, - ) - out_tokens = torch.cat([out_tokens, pad_tokens], dim=0) - attn_mask = torch.cat([attn_mask, pad_attn_mask], dim=0) - - # cu_seqlens - cu_out_tokens = out_tokens.masked_select(attn_mask != 0).unsqueeze(0) - seqlen = attn_mask.sum(dim=1).tolist() - cu_seqlens = torch.cumsum(torch.tensor([0] + seqlen), 0).to( - device=out_tokens.device, dtype=torch.int32 - ) - max_seq_len = max(seqlen) - return Tokens(out_tokens, cu_out_tokens, attn_mask, cu_seqlens, max_seq_len) - - -def flash_attn_func( - q, - k, - v, - dropout_p=0.0, - softmax_scale=None, - causal=True, - return_attn_probs=False, - tp_group_rank=0, - tp_group_size=1, -): - softmax_scale = q.size(-1) ** (-0.5) if softmax_scale is None else softmax_scale - return torch.ops.Optimus.fwd( - q, - k, - v, - None, - dropout_p, - softmax_scale, - causal, - return_attn_probs, - None, - tp_group_rank, - tp_group_size, - )[0] - - -class FlashSelfAttention(torch.nn.Module): - - def __init__( - self, - attention_dropout=0.0, - ): - super().__init__() - self.dropout_p = attention_dropout - - def forward(self, q, k, v, cu_seqlens=None, max_seq_len=None): - if cu_seqlens is None: - output = flash_attn_func(q, k, v, dropout_p=self.dropout_p) - else: - raise ValueError("cu_seqlens is not supported!") - - return output - - -def safediv(n, d): - q, r = divmod(n, d) - assert r == 0 - return q - - -class MultiQueryAttention(nn.Module): - - def __init__(self, cfg, layer_id=None): - super().__init__() - - self.head_dim = cfg.hidden_size // cfg.num_attention_heads - self.max_seq_len = cfg.seq_length - self.use_flash_attention = cfg.use_flash_attn - assert self.use_flash_attention, "FlashAttention is required!" - - self.n_groups = cfg.num_attention_groups - self.tp_size = 1 - self.n_local_heads = cfg.num_attention_heads - self.n_local_groups = self.n_groups - - self.wqkv = nn.Linear( - cfg.hidden_size, - cfg.hidden_size + self.head_dim * 2 * self.n_groups, - bias=False, - ) - self.wo = nn.Linear( - cfg.hidden_size, - cfg.hidden_size, - bias=False, - ) - - # assert self.use_flash_attention, 'non-Flash attention not supported yet.' - self.core_attention = FlashSelfAttention( - attention_dropout=cfg.attention_dropout - ) - # self.core_attention = LocalAttention( - # num_heads = self.n_local_heads, - # head_size = self.head_dim, - # # num_kv_heads = self.n_local_groups, - # casual = True, - # supported_attention_backends = [_Backend.FLASH_ATTN, _Backend.TORCH_SDPA], # RIVER TODO - # ) - self.layer_id = layer_id - - def forward( - self, - x: torch.Tensor, - mask: torch.Tensor | None, - cu_seqlens: torch.Tensor | None, - max_seq_len: torch.Tensor | None, - ): - seqlen, bsz, dim = x.shape - xqkv = self.wqkv(x) - - xq, xkv = torch.split( - xqkv, - (dim // self.tp_size, self.head_dim * 2 * self.n_groups // self.tp_size), - dim=-1, - ) - - # gather on 1st dimension - xq = xq.view(seqlen, bsz, self.n_local_heads, self.head_dim) - xkv = xkv.view(seqlen, bsz, self.n_local_groups, 2 * self.head_dim) - xk, xv = xkv.chunk(2, -1) - - # rotary embedding + flash attn - xq = rearrange(xq, "s b h d -> b s h d") - xk = rearrange(xk, "s b h d -> b s h d") - xv = rearrange(xv, "s b h d -> b s h d") - - # q_per_kv = self.n_local_heads // self.n_local_groups - # if q_per_kv > 1: - # b, s, h, d = xk.size() - # if h == 1: - # xk = xk.expand(b, s, q_per_kv, d) - # xv = xv.expand(b, s, q_per_kv, d) - # else: - # ''' To cover the cases where h > 1, we have - # the following implementation, which is equivalent to: - # xk = xk.repeat_interleave(q_per_kv, dim=-2) - # xv = xv.repeat_interleave(q_per_kv, dim=-2) - # but can avoid calling aten::item() that involves cpu. - # ''' - # idx = torch.arange(q_per_kv * h, device=xk.device).reshape(q_per_kv, -1).permute(1, 0).flatten() - # xk = torch.index_select(xk.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous() - # xv = torch.index_select(xv.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous() - if self.use_flash_attention: - output = self.core_attention(xq, xk, xv) - # reduce-scatter only support first dimension now - output = rearrange(output, "b s h d -> s b (h d)").contiguous() - else: - xq, xk, xv = [ - rearrange(x, "b s ... -> s b ...").contiguous() for x in (xq, xk, xv) - ] - output = self.core_attention(xq, xk, xv) # , mask) - output = self.wo(output) - return output - - -class FeedForward(nn.Module): - - def __init__( - self, - cfg, - dim: int, - hidden_dim: int, - layer_id: int, - multiple_of: int = 256, - ): - super().__init__() - - hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) - - def swiglu(x): - x = torch.chunk(x, 2, dim=-1) - return F.silu(x[0]) * x[1] - - self.swiglu = swiglu - - self.w1 = nn.Linear( - dim, - 2 * hidden_dim, - bias=False, - ) - self.w2 = nn.Linear( - hidden_dim, - dim, - bias=False, - ) - - def forward(self, x): - x = self.swiglu(self.w1(x)) - output = self.w2(x) - return output - - -class TransformerBlock(nn.Module): - - def __init__(self, cfg, layer_id: int): - super().__init__() - - self.n_heads = cfg.num_attention_heads - self.dim = cfg.hidden_size - self.head_dim = cfg.hidden_size // cfg.num_attention_heads - self.attention = MultiQueryAttention( - cfg, - layer_id=layer_id, - ) - - self.feed_forward = FeedForward( - cfg, - dim=cfg.hidden_size, - hidden_dim=cfg.ffn_hidden_size, - layer_id=layer_id, - ) - self.layer_id = layer_id - self.attention_norm = StepVideoRMSNorm( - cfg.hidden_size, - eps=cfg.layernorm_epsilon, - ) - self.ffn_norm = StepVideoRMSNorm( - cfg.hidden_size, - eps=cfg.layernorm_epsilon, - ) - - def forward( - self, - x: torch.Tensor, - mask: torch.Tensor | None, - cu_seqlens: torch.Tensor | None, - max_seq_len: torch.Tensor | None, - ): - residual = self.attention.forward( - self.attention_norm(x), mask, cu_seqlens, max_seq_len - ) - h = x + residual - ffn_res = self.feed_forward.forward(self.ffn_norm(h)) - out = h + ffn_res - return out - - -class Transformer(nn.Module): - - def __init__( - self, - config, - max_seq_size=8192, - ): - super().__init__() - self.num_layers = config.num_layers - self.layers = self._build_layers(config) - - def _build_layers(self, config): - layers = torch.nn.ModuleList() - for layer_id in range(self.num_layers): - layers.append( - TransformerBlock( - config, - layer_id=layer_id + 1, - ) - ) - return layers - - def forward( - self, - hidden_states, - attention_mask, - cu_seqlens=None, - max_seq_len=None, - ): - - if max_seq_len is not None and not isinstance(max_seq_len, torch.Tensor): - max_seq_len = torch.tensor(max_seq_len, dtype=torch.int32, device="cpu") - - for lid, layer in enumerate(self.layers): - hidden_states = layer( - hidden_states, - attention_mask, - cu_seqlens, - max_seq_len, - ) - return hidden_states - - -class Step1Model(PreTrainedModel): - config_class = PretrainedConfig - - @with_empty_init - def __init__( - self, - config, - ): - super().__init__(config) - self.tok_embeddings = LLaMaEmbedding(config) - self.transformer = Transformer(config) - - def forward( - self, - input_ids=None, - attention_mask=None, - ): - - hidden_states = self.tok_embeddings(input_ids) - - hidden_states = self.transformer( - hidden_states, - attention_mask, - ) - return hidden_states - - -class STEP1TextEncoder(torch.nn.Module): - - def __init__(self, model_dir, max_length=320): - super().__init__() - self.max_length = max_length - self.text_tokenizer = Wrapped_StepChatTokenizer( - os.path.join(model_dir, "step1_chat_tokenizer.model") - ) - text_encoder = Step1Model.from_pretrained(model_dir) - self.text_encoder = text_encoder.eval().to(torch.bfloat16) - - @torch.no_grad - def forward(self, prompts, with_mask=True, max_length=None): - self.device = next(self.text_encoder.parameters()).device - - with torch.no_grad(), torch.amp.autocast( - current_platform.device_type, dtype=torch.bfloat16 - ): - if type(prompts) is str: - prompts = [prompts] - txt_tokens = self.text_tokenizer( - prompts, - max_length=max_length or self.max_length, - padding="max_length", - truncation=True, - return_tensors="pt", - ) - y = self.text_encoder( - txt_tokens.input_ids.to(self.device), - attention_mask=( - txt_tokens.attention_mask.to(self.device) if with_mask else None - ), - ) - y_mask = txt_tokens.attention_mask - return y.transpose(0, 1), y_mask - - -EntryClass = STEP1TextEncoder diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/stepvideovae.py b/python/sglang/multimodal_gen/runtime/models/vaes/stepvideovae.py deleted file mode 100644 index d202b7a61..000000000 --- a/python/sglang/multimodal_gen/runtime/models/vaes/stepvideovae.py +++ /dev/null @@ -1,1184 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -# SPDX-License-Identifier: Apache-2.0 -# Copyright 2025 StepFun Inc. All Rights Reserved. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# ============================================================================== -from typing import Any - -import torch -from einops import rearrange -from torch import nn -from torch.nn import functional as F - -from sglang.multimodal_gen.configs.models.vaes import StepVideoVAEConfig -from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE - - -def base_group_norm(x, norm_layer, act_silu=False, channel_last=False) -> torch.Tensor: - if hasattr(base_group_norm, "spatial") and base_group_norm.spatial: - assert channel_last - x_shape = x.shape - x = x.flatten(0, 1) - if channel_last: - # Permute to NCHW format - x = x.permute(0, 3, 1, 2) - - out = F.group_norm( - x.contiguous(), - norm_layer.num_groups, - norm_layer.weight, - norm_layer.bias, - norm_layer.eps, - ) - if act_silu: - out = F.silu(out) - - if channel_last: - # Permute back to NHWC format - out = out.permute(0, 2, 3, 1) - - out = out.view(x_shape) - else: - if channel_last: - # Permute to NCHW format - x = x.permute(0, 3, 1, 2) - out = F.group_norm( - x.contiguous(), - norm_layer.num_groups, - norm_layer.weight, - norm_layer.bias, - norm_layer.eps, - ) - if act_silu: - out = F.silu(out) - if channel_last: - # Permute back to NHWC format - out = out.permute(0, 2, 3, 1) - return out - - -def base_conv2d(x, conv_layer, channel_last=False, residual=None) -> torch.Tensor: - if channel_last: - x = x.permute(0, 3, 1, 2) # NHWC to NCHW - out = F.conv2d( - x, - conv_layer.weight, - conv_layer.bias, - stride=conv_layer.stride, - padding=conv_layer.padding, - ) - if residual is not None: - if channel_last: - residual = residual.permute(0, 3, 1, 2) # NHWC to NCHW - out += residual - if channel_last: - out = out.permute(0, 2, 3, 1) # NCHW to NHWC - return out - - -def base_conv3d( - x, conv_layer, channel_last=False, residual=None, only_return_output=False -) -> torch.Tensor: - if only_return_output: - size = cal_outsize( - x.shape, conv_layer.weight.shape, conv_layer.stride, conv_layer.padding - ) - return torch.empty(size, device=x.device, dtype=x.dtype) - if channel_last: - x = x.permute(0, 4, 1, 2, 3) # NDHWC to NCDHW - out = F.conv3d( - x, - conv_layer.weight, - conv_layer.bias, - stride=conv_layer.stride, - padding=conv_layer.padding, - ) - if residual is not None: - if channel_last: - residual = residual.permute(0, 4, 1, 2, 3) # NDHWC to NCDHW - out += residual - if channel_last: - out = out.permute(0, 2, 3, 4, 1) # NCDHW to NDHWC - return out - - -def cal_outsize(input_sizes, kernel_sizes, stride, padding) -> list: - stride_d, stride_h, stride_w = stride - padding_d, padding_h, padding_w = padding - dilation_d, dilation_h, dilation_w = 1, 1, 1 - - in_d = input_sizes[1] - in_h = input_sizes[2] - in_w = input_sizes[3] - - kernel_d = kernel_sizes[2] - kernel_h = kernel_sizes[3] - kernel_w = kernel_sizes[4] - out_channels = kernel_sizes[0] - - out_d = calc_out_(in_d, padding_d, dilation_d, kernel_d, stride_d) - out_h = calc_out_(in_h, padding_h, dilation_h, kernel_h, stride_h) - out_w = calc_out_(in_w, padding_w, dilation_w, kernel_w, stride_w) - size = [input_sizes[0], out_d, out_h, out_w, out_channels] - return size - - -def calc_out_( - in_size: int, padding: int, dilation: int, kernel: int, stride: int -) -> int: - return (in_size + 2 * padding - dilation * (kernel - 1) - 1) // stride + 1 - - -def base_conv3d_channel_last(x, conv_layer, residual=None) -> torch.Tensor: - in_numel = x.numel() - out_numel = int(x.numel() * conv_layer.out_channels / conv_layer.in_channels) - if (in_numel >= 2**30) or (out_numel >= 2**30): - assert conv_layer.stride[0] == 1, "time split asks time stride = 1" - - B, T, H, W, C = x.shape - K = conv_layer.kernel_size[0] - - chunks = 4 - chunk_size = T // chunks - - if residual is None: - out_nhwc = base_conv3d( - x, - conv_layer, - channel_last=True, - residual=residual, - only_return_output=True, - ) - else: - out_nhwc = residual - - assert B == 1 - for i in range(chunks): - if i == chunks - 1: - xi = x[:1, chunk_size * i :] - out_nhwci = out_nhwc[:1, chunk_size * i :] - else: - xi = x[:1, chunk_size * i : chunk_size * (i + 1) + K - 1] - out_nhwci = out_nhwc[:1, chunk_size * i : chunk_size * (i + 1)] - if residual is not None: - if i == chunks - 1: - ri = residual[:1, chunk_size * i :] - else: - ri = residual[:1, chunk_size * i : chunk_size * (i + 1)] - else: - ri = None - out_nhwci.copy_(base_conv3d(xi, conv_layer, channel_last=True, residual=ri)) - else: - out_nhwc = base_conv3d(x, conv_layer, channel_last=True, residual=residual) - return out_nhwc - - -class Upsample2D(nn.Module): - - def __init__( - self, channels, use_conv=False, use_conv_transpose=False, out_channels=None - ) -> None: - super().__init__() - self.channels = channels - self.out_channels = out_channels or channels - self.use_conv = use_conv - self.use_conv_transpose = use_conv_transpose - - if use_conv: - self.conv = nn.Conv2d(self.channels, self.out_channels, 3, padding=1) - else: - assert "Not Supported" - self.conv = nn.ConvTranspose2d(channels, self.out_channels, 4, 2, 1) - - def forward(self, x, output_size=None) -> torch.Tensor: - assert x.shape[-1] == self.channels - - if self.use_conv_transpose: - return self.conv(x) - - if output_size is None: - x = ( - F.interpolate( - x.permute(0, 3, 1, 2).to(memory_format=torch.channels_last), - scale_factor=2.0, - mode="nearest", - ) - .permute(0, 2, 3, 1) - .contiguous() - ) - else: - x = ( - F.interpolate( - x.permute(0, 3, 1, 2).to(memory_format=torch.channels_last), - size=output_size, - mode="nearest", - ) - .permute(0, 2, 3, 1) - .contiguous() - ) - - # x = self.conv(x) - x = base_conv2d(x, self.conv, channel_last=True) - return x - - -class Downsample2D(nn.Module): - - def __init__(self, channels, use_conv=False, out_channels=None, padding=1) -> None: - super().__init__() - self.channels = channels - self.out_channels = out_channels or channels - self.use_conv = use_conv - self.padding = padding - stride = 2 - - if use_conv: - self.conv = nn.Conv2d( - self.channels, self.out_channels, 3, stride=stride, padding=padding - ) - else: - assert self.channels == self.out_channels - self.conv = nn.AvgPool2d(kernel_size=stride, stride=stride) - - def forward(self, x) -> torch.Tensor: - assert x.shape[-1] == self.channels - if self.use_conv and self.padding == 0: - pad = (0, 0, 0, 1, 0, 1) - x = F.pad(x, pad, mode="constant", value=0) - - assert x.shape[-1] == self.channels - # x = self.conv(x) - x = base_conv2d(x, self.conv, channel_last=True) - return x - - -class CausalConv(nn.Module): - - def __init__(self, chan_in, chan_out, kernel_size, **kwargs) -> None: - super().__init__() - - if isinstance(kernel_size, int): - kernel_size = ( - kernel_size if isinstance(kernel_size, tuple) else ((kernel_size,) * 3) - ) - time_kernel_size, height_kernel_size, width_kernel_size = kernel_size - - self.dilation = kwargs.pop("dilation", 1) - self.stride = kwargs.pop("stride", 1) - if isinstance(self.stride, int): - self.stride = (self.stride, 1, 1) - time_pad = self.dilation * (time_kernel_size - 1) + max((1 - self.stride[0]), 0) - height_pad = height_kernel_size // 2 - width_pad = width_kernel_size // 2 - self.time_causal_padding = ( - width_pad, - width_pad, - height_pad, - height_pad, - time_pad, - 0, - ) - self.time_uncausal_padding = ( - width_pad, - width_pad, - height_pad, - height_pad, - 0, - 0, - ) - - self.conv = nn.Conv3d( - chan_in, - chan_out, - kernel_size, - stride=self.stride, - dilation=self.dilation, - **kwargs, - ) - self.chan_in = chan_in - self.chan_out = chan_out - self.is_first_run = True - - def forward(self, x, is_init=True, residual=None) -> torch.Tensor: - x = nn.functional.pad( - x, self.time_causal_padding if is_init else self.time_uncausal_padding - ) - x = self.conv(x) - if residual is not None: - x.add_(residual) - return x - - -class ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(nn.Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - factor: int, - ) -> None: - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.factor = factor - assert out_channels * factor**3 % in_channels == 0 - self.repeats = out_channels * factor**3 // in_channels - - def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: - x = x.repeat_interleave(self.repeats, dim=1) - x = x.view( - x.size(0), - self.out_channels, - self.factor, - self.factor, - self.factor, - x.size(2), - x.size(3), - x.size(4), - ) - x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() - x = x.view( - x.size(0), - self.out_channels, - x.size(2) * self.factor, - x.size(4) * self.factor, - x.size(6) * self.factor, - ) - x = x[:, :, self.factor - 1 :, :, :] - return x - - -class ConvPixelShuffleUpSampleLayer3D(nn.Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int, - factor: int, - ) -> None: - super().__init__() - self.factor = factor - out_ratio = factor**3 - self.conv = CausalConv( - in_channels, out_channels * out_ratio, kernel_size=kernel_size - ) - - def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: - x = self.conv(x, is_init) - x = self.pixel_shuffle_3d(x, self.factor) - return x - - @staticmethod - def pixel_shuffle_3d(x: torch.Tensor, factor: int) -> torch.Tensor: - batch_size, channels, depth, height, width = x.size() - new_channels = channels // (factor**3) - new_depth = depth * factor - new_height = height * factor - new_width = width * factor - - x = x.view( - batch_size, new_channels, factor, factor, factor, depth, height, width - ) - x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() - x = x.view(batch_size, new_channels, new_depth, new_height, new_width) - x = x[:, :, factor - 1 :, :, :] - return x - - -class ConvPixelUnshuffleDownSampleLayer3D(nn.Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int, - factor: int, - ) -> None: - super().__init__() - self.factor = factor - out_ratio = factor**3 - assert out_channels % out_ratio == 0 - self.conv = CausalConv( - in_channels, out_channels // out_ratio, kernel_size=kernel_size - ) - - def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: - x = self.conv(x, is_init) - x = self.pixel_unshuffle_3d(x, self.factor) - return x - - @staticmethod - def pixel_unshuffle_3d(x: torch.Tensor, factor: int) -> torch.Tensor: - pad = (0, 0, 0, 0, factor - 1, 0) # (left, right, top, bottom, front, back) - x = F.pad(x, pad) - B, C, D, H, W = x.shape - x = x.view(B, C, D // factor, factor, H // factor, factor, W // factor, factor) - x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() - x = x.view(B, C * factor**3, D // factor, H // factor, W // factor) - return x - - -class PixelUnshuffleChannelAveragingDownSampleLayer3D(nn.Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - factor: int, - ) -> None: - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.factor = factor - assert in_channels * factor**3 % out_channels == 0 - self.group_size = in_channels * factor**3 // out_channels - - def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: - pad = ( - 0, - 0, - 0, - 0, - self.factor - 1, - 0, - ) # (left, right, top, bottom, front, back) - x = F.pad(x, pad) - B, C, D, H, W = x.shape - x = x.view( - B, - C, - D // self.factor, - self.factor, - H // self.factor, - self.factor, - W // self.factor, - self.factor, - ) - x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() - x = x.view( - B, C * self.factor**3, D // self.factor, H // self.factor, W // self.factor - ) - x = x.view( - B, - self.out_channels, - self.group_size, - D // self.factor, - H // self.factor, - W // self.factor, - ) - x = x.mean(dim=2) - return x - - -def base_group_norm_with_zero_pad( - x, norm_layer, act_silu=True, pad_size=2 -) -> torch.Tensor: - out_shape = list(x.shape) - out_shape[1] += pad_size - out = torch.empty(out_shape, dtype=x.dtype, device=x.device) - out[:, pad_size:] = base_group_norm( - x, norm_layer, act_silu=act_silu, channel_last=True - ) - out[:, :pad_size] = 0 - return out - - -class CausalConvChannelLast(CausalConv): - time_causal_padding: tuple[Any, ...] - time_uncausal_padding: tuple[Any, ...] - - def __init__(self, chan_in, chan_out, kernel_size, **kwargs) -> None: - super().__init__(chan_in, chan_out, kernel_size, **kwargs) - - self.time_causal_padding = (0, 0) + self.time_causal_padding - self.time_uncausal_padding = (0, 0) + self.time_uncausal_padding - - def forward(self, x, is_init=True, residual=None) -> torch.Tensor: - if self.is_first_run: - self.is_first_run = False - # self.conv.weight = nn.Parameter(self.conv.weight.permute(0,2,3,4,1).contiguous()) - - x = nn.functional.pad( - x, self.time_causal_padding if is_init else self.time_uncausal_padding - ) - - x = base_conv3d_channel_last(x, self.conv, residual=residual) - return x - - -class CausalConvAfterNorm(CausalConv): - - def __init__(self, chan_in, chan_out, kernel_size, **kwargs) -> None: - super().__init__(chan_in, chan_out, kernel_size, **kwargs) - - if self.time_causal_padding == (1, 1, 1, 1, 2, 0): - self.conv = nn.Conv3d( - chan_in, - chan_out, - kernel_size, - stride=self.stride, - dilation=self.dilation, - padding=(0, 1, 1), - **kwargs, - ) - else: - self.conv = nn.Conv3d( - chan_in, - chan_out, - kernel_size, - stride=self.stride, - dilation=self.dilation, - **kwargs, - ) - self.is_first_run = True - - def forward(self, x, is_init=True, residual=None) -> torch.Tensor: - if self.is_first_run: - self.is_first_run = False - - if self.time_causal_padding == (1, 1, 1, 1, 2, 0): - pass - else: - x = nn.functional.pad(x, self.time_causal_padding).contiguous() - - x = base_conv3d_channel_last(x, self.conv, residual=residual) - return x - - -class AttnBlock(nn.Module): - - def __init__(self, in_channels) -> None: - super().__init__() - - self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels) - self.q = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) - self.k = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) - self.v = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) - self.proj_out = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) - - def attention(self, x, is_init=True) -> torch.Tensor: - x = base_group_norm(x, self.norm, act_silu=False, channel_last=True) - q = self.q(x, is_init) - k = self.k(x, is_init) - v = self.v(x, is_init) - - b, t, h, w, c = q.shape - q, k, v = map(lambda x: rearrange(x, "b t h w c -> b 1 (t h w) c"), (q, k, v)) - x = nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True) - x = rearrange(x, "b 1 (t h w) c -> b t h w c", t=t, h=h, w=w) - - return x - - def forward(self, x): - x = x.permute(0, 2, 3, 4, 1).contiguous() - h = self.attention(x) - x = self.proj_out(h, residual=x) - x = x.permute(0, 4, 1, 2, 3) - return x - - -class Resnet3DBlock(nn.Module): - - def __init__( - self, - in_channels, - out_channels=None, - temb_channels=512, - conv_shortcut=False, - ) -> None: - super().__init__() - - self.in_channels = in_channels - out_channels = in_channels if out_channels is None else out_channels - self.out_channels = out_channels - - self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels) - self.conv1 = CausalConvAfterNorm(in_channels, out_channels, kernel_size=3) - if temb_channels > 0: - self.temb_proj = nn.Linear(temb_channels, out_channels) - - self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels) - self.conv2 = CausalConvAfterNorm(out_channels, out_channels, kernel_size=3) - - assert conv_shortcut is False - self.use_conv_shortcut = conv_shortcut - if self.in_channels != self.out_channels: - if self.use_conv_shortcut: - self.conv_shortcut = CausalConvAfterNorm( - in_channels, out_channels, kernel_size=3 - ) - else: - self.nin_shortcut = CausalConvAfterNorm( - in_channels, out_channels, kernel_size=1 - ) - - def forward(self, x, temb=None, is_init=True) -> torch.Tensor: - x = x.permute(0, 2, 3, 4, 1).contiguous() - - h = base_group_norm_with_zero_pad(x, self.norm1, act_silu=True, pad_size=2) - h = self.conv1(h) - if temb is not None: - h = h + self.temb_proj(nn.functional.silu(temb))[:, :, None, None] - - x = self.nin_shortcut(x) if self.in_channels != self.out_channels else x - - h = base_group_norm_with_zero_pad(h, self.norm2, act_silu=True, pad_size=2) - x = self.conv2(h, residual=x) - - x = x.permute(0, 4, 1, 2, 3) - return x - - -class Downsample3D(nn.Module): - - def __init__(self, in_channels, with_conv, stride) -> None: - super().__init__() - - self.with_conv = with_conv - if with_conv: - self.conv = CausalConv( - in_channels, in_channels, kernel_size=3, stride=stride - ) - - def forward(self, x, is_init=True) -> torch.Tensor: - if self.with_conv: - x = self.conv(x, is_init) - else: - x = nn.functional.avg_pool3d(x, kernel_size=2, stride=2) - return x - - -class VideoEncoder(nn.Module): - - def __init__( - self, - ch=32, - ch_mult=(4, 8, 16, 16), - num_res_blocks=2, - in_channels=3, - z_channels=16, - double_z=True, - down_sampling_layer=(1, 2), - resamp_with_conv=True, - version=1, - ) -> None: - super().__init__() - - temb_ch = 0 - - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - - # downsampling - self.conv_in = CausalConv(in_channels, ch, kernel_size=3) - self.down_sampling_layer = down_sampling_layer - - in_ch_mult = (1,) + tuple(ch_mult) - self.down = nn.ModuleList() - for i_level in range(self.num_resolutions): - block = nn.ModuleList() - attn = nn.ModuleList() - block_in = ch * in_ch_mult[i_level] - block_out = ch * ch_mult[i_level] - for i_block in range(self.num_res_blocks): - block.append( - Resnet3DBlock( - in_channels=block_in, - out_channels=block_out, - temb_channels=temb_ch, - ) - ) - block_in = block_out - down = nn.Module() - down.block = block - down.attn = attn - if i_level != self.num_resolutions - 1: - if i_level in self.down_sampling_layer: - down.downsample = Downsample3D( - block_in, resamp_with_conv, stride=(2, 2, 2) - ) - else: - down.downsample = Downsample2D( - block_in, resamp_with_conv, padding=0 - ) # DIFF - self.down.append(down) - - # middle - self.mid = nn.Module() - self.mid.block_1 = Resnet3DBlock( - in_channels=block_in, out_channels=block_in, temb_channels=temb_ch - ) - self.mid.attn_1 = AttnBlock(block_in) - self.mid.block_2 = Resnet3DBlock( - in_channels=block_in, out_channels=block_in, temb_channels=temb_ch - ) - - # end - self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in) - self.version = version - if version == 2: - channels = 4 * z_channels * 2**3 - self.conv_patchify = ConvPixelUnshuffleDownSampleLayer3D( - block_in, channels, kernel_size=3, factor=2 - ) - self.shortcut_pathify = PixelUnshuffleChannelAveragingDownSampleLayer3D( - block_in, channels, 2 - ) - self.shortcut_out = PixelUnshuffleChannelAveragingDownSampleLayer3D( - channels, 2 * z_channels if double_z else z_channels, 1 - ) - self.conv_out = CausalConvChannelLast( - channels, 2 * z_channels if double_z else z_channels, kernel_size=3 - ) - else: - self.conv_out = CausalConvAfterNorm( - block_in, 2 * z_channels if double_z else z_channels, kernel_size=3 - ) - - @torch.inference_mode() - def forward(self, x, video_frame_num, is_init=True) -> torch.Tensor: - # timestep embedding - temb = None - - t = video_frame_num - - # downsampling - h = self.conv_in(x, is_init) - - # make it real channel last, but behave like normal layout - h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) - - for i_level in range(self.num_resolutions): - for i_block in range(self.num_res_blocks): - h = self.down[i_level].block[i_block](h, temb, is_init) - if len(self.down[i_level].attn) > 0: - h = self.down[i_level].attn[i_block](h) - - if i_level != self.num_resolutions - 1: - if isinstance(self.down[i_level].downsample, Downsample2D): - _, _, t, _, _ = h.shape - h = rearrange(h, "b c t h w -> (b t) h w c", t=t) - h = self.down[i_level].downsample(h) - h = rearrange(h, "(b t) h w c -> b c t h w", t=t) - else: - h = self.down[i_level].downsample(h, is_init) - - h = self.mid.block_1(h, temb, is_init) - h = self.mid.attn_1(h) - h = self.mid.block_2(h, temb, is_init) - - h = h.permute(0, 2, 3, 4, 1).contiguous() # b c l h w -> b l h w c - if self.version == 2: - h = base_group_norm(h, self.norm_out, act_silu=True, channel_last=True) - h = h.permute(0, 4, 1, 2, 3).contiguous() - shortcut = self.shortcut_pathify(h, is_init) - h = self.conv_patchify(h, is_init) - h = h.add_(shortcut) - shortcut = self.shortcut_out(h, is_init).permute(0, 2, 3, 4, 1) - h = self.conv_out(h.permute(0, 2, 3, 4, 1).contiguous(), is_init) - h = h.add_(shortcut) - else: - h = base_group_norm_with_zero_pad( - h, self.norm_out, act_silu=True, pad_size=2 - ) - h = self.conv_out(h, is_init) - h = h.permute(0, 4, 1, 2, 3) # b l h w c -> b c l h w - - h = rearrange(h, "b c t h w -> b t c h w") - return h - - -class Res3DBlockUpsample(nn.Module): - - def __init__( - self, input_filters, num_filters, down_sampling_stride, down_sampling=False - ) -> None: - super().__init__() - - self.input_filters = input_filters - self.num_filters = num_filters - - self.act_ = nn.SiLU(inplace=True) - - self.conv1 = CausalConvChannelLast( - num_filters, num_filters, kernel_size=[3, 3, 3] - ) - self.norm1 = nn.GroupNorm(32, num_filters) - - self.conv2 = CausalConvChannelLast( - num_filters, num_filters, kernel_size=[3, 3, 3] - ) - self.norm2 = nn.GroupNorm(32, num_filters) - - self.down_sampling = down_sampling - if down_sampling: - self.down_sampling_stride = down_sampling_stride - else: - self.down_sampling_stride = [1, 1, 1] - - if num_filters != input_filters or down_sampling: - self.conv3 = CausalConvChannelLast( - input_filters, - num_filters, - kernel_size=[1, 1, 1], - stride=self.down_sampling_stride, - ) - self.norm3 = nn.GroupNorm(32, num_filters) - - def forward(self, x, is_init=False) -> torch.Tensor: - x = x.permute(0, 2, 3, 4, 1).contiguous() - - residual = x - - h = self.conv1(x, is_init) - h = base_group_norm(h, self.norm1, act_silu=True, channel_last=True) - - h = self.conv2(h, is_init) - h = base_group_norm(h, self.norm2, act_silu=False, channel_last=True) - - if self.down_sampling or self.num_filters != self.input_filters: - x = self.conv3(x, is_init) - x = base_group_norm(x, self.norm3, act_silu=False, channel_last=True) - - h.add_(x) - h = self.act_(h) - if residual is not None: - h.add_(residual) - - h = h.permute(0, 4, 1, 2, 3) - return h - - -class Upsample3D(nn.Module): - - def __init__(self, in_channels, scale_factor=2) -> None: - super().__init__() - - self.scale_factor = scale_factor - self.conv3d = Res3DBlockUpsample( - input_filters=in_channels, - num_filters=in_channels, - down_sampling_stride=(1, 1, 1), - down_sampling=False, - ) - - def forward(self, x, is_init=True, is_split=True) -> torch.Tensor: - b, c, t, h, w = x.shape - - # x = x.permute(0,2,3,4,1).contiguous().permute(0,4,1,2,3).to(memory_format=torch.channels_last_3d) - if is_split: - split_size = c // 8 - x_slices = torch.split(x, split_size, dim=1) - x = [ - nn.functional.interpolate(x, scale_factor=self.scale_factor) - for x in x_slices - ] - x = torch.cat(x, dim=1) - else: - x = nn.functional.interpolate(x, scale_factor=self.scale_factor) - - x = self.conv3d(x, is_init) - return x - - -class VideoDecoder(nn.Module): - - def __init__( - self, - ch=128, - z_channels=16, - out_channels=3, - ch_mult=(1, 2, 4, 4), - num_res_blocks=2, - temporal_up_layers=(2, 3), - temporal_downsample=4, - resamp_with_conv=True, - version=1, - ) -> None: - super().__init__() - - temb_ch = 0 - - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.temporal_downsample = temporal_downsample - - block_in = ch * ch_mult[self.num_resolutions - 1] - self.version = version - if version == 2: - channels = 4 * z_channels * 2**3 - self.conv_in = CausalConv(z_channels, channels, kernel_size=3) - self.shortcut_in = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D( - z_channels, channels, 1 - ) - self.conv_unpatchify = ConvPixelShuffleUpSampleLayer3D( - channels, block_in, kernel_size=3, factor=2 - ) - self.shortcut_unpathify = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D( - channels, block_in, 2 - ) - else: - self.conv_in = CausalConv(z_channels, block_in, kernel_size=3) - - # middle - self.mid = nn.Module() - self.mid.block_1 = Resnet3DBlock( - in_channels=block_in, out_channels=block_in, temb_channels=temb_ch - ) - self.mid.attn_1 = AttnBlock(block_in) - self.mid.block_2 = Resnet3DBlock( - in_channels=block_in, out_channels=block_in, temb_channels=temb_ch - ) - - # upsampling - self.up_id = len(temporal_up_layers) - self.video_frame_num = 1 - self.cur_video_frame_num = self.video_frame_num // 2**self.up_id + 1 - self.up = nn.ModuleList() - for i_level in reversed(range(self.num_resolutions)): - block = nn.ModuleList() - attn = nn.ModuleList() - block_out = ch * ch_mult[i_level] - for i_block in range(self.num_res_blocks + 1): - block.append( - Resnet3DBlock( - in_channels=block_in, - out_channels=block_out, - temb_channels=temb_ch, - ) - ) - block_in = block_out - up = nn.Module() - up.block = block - up.attn = attn - if i_level != 0: - if i_level in temporal_up_layers: - up.upsample = Upsample3D(block_in) - self.cur_video_frame_num = self.cur_video_frame_num * 2 - else: - up.upsample = Upsample2D(block_in, resamp_with_conv) - self.up.insert(0, up) # prepend to get consistent order - - # end - self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in) - self.conv_out = CausalConvAfterNorm(block_in, out_channels, kernel_size=3) - - @torch.inference_mode() - def forward(self, z, is_init=True) -> torch.Tensor: - z = rearrange(z, "b t c h w -> b c t h w") - h = self.conv_in(z, is_init=is_init) - if self.version == 2: - shortcut = self.shortcut_in(z, is_init=is_init) - h = h.add_(shortcut) - shortcut = self.shortcut_unpathify(h, is_init=is_init) - h = self.conv_unpatchify(h, is_init=is_init) - h = h.add_(shortcut) - - temb = None - - h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) - h = self.mid.block_1(h, temb, is_init=is_init) - h = self.mid.attn_1(h) - h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) - h = self.mid.block_2(h, temb, is_init=is_init) - - # upsampling - for i_level in reversed(range(self.num_resolutions)): - for i_block in range(self.num_res_blocks + 1): - h = h.permute(0, 2, 3, 4, 1).contiguous().permute(0, 4, 1, 2, 3) - h = self.up[i_level].block[i_block](h, temb, is_init=is_init) - if len(self.up[i_level].attn) > 0: - h = self.up[i_level].attn[i_block](h) - if i_level != 0: - if isinstance(self.up[i_level].upsample, Upsample2D): - B = h.size(0) - h = h.permute(0, 2, 3, 4, 1).flatten(0, 1) - h = self.up[i_level].upsample(h) - h = h.unflatten(0, (B, -1)).permute(0, 4, 1, 2, 3) - else: - h = self.up[i_level].upsample(h, is_init=is_init) - - # end - h = h.permute(0, 2, 3, 4, 1) # b c l h w -> b l h w c - h = base_group_norm_with_zero_pad(h, self.norm_out, act_silu=True, pad_size=2) - h = self.conv_out(h) - h = h.permute(0, 4, 1, 2, 3) - - if is_init: - h = h[:, :, (self.temporal_downsample - 1) :] - return h - - -def rms_norm(input, normalized_shape, eps=1e-6) -> torch.Tensor: - dtype = input.dtype - input = input.to(torch.float32) - variance = ( - input.pow(2) - .flatten(-len(normalized_shape)) - .mean(-1)[(...,) + (None,) * len(normalized_shape)] - ) - input = input * torch.rsqrt(variance + eps) - return input.to(dtype) - - -class DiagonalGaussianDistribution: - - def __init__( - self, - parameters, - deterministic=False, - rms_norm_mean=False, - only_return_mean=False, - ) -> None: - self.parameters = parameters - self.mean, self.logvar = torch.chunk(parameters, 2, dim=-3) # N,[X],C,H,W - self.logvar = torch.clamp(self.logvar, -30.0, 20.0) - self.std = torch.exp(0.5 * self.logvar) - self.var = torch.exp(self.logvar) - self.deterministic = deterministic - if self.deterministic: - self.var = self.std = torch.zeros_like( - self.mean, device=self.parameters.device, dtype=self.parameters.dtype - ) - if rms_norm_mean: - self.mean = rms_norm(self.mean, self.mean.size()[1:]) - self.only_return_mean = only_return_mean - - def sample(self, generator=None) -> torch.Tensor: - # make sure sample is on the same device - # as the parameters and has same dtype - sample = torch.randn( - self.mean.shape, generator=generator, device=self.parameters.device - ) - sample = sample.to(dtype=self.parameters.dtype) - x = self.mean + self.std * sample - if self.only_return_mean: - return self.mean - else: - return x - - -class AutoencoderKLStepvideo(nn.Module, ParallelTiledVAE): - - def __init__( - self, - config: StepVideoVAEConfig, - ) -> None: - nn.Module.__init__(self) - ParallelTiledVAE.__init__(self, config) - - self.frame_len = config.frame_len - - if config.version == 2: - self.latent_len = 3 - base_group_norm.spatial = True # type: ignore[attr-defined] - else: - self.latent_len = 5 - base_group_norm.spatial = False # type: ignore[attr-defined] - - self.encoder = VideoEncoder( - in_channels=config.in_channels, - z_channels=config.z_channels, - num_res_blocks=config.num_res_blocks, - version=config.version, - ) - - self.decoder = VideoDecoder( - z_channels=config.z_channels, - out_channels=config.out_channels, - num_res_blocks=config.num_res_blocks, - version=config.version, - ) - - self.world_size = config.world_size - # self.is_init = True - - def load_state_dict(self, state_dict, strict=True): - remapped = {} - for key, value in state_dict.items(): - if key.startswith("decoder.conv_out."): - # move “decoder.conv_out.weight” → “decoder.conv_out.conv.weight” - suffix = key[len("decoder.conv_out.") :] - remapped[f"decoder.conv_out.conv.{suffix}"] = value - else: - remapped[key] = value - super().load_state_dict(remapped, strict=strict) - - def _encode(self, x, is_init_image=True) -> torch.Tensor: - # b, len, c, h, w = x.size() - b, c, len, h, w = x.size() - # x = rearrange(x, 'b l c h w -> b c l h w').contiguous() - z = self.encoder(x, len, True) # 下采样[1, 4, 8, 16, 16] - return z - - @torch.inference_mode() - def encode(self, x): - # b (nc cf) c h w -> (b nc) cf c h w -> encode -> (b nc) cf c h w -> b (nc cf) c h w - chunks = list(x.split(self.frame_len, dim=1)) - for i in range(len(chunks)): - chunks[i] = self._encode(chunks[i], True) - z = torch.cat(chunks, dim=1) - - posterior = DiagonalGaussianDistribution(z) - return posterior.sample() - - def _decode(self, z) -> torch.Tensor: - - chunks = list(z.split(self.latent_len, dim=2)) - for i in range(len(chunks)): - chunks[i] = chunks[i].permute(0, 2, 1, 3, 4) - chunks[i] = chunks[i].to(next(self.decoder.parameters()).dtype) - chunks[i] = self.decoder(chunks[i], is_init=True) - x = torch.cat(chunks, dim=2) - return x - - def decode(self, z) -> torch.Tensor: - num_frames = z.size(2) - dec = ParallelTiledVAE.decode(self, z).permute(0, 2, 1, 3, 4) - dec = self.mix(dec).permute(0, 2, 1, 3, 4) - num_sample_frames = num_frames // 3 * 17 - return dec[:, :, :num_sample_frames] - - def mix(self, x) -> torch.Tensor: - remain_scale = 0.6 - mix_scale = 1.0 - remain_scale - front = slice(self.frame_len - 1, x.size(1) - 1, self.frame_len) - back = slice(self.frame_len, x.size(1), self.frame_len) - x[:, back] = x[:, back] * remain_scale + x[:, front] * mix_scale - x[:, front] = x[:, front] * remain_scale + x[:, back] * mix_scale - return x - - def forward( - self, - sample: torch.Tensor, - sample_posterior: bool = False, - generator: torch.Generator | None = None, - ) -> torch.Tensor: - """ - Args: - sample (`torch.Tensor`): Input sample. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`DecoderOutput`] instead of a plain tuple. - """ - x = sample - posterior = self.encode(x).latent_dist - if sample_posterior: - z = posterior.sample(generator=generator) - else: - z = posterior.mode() - dec = self.decode(z) - return dec - - -EntryClass = AutoencoderKLStepvideo diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stepvideo_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/stepvideo_pipeline.py deleted file mode 100644 index 9d2e80c23..000000000 --- a/python/sglang/multimodal_gen/runtime/pipelines/stepvideo_pipeline.py +++ /dev/null @@ -1,182 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -# SPDX-License-Identifier: Apache-2.0 -# type: ignore -# SPDX-License-Identifier: Apache-2.0 -""" -Hunyuan video diffusion pipeline implementation. - -This module contains an implementation of the Hunyuan video diffusion pipeline -using the modular pipeline architecture. -""" - -import os -from typing import Any - -import torch -from huggingface_hub import hf_hub_download - -from sglang.multimodal_gen.runtime.distributed import get_local_torch_device -from sglang.multimodal_gen.runtime.loader.component_loader import ( - PipelineComponentLoader, -) -from sglang.multimodal_gen.runtime.models.encoders.bert import ( - HunyuanClip, # type: ignore -) -from sglang.multimodal_gen.runtime.models.encoders.stepllm import STEP1TextEncoder -from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( - ComposedPipelineBase, -) -from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline -from sglang.multimodal_gen.runtime.pipelines_core.stages import ( - DecodingStage, - DenoisingStage, - InputValidationStage, - LatentPreparationStage, - StepvideoPromptEncodingStage, - TimestepPreparationStage, -) -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 StepVideoPipeline(LoRAPipeline, ComposedPipelineBase): - pipeline_name = "StepVideoPipeline" - - _required_config_modules = ["transformer", "scheduler", "vae"] - - def create_pipeline_stages(self, server_args: ServerArgs): - """Set up pipeline stages with proper dependency injection.""" - - self.add_stage( - stage_name="input_validation_stage", stage=InputValidationStage() - ) - - self.add_stage( - stage_name="prompt_encoding_stage", - stage=StepvideoPromptEncodingStage( - stepllm=self.get_module("text_encoder"), - clip=self.get_module("text_encoder_2"), - ), - ) - - self.add_stage( - stage_name="timestep_preparation_stage", - stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")), - ) - - self.add_stage( - stage_name="latent_preparation_stage", - stage=LatentPreparationStage( - scheduler=self.get_module("scheduler"), - transformer=self.get_module("transformer"), - ), - ) - - self.add_stage( - stage_name="denoising_stage", - stage=DenoisingStage( - transformer=self.get_module("transformer"), - scheduler=self.get_module("scheduler"), - ), - ) - - self.add_stage( - stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")) - ) - - def build_llm(self, model_dir, device) -> torch.nn.Module: - text_encoder = ( - STEP1TextEncoder(model_dir, max_length=320).to(torch.bfloat16).eval() - ) - return text_encoder - - def build_clip(self, model_dir, device) -> HunyuanClip: - clip = HunyuanClip(model_dir, max_length=77).eval() - return clip - - def initialize_pipeline(self, server_args: ServerArgs): - """ - Initialize the pipeline. - """ - target_device = get_local_torch_device() - llm_dir = os.path.join(self.model_path, "step_llm") - clip_dir = os.path.join(self.model_path, "hunyuan_clip") - text_enc = self.build_llm(llm_dir, target_device) - clip_enc = self.build_clip(clip_dir, target_device) - self.add_module("text_encoder", text_enc) - self.add_module("text_encoder_2", clip_enc) - lib_path = ( - os.path.join( - server_args.model_path, - "lib/liboptimus_ths-torch2.5-cu124.cpython-310-x86_64-linux-gnu.so", - ) - if os.path.isdir(server_args.model_path) # local checkout - else hf_hub_download( - repo_id=server_args.model_path, - filename="lib/liboptimus_ths-torch2.5-cu124.cpython-310-x86_64-linux-gnu.so", - ) - ) - torch.ops.load_library(lib_path) - - def load_modules( - self, - server_args: ServerArgs, - loaded_modules: dict[str, torch.nn.Module] | None = None, - ) -> dict[str, Any]: - """ - Load the modules from the config. - """ - model_index = self._load_config() - logger.info("Loading pipeline modules from config: %s", model_index) - - # remove keys that are not pipeline modules - model_index.pop("_class_name") - model_index.pop("_diffusers_version") - - # some sanity checks - assert ( - len(model_index) > 1 - ), "model_index.json must contain at least one pipeline module" - - required_modules = ["transformer", "scheduler", "vae"] - for module_name in required_modules: - if module_name not in model_index: - raise ValueError( - f"model_index.json must contain a {module_name} module" - ) - logger.info("Diffusers config passed sanity checks") - - # all the component models used by the pipeline - modules = {} - for module_name, ( - transformers_or_diffusers, - architecture, - ) in model_index.items(): - component_model_path = os.path.join(self.model_path, module_name) - module = PipelineComponentLoader.load_module( - module_name=module_name, - component_model_path=component_model_path, - transformers_or_diffusers=transformers_or_diffusers, - server_args=server_args, - ) - logger.info("Loaded module %s from %s", module_name, component_model_path) - - if module_name in modules: - logger.warning("Overwriting module %s", module_name) - modules[module_name] = module - - required_modules = self.required_config_modules - # Check if all required modules were loaded - for module_name in required_modules: - if module_name not in modules or modules[module_name] is None: - raise ValueError( - f"Required module {module_name} was not loaded properly" - ) - - return modules - - -EntryClass = StepVideoPipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py index 3c962b0b9..f395f5525 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py @@ -31,9 +31,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import ( LatentPreparationStage, ) -from sglang.multimodal_gen.runtime.pipelines_core.stages.stepvideo_encoding import ( - StepvideoPromptEncodingStage, -) from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import ( TextEncodingStage, ) @@ -55,5 +52,4 @@ __all__ = [ "ImageEncodingStage", "ImageVAEEncodingStage", "TextEncodingStage", - "StepvideoPromptEncodingStage", ] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py index acd901b3b..7d38e0a7a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py @@ -128,8 +128,6 @@ class LatentPreparationStage(PipelineStage): server_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio ) latent_num_frames = (video_length - 1) // temporal_scale_factor + 1 - else: # stepvideo only - latent_num_frames = video_length // 17 * 3 return int(latent_num_frames) def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/stepvideo_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/stepvideo_encoding.py deleted file mode 100644 index eed0edd99..000000000 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/stepvideo_encoding.py +++ /dev/null @@ -1,99 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -# SPDX-License-Identifier: Apache-2.0 - -import torch - -from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context -from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req -from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage -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.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger - -logger = init_logger(__name__) - - -# The dedicated stepvideo prompt encoding stage. -class StepvideoPromptEncodingStage(PipelineStage): - """ - Stage for encoding prompts using the remote caption API. - - This stage applies the magic string transformations and calls - the remote caption service asynchronously to get: - - primary prompt embeddings, - - an attention mask, - - and a clip embedding. - """ - - def __init__(self, stepllm, clip) -> None: - super().__init__() - # self.caption_client = caption_client # This should have a call_caption(prompts: List[str]) method. - self.stepllm = stepllm - self.clip = clip - - @torch.no_grad() - def forward(self, batch: Req, server_args) -> Req: - - prompts = [batch.prompt + server_args.pipeline_config.pos_magic] - bs = len(prompts) - prompts += [server_args.pipeline_config.neg_magic] * bs - with set_forward_context(current_timestep=0, attn_metadata=None): - y, y_mask = self.stepllm(prompts) - clip_emb, _ = self.clip(prompts) - len_clip = clip_emb.shape[1] - y_mask = torch.nn.functional.pad(y_mask, (len_clip, 0), value=1) - pos_clip, neg_clip = clip_emb[:bs], clip_emb[bs:] - - # split positive vs negative text - batch.prompt_embeds = y[:bs] # [bs, seq_len, dim] - batch.negative_prompt_embeds = y[bs : 2 * bs] # [bs, seq_len, dim] - batch.prompt_attention_mask = y_mask[:bs] # [bs, seq_len] - batch.negative_attention_mask = y_mask[bs : 2 * bs] # [bs, seq_len] - batch.clip_embedding_pos = pos_clip - batch.clip_embedding_neg = neg_clip - return batch - - def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: - """Verify stepvideo encoding stage inputs.""" - result = VerificationResult() - result.add_check("prompt", batch.prompt, V.string_not_empty) - return result - - def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: - """Verify stepvideo encoding stage outputs.""" - result = VerificationResult() - result.add_check( - "prompt_embeds", batch.prompt_embeds, [V.is_tensor, V.with_dims(3)] - ) - result.add_check( - "negative_prompt_embeds", - batch.negative_prompt_embeds, - [V.is_tensor, V.with_dims(3)], - ) - result.add_check( - "prompt_attention_mask", - batch.prompt_attention_mask, - [V.is_tensor, V.with_dims(2)], - ) - result.add_check( - "negative_attention_mask", - batch.negative_attention_mask, - [V.is_tensor, V.with_dims(2)], - ) - result.add_check( - "clip_embedding_pos", - batch.clip_embedding_pos, - [V.is_tensor, V.with_dims(2)], - ) - result.add_check( - "clip_embedding_neg", - batch.clip_embedding_neg, - [V.is_tensor, V.with_dims(2)], - ) - return result diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 9b5b24e6e..4f9f92605 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -342,11 +342,6 @@ class ServerArgs: type=str, help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.", ) - parser.add_argument( - "--model-dir", - type=str, - help="Directory containing StepVideo model", - ) parser.add_argument( "--vae-path", type=str,