[diffusion] chore: remove stepvideo code (#15918)
This commit is contained in:
Yuhao Yang
@@ -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"]
|
||||
|
||||
@@ -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"
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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"
|
||||
@@ -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)",
|
||||
|
||||
@@ -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 = "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。"
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user