WIP: initial multimodal-gen support (#12484)

Co-authored-by: yhyang201 <yhyang201@gmail.com>
Co-authored-by: yizhang2077 <1109276519@qq.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
Co-authored-by: JiLi <leege233@gmail.com>
Co-authored-by: CHEN Xi <78632976+RubiaCx@users.noreply.github.com>
Co-authored-by: laixin <xielx@shanghaitech.edu.cn>
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
Co-authored-by: jzhang38 <a1286225768@gmail.com>
Co-authored-by: BrianChen1129 <yongqichcd@gmail.com>
Co-authored-by: Kevin Lin <42618777+kevin314@users.noreply.github.com>
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
Co-authored-by: rlsu9 <r3su@ucsd.edu>
Co-authored-by: Jinzhe Pan <48981407+eigensystem@users.noreply.github.com>
Co-authored-by: foreverpiano <pianoqwz@qq.com>
Co-authored-by: RandNMR73 <notomatthew31@gmail.com>
Co-authored-by: PorridgeSwim <yz3883@columbia.edu>
Co-authored-by: Jiali Chen <90408393+gary-chenjl@users.noreply.github.com>
This commit is contained in:
Mick
2025-11-05 12:28:52 -08:00
committed by GitHub
co-authored by yhyang201 yizhang2077 Xinyuan Tong ispobock JiLi CHEN Xi laixin SolitaryThinker jzhang38 BrianChen1129 Kevin Lin Edenzzzz rlsu9 Jinzhe Pan foreverpiano RandNMR73 PorridgeSwim Jiali Chen
parent 4fe53e5888
commit 7bc1dae095
249 changed files with 63750 additions and 11 deletions
@@ -0,0 +1,3 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# Configs for pipelines, and pipeline modules (in models folder)
@@ -0,0 +1,16 @@
{
"temporal_chunk_size": 2,
"temporal_topk": 2,
"spatial_chunk_size": [4, 13],
"spatial_topk": 6,
"st_chunk_size": [4, 4, 13],
"st_topk": 18,
"moba_select_mode": "topk",
"moba_threshold": 0.25,
"moba_threshold_type": "query_head",
"first_full_layer": 0,
"first_full_step": 12,
"temporal_layer": 1,
"spatial_layer": 1,
"st_layer": 1
}
@@ -0,0 +1,16 @@
{
"temporal_chunk_size": 2,
"temporal_topk": 3,
"spatial_chunk_size": [3, 4],
"spatial_topk": 20,
"st_chunk_size": [4, 6, 4],
"st_topk": 15,
"moba_select_mode": "threshold",
"moba_threshold": 0.25,
"moba_threshold_type": "query_head",
"first_full_layer": 0,
"first_full_step": 12,
"temporal_layer": 1,
"spatial_layer": 1,
"st_layer": 1
}
@@ -0,0 +1,258 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import dataclasses
from enum import Enum
from typing import Any, Optional
from sglang.multimodal_gen.configs.utils import update_config_from_args
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.utils import FlexibleArgumentParser, StoreBoolean
logger = init_logger(__name__)
class DatasetType(str, Enum):
"""
Enumeration for different dataset types.
"""
HF = "hf"
MERGED = "merged"
@classmethod
def from_string(cls, value: str) -> "DatasetType":
"""Convert string to DatasetType enum."""
try:
return cls(value.lower())
except ValueError:
raise ValueError(
f"Invalid dataset type: {value}. Must be one of: {', '.join([m.value for m in cls])}"
) from None
@classmethod
def choices(cls) -> list[str]:
"""Get all available choices as strings for argparse."""
return [dataset_type.value for dataset_type in cls]
class VideoLoaderType(str, Enum):
"""
Enumeration for different video loaders.
"""
TORCHCODEC = "torchcodec"
TORCHVISION = "torchvision"
@classmethod
def from_string(cls, value: str) -> "VideoLoaderType":
"""Convert string to VideoLoader enum."""
try:
return cls(value.lower())
except ValueError:
raise ValueError(
f"Invalid video loader: {value}. Must be one of: {', '.join([m.value for m in cls])}"
) from None
@classmethod
def choices(cls) -> list[str]:
"""Get all available choices as strings for argparse."""
return [video_loader.value for video_loader in cls]
@dataclasses.dataclass
class PreprocessConfig:
"""Configuration for preprocessing operations."""
# Model and dataset configuration
model_path: str = ""
dataset_path: str = ""
dataset_type: DatasetType = DatasetType.HF
dataset_output_dir: str = "./output"
# Dataloader configuration
dataloader_num_workers: int = 1
preprocess_video_batch_size: int = 2
# Saver configuration
samples_per_file: int = 64
flush_frequency: int = 256
# Video processing parameters
video_loader_type: VideoLoaderType = VideoLoaderType.TORCHCODEC
max_height: int = 480
max_width: int = 848
num_frames: int = 163
video_length_tolerance_range: float = 2.0
train_fps: int = 30
speed_factor: float = 1.0
drop_short_ratio: float = 1.0
do_temporal_sample: bool = False
# Model configuration
training_cfg_rate: float = 0.0
# framework configuration
seed: int = 42
@staticmethod
def add_cli_args(
parser: FlexibleArgumentParser, prefix: str = "preprocess"
) -> FlexibleArgumentParser:
"""Add preprocessing configuration arguments to the parser."""
prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else ""
preprocess_args = parser.add_argument_group("Preprocessing Arguments")
# Model & Dataset
preprocess_args.add_argument(
f"--{prefix_with_dot}model-path",
type=str,
default=PreprocessConfig.model_path,
help="Path to the model for preprocessing",
)
preprocess_args.add_argument(
f"--{prefix_with_dot}dataset-path",
type=str,
default=PreprocessConfig.dataset_path,
help="Path to the dataset directory for preprocessing",
)
preprocess_args.add_argument(
f"--{prefix_with_dot}dataset-type",
type=str,
choices=DatasetType.choices(),
default=PreprocessConfig.dataset_type.value,
help="Type of the dataset",
)
preprocess_args.add_argument(
f"--{prefix_with_dot}dataset-output-dir",
type=str,
default=PreprocessConfig.dataset_output_dir,
help="The output directory where the dataset will be written.",
)
# Dataloader
preprocess_args.add_argument(
f"--{prefix_with_dot}dataloader-num-workers",
type=int,
default=PreprocessConfig.dataloader_num_workers,
help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
)
preprocess_args.add_argument(
f"--{prefix_with_dot}preprocess-video-batch-size",
type=int,
default=PreprocessConfig.preprocess_video_batch_size,
help="Batch size (per device) for the training dataloader.",
)
# Saver
preprocess_args.add_argument(
f"--{prefix_with_dot}samples-per-file",
type=int,
default=PreprocessConfig.samples_per_file,
help="Number of samples per output file",
)
preprocess_args.add_argument(
f"--{prefix_with_dot}flush-frequency",
type=int,
default=PreprocessConfig.flush_frequency,
help="How often to save to parquet files",
)
# Video processing parameters
preprocess_args.add_argument(
f"--{prefix_with_dot}video-loader-type",
type=str,
choices=VideoLoaderType.choices(),
default=PreprocessConfig.video_loader_type.value,
help="Type of the video loader",
)
preprocess_args.add_argument(
f"--{prefix_with_dot}max-height",
type=int,
default=PreprocessConfig.max_height,
help="Maximum height for video processing",
)
preprocess_args.add_argument(
f"--{prefix_with_dot}max-width",
type=int,
default=PreprocessConfig.max_width,
help="Maximum width for video processing",
)
preprocess_args.add_argument(
f"--{prefix_with_dot}num-frames",
type=int,
default=PreprocessConfig.num_frames,
help="Number of frames to process",
)
preprocess_args.add_argument(
f"--{prefix_with_dot}video-length-tolerance-range",
type=float,
default=PreprocessConfig.video_length_tolerance_range,
help="Video length tolerance range",
)
preprocess_args.add_argument(
f"--{prefix_with_dot}train-fps",
type=int,
default=PreprocessConfig.train_fps,
help="Training FPS",
)
preprocess_args.add_argument(
f"--{prefix_with_dot}speed-factor",
type=float,
default=PreprocessConfig.speed_factor,
help="Speed factor for video processing",
)
preprocess_args.add_argument(
f"--{prefix_with_dot}drop-short-ratio",
type=float,
default=PreprocessConfig.drop_short_ratio,
help="Ratio for dropping short videos",
)
preprocess_args.add_argument(
f"--{prefix_with_dot}do-temporal-sample",
action=StoreBoolean,
default=PreprocessConfig.do_temporal_sample,
help="Whether to do temporal sampling",
)
# Model Training configuration
preprocess_args.add_argument(
f"--{prefix_with_dot}training-cfg-rate",
type=float,
default=PreprocessConfig.training_cfg_rate,
help="Training CFG rate",
)
preprocess_args.add_argument(
f"--{prefix_with_dot}seed",
type=int,
default=PreprocessConfig.seed,
help="Seed for random number generator",
)
return parser
@classmethod
def from_kwargs(cls, kwargs: dict[str, Any]) -> Optional["PreprocessConfig"]:
"""Create PreprocessConfig from keyword arguments."""
if "dataset_type" in kwargs and isinstance(kwargs["dataset_type"], str):
kwargs["dataset_type"] = DatasetType.from_string(kwargs["dataset_type"])
if "video_loader_type" in kwargs and isinstance(
kwargs["video_loader_type"], str
):
kwargs["video_loader_type"] = VideoLoaderType.from_string(
kwargs["video_loader_type"]
)
preprocess_config = cls()
if not update_config_from_args(
preprocess_config, kwargs, prefix="preprocess", pop_args=True
):
return None
return preprocess_config
def check_preprocess_config(self) -> None:
if self.dataset_path == "":
raise ValueError("dataset_path must be set for preprocess mode")
if self.samples_per_file <= 0:
raise ValueError("samples_per_file must be greater than 0")
if self.flush_frequency <= 0:
raise ValueError("flush_frequency must be greater than 0")
@@ -0,0 +1,48 @@
{
"embedded_cfg_scale": 6,
"flow_shift": 17,
"dit_cpu_offload": false,
"disable_autocast": false,
"precision": "bf16",
"vae_precision": "fp32",
"vae_tiling": true,
"vae_sp": true,
"vae_config": {
"load_encoder": false,
"load_decoder": true,
"tile_sample_min_height": 256,
"tile_sample_min_width": 256,
"tile_sample_min_num_frames": 16,
"tile_sample_stride_height": 192,
"tile_sample_stride_width": 192,
"tile_sample_stride_num_frames": 12,
"blend_num_frames": 4,
"use_tiling": true,
"use_temporal_tiling": true,
"use_parallel_tiling": true
},
"dit_config": {
"prefix": "Hunyuan",
"quant_config": null
},
"text_encoder_precisions": [
"fp16",
"fp16"
],
"text_encoder_configs": [
{
"prefix": "llama",
"quant_config": null,
"lora_config": null
},
{
"prefix": "clip",
"quant_config": null,
"lora_config": null,
"num_hidden_layers_override": null,
"require_post_norm": null
}
],
"mask_strategy_file_path": null,
"enable_torch_compile": false
}
@@ -0,0 +1,8 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
from sglang.multimodal_gen.configs.models.base import ModelConfig
from sglang.multimodal_gen.configs.models.dits.base import DiTConfig
from sglang.multimodal_gen.configs.models.encoders.base import EncoderConfig
from sglang.multimodal_gen.configs.models.vaes.base import VAEConfig
__all__ = ["ModelConfig", "VAEConfig", "DiTConfig", "EncoderConfig"]
@@ -0,0 +1,105 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field, fields
from typing import Any, Dict
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
# 1. ArchConfig contains all fields from diffuser's/transformer's config.json (i.e. all fields related to the architecture of the model)
# 2. ArchConfig should be inherited & overridden by each model arch_config
# 3. Any field in ArchConfig is fixed upon initialization, and should be hidden away from users
@dataclass
class ArchConfig:
stacked_params_mapping: list[tuple[str, str, str]] = field(
default_factory=list
) # mapping from huggingface weight names to custom names
extra_attrs: Dict[str, Any] = field(default_factory=dict)
def __getattr__(self, name: str):
d = object.__getattribute__(self, "__dict__")
extras = d.get("extra_attrs")
if extras is not None and name in extras:
return extras[name]
raise AttributeError(
f"'{self.__class__.__name__}' object has no attribute '{name}'"
)
def __setattr__(self, key, value):
if key in type(self).__dataclass_fields__:
object.__setattr__(self, key, value)
else:
d = object.__getattribute__(self, "__dict__")
extras = d.get("extra_attrs")
if extras is None:
extras = {}
d["extra_attrs"] = extras
extras[key] = value
@dataclass
class ModelConfig:
# Every model config parameter can be categorized into either ArchConfig or everything else
# Diffuser/Transformer parameters
arch_config: ArchConfig = field(default_factory=ArchConfig)
# sgl-diffusion-specific parameters here
# i.e. STA, quantization, teacache
def __getattr__(self, name):
# Only called if 'name' is not found in ModelConfig directly
if hasattr(self.arch_config, name):
return getattr(self.arch_config, name)
raise AttributeError(
f"'{type(self).__name__}' object has no attribute '{name}'"
)
def __getstate__(self):
# Return a dictionary of attributes to pickle
# Convert to dict and exclude any problematic attributes
state = self.__dict__.copy()
return state
def __setstate__(self, state):
# Restore instance attributes from the unpickled state
self.__dict__.update(state)
# This should be used only when loading from transformers/diffusers
def update_model_arch(self, source_model_dict: dict[str, Any]) -> None:
"""
Update arch_config with source_model_dict
"""
arch_config = self.arch_config
valid_fields = {f.name for f in fields(arch_config)}
for key, value in source_model_dict.items():
setattr(arch_config, key, value)
# else:
# raise AttributeError(
# f"{type(arch_config).__name__} has no field '{key}'"
# )
if hasattr(arch_config, "__post_init__"):
arch_config.__post_init__()
def update_model_config(self, source_model_dict: dict[str, Any]) -> None:
assert (
"arch_config" not in source_model_dict
), "Source model config shouldn't contain arch_config."
valid_fields = {f.name for f in fields(self)}
for key, value in source_model_dict.items():
if key in valid_fields:
setattr(self, key, value)
else:
logger.warning(
"%s does not contain field '%s'!", type(self).__name__, key
)
raise AttributeError(f"Invalid field: {key}")
if hasattr(self, "__post_init__"):
self.__post_init__()
@@ -0,0 +1,7 @@
# 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"]
@@ -0,0 +1,69 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from typing import Any
from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig
from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
@dataclass
class DiTArchConfig(ArchConfig):
_fsdp_shard_conditions: list = field(default_factory=list)
_compile_conditions: list = field(default_factory=list)
param_names_mapping: dict = field(default_factory=dict)
reverse_param_names_mapping: dict = field(default_factory=dict)
lora_param_names_mapping: dict = field(default_factory=dict)
_supported_attention_backends: set[AttentionBackendEnum] = field(
default_factory=lambda: {
AttentionBackendEnum.SLIDING_TILE_ATTN,
AttentionBackendEnum.SAGE_ATTN,
AttentionBackendEnum.FA3,
AttentionBackendEnum.TORCH_SDPA,
AttentionBackendEnum.VIDEO_SPARSE_ATTN,
AttentionBackendEnum.VMOBA_ATTN,
AttentionBackendEnum.SAGE_ATTN_THREE,
}
)
hidden_size: int = 0
num_attention_heads: int = 0
num_channels_latents: int = 0
exclude_lora_layers: list[str] = field(default_factory=list)
boundary_ratio: float | None = None
def __post_init__(self) -> None:
if not self._compile_conditions:
self._compile_conditions = self._fsdp_shard_conditions.copy()
@dataclass
class DiTConfig(ModelConfig):
arch_config: DiTArchConfig = field(default_factory=DiTArchConfig)
# sgl-diffusionDiT-specific parameters
prefix: str = ""
quant_config: QuantizationConfig | None = None
@staticmethod
def add_cli_args(parser: Any, prefix: str = "dit-config") -> Any:
"""Add CLI arguments for DiTConfig fields"""
parser.add_argument(
f"--{prefix}.prefix",
type=str,
dest=f"{prefix.replace('-', '_')}.prefix",
default=DiTConfig.prefix,
help="Prefix for the DiT model",
)
parser.add_argument(
f"--{prefix}.quant-config",
type=str,
dest=f"{prefix.replace('-', '_')}.quant_config",
default=None,
help="Quantization configuration for the DiT model",
)
return parser
@@ -0,0 +1,36 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from typing import Tuple
from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
@dataclass
class FluxArchConfig(DiTArchConfig):
patch_size: int = 1
in_channels: int = 64
out_channels: int | None = None
num_layers: int = 19
num_single_layers: int = 38
attention_head_dim: int = 128
num_attention_heads: int = 24
joint_attention_dim: int = 4096
pooled_projection_dim: int = 768
guidance_embeds: bool = False
axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)
def __post_init__(self):
super().__post_init__()
self.out_channels = self.out_channels or self.in_channels
self.hidden_size = self.num_attention_heads * self.attention_head_dim
self.num_channels_latents = self.out_channels
@dataclass
class FluxConfig(DiTConfig):
arch_config: DiTArchConfig = field(default_factory=FluxArchConfig)
prefix: str = "Flux"
@@ -0,0 +1,185 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
import torch
from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
def is_double_block(n: str, m) -> bool:
return "double" in n and str.isdigit(n.split(".")[-1])
def is_single_block(n: str, m) -> bool:
return "single" in n and str.isdigit(n.split(".")[-1])
def is_refiner_block(n: str, m) -> bool:
return "refiner" in n and str.isdigit(n.split(".")[-1])
def is_txt_in(n: str, m) -> bool:
return n.split(".")[-1] == "txt_in"
@dataclass
class HunyuanVideoArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(
default_factory=lambda: [is_double_block, is_single_block, is_refiner_block]
)
_compile_conditions: list = field(
default_factory=lambda: [is_double_block, is_single_block, is_txt_in]
)
param_names_mapping: dict = field(
default_factory=lambda: {
# 1. context_embedder.time_text_embed submodules (specific rules, applied first):
r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$": r"txt_in.t_embedder.mlp.fc_in.\1",
r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_2\.(.*)$": r"txt_in.t_embedder.mlp.fc_out.\1",
r"^context_embedder\.proj_in\.(.*)$": r"txt_in.input_embedder.\1",
r"^context_embedder\.time_text_embed\.text_embedder\.linear_1\.(.*)$": r"txt_in.c_embedder.fc_in.\1",
r"^context_embedder\.time_text_embed\.text_embedder\.linear_2\.(.*)$": r"txt_in.c_embedder.fc_out.\1",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm1\.(.*)$": r"txt_in.refiner_blocks.\1.norm1.\2",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm2\.(.*)$": r"txt_in.refiner_blocks.\1.norm2.\2",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_q\.(.*)$": (
r"txt_in.refiner_blocks.\1.self_attn_qkv.\2",
0,
3,
),
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_k\.(.*)$": (
r"txt_in.refiner_blocks.\1.self_attn_qkv.\2",
1,
3,
),
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_v\.(.*)$": (
r"txt_in.refiner_blocks.\1.self_attn_qkv.\2",
2,
3,
),
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"txt_in.refiner_blocks.\1.self_attn_proj.\2",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$": r"txt_in.refiner_blocks.\1.mlp.fc_in.\2",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$": r"txt_in.refiner_blocks.\1.mlp.fc_out.\2",
r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm_out\.linear\.(.*)$": r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2",
# 3. x_embedder mapping:
r"^x_embedder\.proj\.(.*)$": r"img_in.proj.\1",
# 4. Top-level time_text_embed mappings:
r"^time_text_embed\.timestep_embedder\.linear_1\.(.*)$": r"time_in.mlp.fc_in.\1",
r"^time_text_embed\.timestep_embedder\.linear_2\.(.*)$": r"time_in.mlp.fc_out.\1",
r"^time_text_embed\.guidance_embedder\.linear_1\.(.*)$": r"guidance_in.mlp.fc_in.\1",
r"^time_text_embed\.guidance_embedder\.linear_2\.(.*)$": r"guidance_in.mlp.fc_out.\1",
r"^time_text_embed\.text_embedder\.linear_1\.(.*)$": r"vector_in.fc_in.\1",
r"^time_text_embed\.text_embedder\.linear_2\.(.*)$": r"vector_in.fc_out.\1",
# 5. transformer_blocks mapping:
r"^transformer_blocks\.(\d+)\.norm1\.linear\.(.*)$": r"double_blocks.\1.img_mod.linear.\2",
r"^transformer_blocks\.(\d+)\.norm1_context\.linear\.(.*)$": r"double_blocks.\1.txt_mod.linear.\2",
r"^transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$": r"double_blocks.\1.img_attn_q_norm.\2",
r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$": r"double_blocks.\1.img_attn_k_norm.\2",
r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": (
r"double_blocks.\1.img_attn_qkv.\2",
0,
3,
),
r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": (
r"double_blocks.\1.img_attn_qkv.\2",
1,
3,
),
r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": (
r"double_blocks.\1.img_attn_qkv.\2",
2,
3,
),
r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$": (
r"double_blocks.\1.txt_attn_qkv.\2",
0,
3,
),
r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$": (
r"double_blocks.\1.txt_attn_qkv.\2",
1,
3,
),
r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$": (
r"double_blocks.\1.txt_attn_qkv.\2",
2,
3,
),
r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"double_blocks.\1.img_attn_proj.\2",
# Corrected: merge attn.to_add_out into the main projection.
r"^transformer_blocks\.(\d+)\.attn\.to_add_out\.(.*)$": r"double_blocks.\1.txt_attn_proj.\2",
r"^transformer_blocks\.(\d+)\.attn\.norm_added_q\.(.*)$": r"double_blocks.\1.txt_attn_q_norm.\2",
r"^transformer_blocks\.(\d+)\.attn\.norm_added_k\.(.*)$": r"double_blocks.\1.txt_attn_k_norm.\2",
r"^transformer_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$": r"double_blocks.\1.img_mlp.fc_in.\2",
r"^transformer_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$": r"double_blocks.\1.img_mlp.fc_out.\2",
r"^transformer_blocks\.(\d+)\.ff_context\.net\.0(?:\.proj)?\.(.*)$": r"double_blocks.\1.txt_mlp.fc_in.\2",
r"^transformer_blocks\.(\d+)\.ff_context\.net\.2(?:\.proj)?\.(.*)$": r"double_blocks.\1.txt_mlp.fc_out.\2",
# 6. single_transformer_blocks mapping:
r"^single_transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$": r"single_blocks.\1.q_norm.\2",
r"^single_transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$": r"single_blocks.\1.k_norm.\2",
r"^single_transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": (
r"single_blocks.\1.linear1.\2",
0,
4,
),
r"^single_transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": (
r"single_blocks.\1.linear1.\2",
1,
4,
),
r"^single_transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": (
r"single_blocks.\1.linear1.\2",
2,
4,
),
r"^single_transformer_blocks\.(\d+)\.proj_mlp\.(.*)$": (
r"single_blocks.\1.linear1.\2",
3,
4,
),
# Corrected: map proj_out to modulation.linear rather than a separate proj_out branch.
r"^single_transformer_blocks\.(\d+)\.proj_out\.(.*)$": r"single_blocks.\1.linear2.\2",
r"^single_transformer_blocks\.(\d+)\.norm\.linear\.(.*)$": r"single_blocks.\1.modulation.linear.\2",
# 7. Final layers mapping:
r"^norm_out\.linear\.(.*)$": r"final_layer.adaLN_modulation.linear.\1",
r"^proj_out\.(.*)$": r"final_layer.linear.\1",
}
)
# Reverse mapping for saving checkpoints: custom -> hf
reverse_param_names_mapping: dict = field(default_factory=lambda: {})
patch_size: int = 2
patch_size_t: int = 1
in_channels: int = 16
out_channels: int = 16
num_attention_heads: int = 24
attention_head_dim: int = 128
mlp_ratio: float = 4.0
num_layers: int = 20
num_single_layers: int = 40
num_refiner_layers: int = 2
rope_axes_dim: tuple[int, int, int] = (16, 56, 56)
guidance_embeds: bool = False
dtype: torch.dtype | None = None
text_embed_dim: int = 4096
pooled_projection_dim: int = 768
rope_theta: int = 256
qk_norm: str = "rms_norm"
exclude_lora_layers: list[str] = field(
default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"]
)
def __post_init__(self):
super().__post_init__()
self.hidden_size: int = self.attention_head_dim * self.num_attention_heads
self.num_channels_latents: int = self.in_channels
@dataclass
class HunyuanVideoConfig(DiTConfig):
arch_config: DiTArchConfig = field(default_factory=HunyuanVideoArchConfig)
prefix: str = "Hunyuan"
@@ -0,0 +1,36 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from typing import Tuple
from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
@dataclass
class QwenImageArchConfig(DiTArchConfig):
patch_size: int = 1
in_channels: int = 64
out_channels: int | None = None
num_layers: int = 19
num_single_layers: int = 38
attention_head_dim: int = 128
num_attention_heads: int = 24
joint_attention_dim: int = 4096
pooled_projection_dim: int = 768
guidance_embeds: bool = False
axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)
def __post_init__(self):
super().__post_init__()
self.out_channels = self.out_channels or self.in_channels
self.hidden_size = self.num_attention_heads * self.attention_head_dim
self.num_channels_latents = self.out_channels
@dataclass
class QwenImageDitConfig(DiTConfig):
arch_config: DiTArchConfig = field(default_factory=QwenImageArchConfig)
prefix: str = "qwenimage"
@@ -0,0 +1,64 @@
# 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"
@@ -0,0 +1,103 @@
# 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_blocks(n: str, m) -> bool:
return "blocks" in n and str.isdigit(n.split(".")[-1])
@dataclass
class WanVideoArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
param_names_mapping: dict = field(
default_factory=lambda: {
r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1",
r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1",
r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$": r"condition_embedder.image_embedder.ff.fc_in.\1",
r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$": r"condition_embedder.image_embedder.ff.fc_out.\1",
r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$": r"blocks.\1.to_q.\2",
r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$": r"blocks.\1.to_k.\2",
r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$": r"blocks.\1.to_v.\2",
r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": r"blocks.\1.to_out.\2",
r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2",
r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2",
r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"blocks.\1.attn2.to_out.\2",
r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2",
}
)
# Reverse mapping for saving checkpoints: custom -> hf
reverse_param_names_mapping: dict = field(default_factory=lambda: {})
# Some LoRA adapters use the original official layer names instead of hf layer names,
# so apply this before the param_names_mapping
lora_param_names_mapping: dict = field(
default_factory=lambda: {
r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2",
r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2",
r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.attn1.to_v.\2",
r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.attn1.to_out.0.\2",
r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2",
r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2",
r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2",
r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.0.\2",
r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
}
)
patch_size: tuple[int, int, int] = (1, 2, 2)
text_len = 512
num_attention_heads: int = 40
attention_head_dim: int = 128
in_channels: int = 16
out_channels: int = 16
text_dim: int = 4096
freq_dim: int = 256
ffn_dim: int = 13824
num_layers: int = 40
cross_attn_norm: bool = True
qk_norm: str = "rms_norm_across_heads"
eps: float = 1e-6
image_dim: int | None = None
added_kv_proj_dim: int | None = None
rope_max_seq_len: int = 1024
pos_embed_seq_len: int | None = None
exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
# Wan MoE
boundary_ratio: float | None = None
# Causal Wan
local_attn_size: int = (
-1
) # Window size for temporal local attention (-1 indicates global attention)
sink_size: int = (
0 # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache
)
num_frames_per_block: int = 3
sliding_window_num_frames: int = 21
def __post_init__(self):
super().__post_init__()
self.out_channels = self.out_channels or self.in_channels
self.hidden_size = self.num_attention_heads * self.attention_head_dim
self.num_channels_latents = self.out_channels
@dataclass
class WanVideoConfig(DiTConfig):
arch_config: DiTArchConfig = field(default_factory=WanVideoArchConfig)
prefix: str = "Wan"
@@ -0,0 +1,25 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
from sglang.multimodal_gen.configs.models.encoders.base import (
BaseEncoderOutput,
EncoderConfig,
ImageEncoderConfig,
TextEncoderConfig,
)
from sglang.multimodal_gen.configs.models.encoders.clip import (
CLIPTextConfig,
CLIPVisionConfig,
)
from sglang.multimodal_gen.configs.models.encoders.llama import LlamaConfig
from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config
__all__ = [
"EncoderConfig",
"TextEncoderConfig",
"ImageEncoderConfig",
"BaseEncoderOutput",
"CLIPTextConfig",
"CLIPVisionConfig",
"LlamaConfig",
"T5Config",
]
@@ -0,0 +1,85 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from typing import Any
import torch
from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig
from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
@dataclass
class EncoderArchConfig(ArchConfig):
architectures: list[str] = field(default_factory=lambda: [])
_supported_attention_backends: set[AttentionBackendEnum] = field(
default_factory=lambda: {
AttentionBackendEnum.FA3,
AttentionBackendEnum.TORCH_SDPA,
}
)
output_hidden_states: bool = False
use_return_dict: bool = True
@dataclass
class TextEncoderArchConfig(EncoderArchConfig):
vocab_size: int = 0
hidden_size: int = 0
num_hidden_layers: int = 0
num_attention_heads: int = 0
pad_token_id: int = 0
eos_token_id: int = 0
text_len: int = 0
hidden_state_skip_layer: int = 0
decoder_start_token_id: int = 0
output_past: bool = True
scalable_attention: bool = True
tie_word_embeddings: bool = False
stacked_params_mapping: list[tuple[str, str, str]] = field(
default_factory=list
) # mapping from huggingface weight names to custom names
tokenizer_kwargs: dict[str, Any] = field(default_factory=dict)
_fsdp_shard_conditions: list = field(default_factory=lambda: [])
def __post_init__(self) -> None:
self.tokenizer_kwargs = {
"truncation": True,
"max_length": self.text_len,
"return_tensors": "pt",
}
@dataclass
class ImageEncoderArchConfig(EncoderArchConfig):
pass
@dataclass
class BaseEncoderOutput:
last_hidden_state: torch.FloatTensor | None = None
pooler_output: torch.FloatTensor | None = None
hidden_states: tuple[torch.FloatTensor, ...] | None = None
attentions: tuple[torch.FloatTensor, ...] | None = None
attention_mask: torch.Tensor | None = None
@dataclass
class EncoderConfig(ModelConfig):
arch_config: ArchConfig = field(default_factory=EncoderArchConfig)
prefix: str = ""
quant_config: QuantizationConfig | None = None
lora_config: Any | None = None
@dataclass
class TextEncoderConfig(EncoderConfig):
arch_config: ArchConfig = field(default_factory=TextEncoderArchConfig)
@dataclass
class ImageEncoderConfig(EncoderConfig):
arch_config: ArchConfig = field(default_factory=ImageEncoderArchConfig)
@@ -0,0 +1,95 @@
# 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.encoders.base import (
ImageEncoderArchConfig,
ImageEncoderConfig,
TextEncoderArchConfig,
TextEncoderConfig,
)
def _is_transformer_layer(n: str, m) -> bool:
return "layers" in n and str.isdigit(n.split(".")[-1])
def _is_embeddings(n: str, m) -> bool:
return n.endswith("embeddings")
@dataclass
class CLIPTextArchConfig(TextEncoderArchConfig):
vocab_size: int = 49408
hidden_size: int = 512
intermediate_size: int = 2048
projection_dim: int = 512
num_hidden_layers: int = 12
num_attention_heads: int = 8
max_position_embeddings: int = 77
hidden_act: str = "quick_gelu"
layer_norm_eps: float = 1e-5
dropout: float = 0.0
attention_dropout: float = 0.0
initializer_range: float = 0.02
initializer_factor: float = 1.0
pad_token_id: int = 1
bos_token_id: int = 49406
eos_token_id: int = 49407
text_len: int = 77
stacked_params_mapping: list[tuple[str, str, str]] = field(
default_factory=lambda: [
# (param_name, shard_name, shard_id)
("qkv_proj", "q_proj", "q"),
("qkv_proj", "k_proj", "k"),
("qkv_proj", "v_proj", "v"),
]
)
_fsdp_shard_conditions: list = field(
default_factory=lambda: [_is_transformer_layer, _is_embeddings]
)
@dataclass
class CLIPVisionArchConfig(ImageEncoderArchConfig):
hidden_size: int = 768
intermediate_size: int = 3072
projection_dim: int = 512
num_hidden_layers: int = 12
num_attention_heads: int = 12
num_channels: int = 3
image_size: int = 224
patch_size: int = 32
hidden_act: str = "quick_gelu"
layer_norm_eps: float = 1e-5
dropout: float = 0.0
attention_dropout: float = 0.0
initializer_range: float = 0.02
initializer_factor: float = 1.0
stacked_params_mapping: list[tuple[str, str, str]] = field(
default_factory=lambda: [
# (param_name, shard_name, shard_id)
("qkv_proj", "q_proj", "q"),
("qkv_proj", "k_proj", "k"),
("qkv_proj", "v_proj", "v"),
]
)
@dataclass
class CLIPTextConfig(TextEncoderConfig):
arch_config: TextEncoderArchConfig = field(default_factory=CLIPTextArchConfig)
num_hidden_layers_override: int | None = None
require_post_norm: bool | None = None
prefix: str = "clip"
@dataclass
class CLIPVisionConfig(ImageEncoderConfig):
arch_config: ImageEncoderArchConfig = field(default_factory=CLIPVisionArchConfig)
num_hidden_layers_override: int | None = None
require_post_norm: bool | None = None
prefix: str = "clip"
@@ -0,0 +1,69 @@
# 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.encoders.base import (
TextEncoderArchConfig,
TextEncoderConfig,
)
def _is_transformer_layer(n: str, m) -> bool:
return "layers" in n and str.isdigit(n.split(".")[-1])
def _is_embeddings(n: str, m) -> bool:
return n.endswith("embed_tokens")
def _is_final_norm(n: str, m) -> bool:
return n.endswith("norm")
@dataclass
class LlamaArchConfig(TextEncoderArchConfig):
vocab_size: int = 32000
hidden_size: int = 4096
intermediate_size: int = 11008
num_hidden_layers: int = 32
num_attention_heads: int = 32
num_key_value_heads: int | None = None
hidden_act: str = "silu"
max_position_embeddings: int = 2048
initializer_range: float = 0.02
rms_norm_eps: float = 1e-6
use_cache: bool = True
pad_token_id: int = 0
bos_token_id: int = 1
eos_token_id: int = 2
pretraining_tp: int = 1
tie_word_embeddings: bool = False
rope_theta: float = 10000.0
rope_scaling: float | None = None
attention_bias: bool = False
attention_dropout: float = 0.0
mlp_bias: bool = False
head_dim: int | None = None
hidden_state_skip_layer: int = 2
text_len: int = 256
stacked_params_mapping: list[tuple[str, str, str]] = field(
default_factory=lambda: [
# (param_name, shard_name, shard_id)
(".qkv_proj", ".q_proj", "q"),
(".qkv_proj", ".k_proj", "k"),
(".qkv_proj", ".v_proj", "v"),
(".gate_up_proj", ".gate_proj", 0), # type: ignore
(".gate_up_proj", ".up_proj", 1), # type: ignore
]
)
_fsdp_shard_conditions: list = field(
default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm]
)
@dataclass
class LlamaConfig(TextEncoderConfig):
arch_config: TextEncoderArchConfig = field(default_factory=LlamaArchConfig)
prefix: str = "llama"
@@ -0,0 +1,67 @@
# 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.encoders.base import (
TextEncoderArchConfig,
TextEncoderConfig,
)
def _is_transformer_layer(n: str, m) -> bool:
return "layers" in n and str.isdigit(n.split(".")[-1])
def _is_embeddings(n: str, m) -> bool:
return n.endswith("embed_tokens")
def _is_final_norm(n: str, m) -> bool:
return n.endswith("norm")
@dataclass
class QwenImageArchConfig(TextEncoderArchConfig):
vocab_size: int = 32000
hidden_size: int = 4096
intermediate_size: int = 11008
num_hidden_layers: int = 32
num_attention_heads: int = 32
num_key_value_heads: int | None = None
hidden_act: str = "silu"
max_position_embeddings: int = 2048
initializer_range: float = 0.02
rms_norm_eps: float = 1e-6
use_cache: bool = True
pad_token_id: int = -1
eos_token_id: int = 2
pretraining_tp: int = 1
tie_word_embeddings: bool = False
rope_theta: float = 10000.0
rope_scaling: float | None = None
attention_bias: bool = False
attention_dropout: float = 0.0
mlp_bias: bool = False
head_dim: int | None = None
hidden_state_skip_layer: int = 2
text_len: int = 256
stacked_params_mapping: list[tuple[str, str, str]] = field(
default_factory=lambda: [
# (param_name, shard_name, shard_id)
(".qkv_proj", ".q_proj", "q"),
(".qkv_proj", ".k_proj", "k"),
(".qkv_proj", ".v_proj", "v"),
(".gate_up_proj", ".gate_proj", 0), # type: ignore
(".gate_up_proj", ".up_proj", 1), # type: ignore
]
)
_fsdp_shard_conditions: list = field(
default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm]
)
@dataclass
class Qwen2_5VLConfig(TextEncoderConfig):
arch_config: TextEncoderArchConfig = field(default_factory=QwenImageArchConfig)
# prefix: str = "qwen_image"
@@ -0,0 +1,86 @@
# 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.encoders.base import (
TextEncoderArchConfig,
TextEncoderConfig,
)
def _is_transformer_layer(n: str, m) -> bool:
return "block" in n and str.isdigit(n.split(".")[-1])
def _is_embeddings(n: str, m) -> bool:
return n.endswith("shared")
def _is_final_layernorm(n: str, m) -> bool:
return n.endswith("final_layer_norm")
@dataclass
class T5ArchConfig(TextEncoderArchConfig):
vocab_size: int = 32128
d_model: int = 512
d_kv: int = 64
d_ff: int = 2048
num_layers: int = 6
num_decoder_layers: int | None = None
num_heads: int = 8
relative_attention_num_buckets: int = 32
relative_attention_max_distance: int = 128
dropout_rate: float = 0.1
layer_norm_epsilon: float = 1e-6
initializer_factor: float = 1.0
feed_forward_proj: str = "relu"
dense_act_fn: str = ""
is_gated_act: bool = False
is_encoder_decoder: bool = True
use_cache: bool = True
pad_token_id: int = 0
eos_token_id: int = 1
classifier_dropout: float = 0.0
text_len: int = 512
stacked_params_mapping: list[tuple[str, str, str]] = field(
default_factory=lambda: [
# (param_name, shard_name, shard_id)
(".qkv_proj", ".q", "q"),
(".qkv_proj", ".k", "k"),
(".qkv_proj", ".v", "v"),
]
)
_fsdp_shard_conditions: list = field(
default_factory=lambda: [
_is_transformer_layer,
_is_embeddings,
_is_final_layernorm,
]
)
# Referenced from https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/configuration_t5.py
def __post_init__(self):
super().__post_init__()
act_info = self.feed_forward_proj.split("-")
self.dense_act_fn: str = act_info[-1]
self.is_gated_act: bool = act_info[0] == "gated"
if self.feed_forward_proj == "gated-gelu":
self.dense_act_fn = "gelu_new"
self.tokenizer_kwargs = {
"padding": "max_length",
"truncation": True,
"max_length": self.text_len,
"add_special_tokens": True,
"return_attention_mask": True,
"return_tensors": "pt",
}
@dataclass
class T5Config(TextEncoderConfig):
arch_config: TextEncoderArchConfig = field(default_factory=T5ArchConfig)
prefix: str = "t5"
@@ -0,0 +1,11 @@
# 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",
]
@@ -0,0 +1,158 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
import argparse
import dataclasses
from dataclasses import dataclass, field
from typing import Any
import torch
from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig
from sglang.multimodal_gen.runtime.models.vision_utils import get_default_height_width
from sglang.multimodal_gen.utils import StoreBoolean
@dataclass
class VAEArchConfig(ArchConfig):
scaling_factor: float | torch.Tensor = 0
temporal_compression_ratio: int = 4
# or vae_scale_factor?
spatial_compression_ratio: int = 8
@dataclass
class VAEConfig(ModelConfig):
arch_config: VAEArchConfig = field(default_factory=VAEArchConfig)
# sgl-diffusionVAE-specific parameters
load_encoder: bool = True
load_decoder: bool = True
tile_sample_min_height: int = 256
tile_sample_min_width: int = 256
tile_sample_min_num_frames: int = 16
tile_sample_stride_height: int = 192
tile_sample_stride_width: int = 192
tile_sample_stride_num_frames: int = 12
blend_num_frames: int = 0
use_tiling: bool = True
use_temporal_tiling: bool = True
use_parallel_tiling: bool = True
use_temporal_scaling_frames: bool = True
def __post_init__(self):
self.blend_num_frames = (
self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
)
def post_init(self):
pass
# returns width, height
def calculate_dimensions(
self, image, vae_scale_factor, width, height
) -> tuple[int, int]:
height, width = get_default_height_width(image, vae_scale_factor, height, width)
return width, height
@staticmethod
def add_cli_args(parser: Any, prefix: str = "vae-config") -> Any:
"""Add CLI arguments for VAEConfig fields"""
parser.add_argument(
f"--{prefix}.load-encoder",
action=StoreBoolean,
dest=f"{prefix.replace('-', '_')}.load_encoder",
default=VAEConfig.load_encoder,
help="Whether to load the VAE encoder",
)
parser.add_argument(
f"--{prefix}.load-decoder",
action=StoreBoolean,
dest=f"{prefix.replace('-', '_')}.load_decoder",
default=VAEConfig.load_decoder,
help="Whether to load the VAE decoder",
)
parser.add_argument(
f"--{prefix}.tile-sample-min-height",
type=int,
dest=f"{prefix.replace('-', '_')}.tile_sample_min_height",
default=VAEConfig.tile_sample_min_height,
help="Minimum height for VAE tile sampling",
)
parser.add_argument(
f"--{prefix}.tile-sample-min-width",
type=int,
dest=f"{prefix.replace('-', '_')}.tile_sample_min_width",
default=VAEConfig.tile_sample_min_width,
help="Minimum width for VAE tile sampling",
)
parser.add_argument(
f"--{prefix}.tile-sample-min-num-frames",
type=int,
dest=f"{prefix.replace('-', '_')}.tile_sample_min_num_frames",
default=VAEConfig.tile_sample_min_num_frames,
help="Minimum number of frames for VAE tile sampling",
)
parser.add_argument(
f"--{prefix}.tile-sample-stride-height",
type=int,
dest=f"{prefix.replace('-', '_')}.tile_sample_stride_height",
default=VAEConfig.tile_sample_stride_height,
help="Stride height for VAE tile sampling",
)
parser.add_argument(
f"--{prefix}.tile-sample-stride-width",
type=int,
dest=f"{prefix.replace('-', '_')}.tile_sample_stride_width",
default=VAEConfig.tile_sample_stride_width,
help="Stride width for VAE tile sampling",
)
parser.add_argument(
f"--{prefix}.tile-sample-stride-num-frames",
type=int,
dest=f"{prefix.replace('-', '_')}.tile_sample_stride_num_frames",
default=VAEConfig.tile_sample_stride_num_frames,
help="Stride number of frames for VAE tile sampling",
)
parser.add_argument(
f"--{prefix}.blend-num-frames",
type=int,
dest=f"{prefix.replace('-', '_')}.blend_num_frames",
default=VAEConfig.blend_num_frames,
help="Number of frames to blend for VAE tile sampling",
)
parser.add_argument(
f"--{prefix}.use-tiling",
action=StoreBoolean,
dest=f"{prefix.replace('-', '_')}.use_tiling",
default=VAEConfig.use_tiling,
help="Whether to use tiling for VAE",
)
parser.add_argument(
f"--{prefix}.use-temporal-tiling",
action=StoreBoolean,
dest=f"{prefix.replace('-', '_')}.use_temporal_tiling",
default=VAEConfig.use_temporal_tiling,
help="Whether to use temporal tiling for VAE",
)
parser.add_argument(
f"--{prefix}.use-parallel-tiling",
action=StoreBoolean,
dest=f"{prefix.replace('-', '_')}.use_parallel_tiling",
default=VAEConfig.use_parallel_tiling,
help="Whether to use parallel tiling for VAE",
)
return parser
@classmethod
def from_cli_args(cls, args: argparse.Namespace) -> "VAEConfig":
kwargs = {}
for attr in dataclasses.fields(cls):
value = getattr(args, attr.name, None)
if value is not None:
kwargs[attr.name] = value
return cls(**kwargs)
@@ -0,0 +1,50 @@
# 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 FluxVAEArchConfig(VAEArchConfig):
spatial_compression_ratio: int = 1
base_dim: int = 96
decoder_base_dim: int | None = None
z_dim: int = 16
dim_mult: tuple[int, ...] = (1, 2, 4, 4)
num_res_blocks: int = 2
attn_scales: tuple[float, ...] = ()
temperal_downsample: tuple[bool, ...] = (False, True, True)
dropout: float = 0.0
is_residual: bool = False
in_channels: int = 3
out_channels: int = 3
patch_size: int | None = None
scale_factor_temporal: int = 4
scale_factor_spatial: int = 8
clip_output: bool = True
@dataclass
class FluxVAEConfig(VAEConfig):
arch_config: FluxVAEArchConfig = field(default_factory=FluxVAEArchConfig)
use_feature_cache: bool = True
use_tiling: bool = False
use_temporal_tiling: bool = False
use_parallel_tiling: bool = False
def __post_init__(self):
self.blend_num_frames = (
self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
) * 2
def post_init(self):
self.arch_config.vae_scale_factor = 2 ** (
len(self.arch_config.block_out_channels) - 1
)
self.arch_config.spatial_compression_ratio = self.arch_config.vae_scale_factor
@@ -0,0 +1,41 @@
# 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 HunyuanVAEArchConfig(VAEArchConfig):
in_channels: int = 3
out_channels: int = 3
latent_channels: int = 16
down_block_types: tuple[str, ...] = (
"HunyuanVideoDownBlock3D",
"HunyuanVideoDownBlock3D",
"HunyuanVideoDownBlock3D",
"HunyuanVideoDownBlock3D",
)
up_block_types: tuple[str, ...] = (
"HunyuanVideoUpBlock3D",
"HunyuanVideoUpBlock3D",
"HunyuanVideoUpBlock3D",
"HunyuanVideoUpBlock3D",
)
block_out_channels: tuple[int, ...] = (128, 256, 512, 512)
layers_per_block: int = 2
act_fn: str = "silu"
norm_num_groups: int = 32
scaling_factor: float = 0.476986
spatial_compression_ratio: int = 8
temporal_compression_ratio: int = 4
mid_block_add_attention: bool = True
def __post_init__(self):
self.spatial_compression_ratio: int = 2 ** (len(self.block_out_channels) - 1)
@dataclass
class HunyuanVAEConfig(VAEConfig):
arch_config: VAEArchConfig = field(default_factory=HunyuanVAEArchConfig)
@@ -0,0 +1,61 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit import calculate_dimensions
from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig
@dataclass
class QwenImageVAEArchConfig(VAEArchConfig):
spatial_compression_ratio: int = 1
base_dim: int = 96
decoder_base_dim: int | None = None
z_dim: int = 16
dim_mult: tuple[int, ...] = (1, 2, 4, 4)
num_res_blocks: int = 2
attn_scales: tuple[float, ...] = ()
temperal_downsample: tuple[bool, ...] = (False, True, True)
dropout: float = 0.0
is_residual: bool = False
in_channels: int = 3
out_channels: int = 3
patch_size: int | None = None
scale_factor_temporal: int = 4
scale_factor_spatial: int = 8
clip_output: bool = True
def __post_init__(self):
self.vae_scale_factor = 2 ** len(self.temperal_downsample)
@dataclass
class QwenImageVAEConfig(VAEConfig):
arch_config: QwenImageVAEArchConfig = field(default_factory=QwenImageVAEArchConfig)
use_feature_cache: bool = True
use_tiling: bool = False
use_temporal_tiling: bool = False
use_parallel_tiling: bool = False
def calculate_dimensions(self, image, vae_scale_factor, width, height):
width = image.size[0]
height = image.size[1]
width, height, _ = calculate_dimensions(1024 * 1024, width / height)
return width, height
def __post_init__(self):
self.blend_num_frames = (
self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
) * 2
def post_init(self):
self.arch_config.vae_scale_factor = 2 ** (
len(self.arch_config.temperal_downsample)
)
self.arch_config.spatial_compression_ratio = self.arch_config.vae_scale_factor
@@ -0,0 +1,31 @@
# 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
@@ -0,0 +1,88 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
import torch
from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig
@dataclass
class WanVAEArchConfig(VAEArchConfig):
base_dim: int = 96
decoder_base_dim: int | None = None
z_dim: int = 16
dim_mult: tuple[int, ...] = (1, 2, 4, 4)
num_res_blocks: int = 2
attn_scales: tuple[float, ...] = ()
temperal_downsample: tuple[bool, ...] = (False, True, True)
dropout: float = 0.0
latents_mean: tuple[float, ...] = (
-0.7571,
-0.7089,
-0.9113,
0.1075,
-0.1745,
0.9653,
-0.1517,
1.5508,
0.4134,
-0.0715,
0.5517,
-0.3632,
-0.1922,
-0.9497,
0.2503,
-0.2921,
)
latents_std: tuple[float, ...] = (
2.8184,
1.4541,
2.3275,
2.6558,
1.2196,
1.7708,
2.6052,
2.0743,
3.2687,
2.1526,
2.8652,
1.5579,
1.6382,
1.1253,
2.8251,
1.9160,
)
is_residual: bool = False
in_channels: int = 3
out_channels: int = 3
patch_size: int | None = None
scale_factor_temporal: int = 4
scale_factor_spatial: int = 8
clip_output: bool = True
def __post_init__(self):
self.scaling_factor: torch.tensor = 1.0 / torch.tensor(self.latents_std).view(
1, self.z_dim, 1, 1, 1
)
self.shift_factor: torch.tensor = torch.tensor(self.latents_mean).view(
1, self.z_dim, 1, 1, 1
)
self.temporal_compression_ratio = self.scale_factor_temporal
self.spatial_compression_ratio = self.scale_factor_spatial
@dataclass
class WanVAEConfig(VAEConfig):
arch_config: WanVAEArchConfig = field(default_factory=WanVAEArchConfig)
use_feature_cache: bool = True
use_tiling: bool = False
use_temporal_tiling: bool = False
use_parallel_tiling: bool = False
def __post_init__(self):
self.blend_num_frames = (
self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
) * 2
@@ -0,0 +1,37 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
from sglang.multimodal_gen.configs.pipelines.base import (
PipelineConfig,
SlidingTileAttnConfig,
)
from sglang.multimodal_gen.configs.pipelines.flux import FluxPipelineConfig
from sglang.multimodal_gen.configs.pipelines.hunyuan import (
FastHunyuanConfig,
HunyuanConfig,
)
from sglang.multimodal_gen.configs.pipelines.registry import (
get_pipeline_config_cls_from_name,
)
from sglang.multimodal_gen.configs.pipelines.stepvideo import StepVideoT2VConfig
from sglang.multimodal_gen.configs.pipelines.wan import (
SelfForcingWanT2V480PConfig,
WanI2V480PConfig,
WanI2V720PConfig,
WanT2V480PConfig,
WanT2V720PConfig,
)
__all__ = [
"HunyuanConfig",
"FastHunyuanConfig",
"FluxPipelineConfig",
"PipelineConfig",
"SlidingTileAttnConfig",
"WanT2V480PConfig",
"WanI2V480PConfig",
"WanT2V720PConfig",
"WanI2V720PConfig",
"StepVideoT2VConfig",
"SelfForcingWanT2V480PConfig",
"get_pipeline_config_cls_from_name",
]
@@ -0,0 +1,485 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
import json
from collections.abc import Callable
from dataclasses import asdict, dataclass, field, fields
from enum import Enum
from typing import Any, cast
import torch
from diffusers.image_processor import VaeImageProcessor
from sglang.multimodal_gen.configs.models import (
DiTConfig,
EncoderConfig,
ModelConfig,
VAEConfig,
)
from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput
from sglang.multimodal_gen.configs.utils import update_config_from_args
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.utils import (
FlexibleArgumentParser,
StoreBoolean,
shallow_asdict,
)
logger = init_logger(__name__)
class STA_Mode(str, Enum):
"""STA (Sliding Tile Attention) modes."""
STA_INFERENCE = "STA_inference"
STA_SEARCHING = "STA_searching"
STA_TUNING = "STA_tuning"
STA_TUNING_CFG = "STA_tuning_cfg"
NONE = None
def preprocess_text(prompt: str) -> str:
return prompt
def postprocess_text(output: BaseEncoderOutput, _text_inputs) -> torch.tensor:
raise NotImplementedError
# config for a single pipeline
@dataclass
class PipelineConfig:
"""Base configuration for all pipeline architectures."""
model_path: str = ""
pipeline_config_path: str | None = None
is_image_gen: bool = False
# generation parameters
# controls the timestep embedding generation
should_use_guidance: bool = True
embedded_cfg_scale: float = 6.0
flow_shift: float | None = None
disable_autocast: bool = False
# Model configuration
dit_config: DiTConfig = field(default_factory=DiTConfig)
dit_precision: str = "bf16"
# VAE configuration
vae_config: VAEConfig = field(default_factory=VAEConfig)
vae_precision: str = "fp32"
vae_tiling: bool = True
vae_sp: bool = True
# Image encoder configuration
image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig)
image_encoder_precision: str = "fp32"
# Text encoder configuration
DEFAULT_TEXT_ENCODER_PRECISIONS = ("fp32",)
text_encoder_configs: tuple[EncoderConfig, ...] = field(
default_factory=lambda: (EncoderConfig(),)
)
# See PRECISION_TO_TYPE for detailed mapping
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32",))
text_encoder_extra_args: list[dict] = field(default_factory=lambda: [{}])
# image encoding
image_encoder_extra_args: dict = field(default_factory=lambda: {})
def postprocess_image(self, image):
return image.last_hidden_state
preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
default_factory=lambda: (preprocess_text,)
)
postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.tensor], ...] = (
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
skip_time_steps: int = 15
# DMD parameters
dmd_denoising_steps: list[int] | None = field(default=None)
# Wan2.2 TI2V parameters
ti2v_task: bool = False
i2v_task: bool = False
ti2i_task: bool = False
boundary_ratio: float | None = None
# Compilation
# enable_torch_compile: bool = False
def slice_noise_pred(self, noise, latents):
return noise
def set_width_and_height(self, width, height, image):
"""
image: input image
"""
return width, height
# called in ImageEncodingStage, preprocess the image
def preprocess_image(self, image, image_processor: VaeImageProcessor):
return image
def prepare_latent_shape(self, batch, batch_size, num_frames):
height = batch.height // self.vae_config.arch_config.spatial_compression_ratio
width = batch.width // self.vae_config.arch_config.spatial_compression_ratio
# Calculate latent shape
shape = (
batch_size,
self.dit_config.num_channels_latents,
num_frames,
height,
width,
)
return shape
# called after latents are prepared
def pack_latents(self, latents, batch_size, batch):
return latents
def get_pos_prompt_embeds(self, batch):
return batch.prompt_embeds
def get_neg_prompt_embeds(self, batch):
return batch.negative_prompt_embeds
def post_denoising_loop(self, latents, batch):
return latents
def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
return {}
def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
return {}
@staticmethod
def add_cli_args(
parser: FlexibleArgumentParser, prefix: str = ""
) -> FlexibleArgumentParser:
prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else ""
# model_path will be conflicting with the model_path in ServerArgs,
# so we add it separately if prefix is not empty
if prefix_with_dot != "":
parser.add_argument(
f"--{prefix_with_dot}model-path",
type=str,
dest=f"{prefix_with_dot.replace('-', '_')}model_path",
default=PipelineConfig.model_path,
help="Path to the pretrained model",
)
parser.add_argument(
f"--{prefix_with_dot}pipeline-config-path",
type=str,
dest=f"{prefix_with_dot.replace('-', '_')}pipeline_config_path",
default=PipelineConfig.pipeline_config_path,
help="Path to the pipeline config",
)
parser.add_argument(
f"--{prefix_with_dot}embedded-cfg-scale",
type=float,
dest=f"{prefix_with_dot.replace('-', '_')}embedded_cfg_scale",
default=PipelineConfig.embedded_cfg_scale,
help="Embedded CFG scale",
)
parser.add_argument(
f"--{prefix_with_dot}flow-shift",
type=float,
dest=f"{prefix_with_dot.replace('-', '_')}flow_shift",
default=PipelineConfig.flow_shift,
help="Flow shift parameter",
)
# DiT configuration
parser.add_argument(
f"--{prefix_with_dot}dit-precision",
type=str,
dest=f"{prefix_with_dot.replace('-', '_')}dit_precision",
default=PipelineConfig.dit_precision,
choices=["fp32", "fp16", "bf16"],
help="Precision for the DiT model",
)
# VAE configuration
parser.add_argument(
f"--{prefix_with_dot}vae-precision",
type=str,
dest=f"{prefix_with_dot.replace('-', '_')}vae_precision",
default=PipelineConfig.vae_precision,
choices=["fp32", "fp16", "bf16"],
help="Precision for VAE",
)
parser.add_argument(
f"--{prefix_with_dot}vae-tiling",
action=StoreBoolean,
dest=f"{prefix_with_dot.replace('-', '_')}vae_tiling",
default=PipelineConfig.vae_tiling,
help="Enable VAE tiling",
)
parser.add_argument(
f"--{prefix_with_dot}vae-sp",
action=StoreBoolean,
dest=f"{prefix_with_dot.replace('-', '_')}vae_sp",
help="Enable VAE spatial parallelism",
)
# Text encoder configuration
parser.add_argument(
f"--{prefix_with_dot}text-encoder-precisions",
nargs="+",
type=str,
dest=f"{prefix_with_dot.replace('-', '_')}text_encoder_precisions",
default=PipelineConfig.DEFAULT_TEXT_ENCODER_PRECISIONS,
choices=["fp32", "fp16", "bf16"],
help="Precision for each text encoder",
)
# Image encoder configuration
parser.add_argument(
f"--{prefix_with_dot}image-encoder-precision",
type=str,
dest=f"{prefix_with_dot.replace('-', '_')}image_encoder_precision",
default=PipelineConfig.image_encoder_precision,
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(
f"--{prefix_with_dot}dmd-denoising-steps",
type=parse_int_list,
default=PipelineConfig.dmd_denoising_steps,
help="Comma-separated list of denoising steps (e.g., '1000,757,522')",
)
# Add VAE configuration arguments
from sglang.multimodal_gen.configs.models.vaes.base import VAEConfig
VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config")
# Add DiT configuration arguments
from sglang.multimodal_gen.configs.models.dits.base import DiTConfig
DiTConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}dit-config")
return parser
def update_config_from_dict(self, args: dict[str, Any], prefix: str = "") -> None:
prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else ""
update_config_from_args(self, args, prefix, pop_args=True)
update_config_from_args(
self.vae_config, args, f"{prefix_with_dot}vae_config", pop_args=True
)
update_config_from_args(
self.dit_config, args, f"{prefix_with_dot}dit_config", pop_args=True
)
@classmethod
def from_pretrained(cls, model_path: str) -> "PipelineConfig":
"""
use the pipeline class setting from model_path to match the pipeline config
"""
from sglang.multimodal_gen.configs.pipelines.registry import (
get_pipeline_config_cls_from_name,
)
pipeline_config_cls = get_pipeline_config_cls_from_name(model_path)
return cast(PipelineConfig, pipeline_config_cls(model_path=model_path))
@classmethod
def from_kwargs(
cls, kwargs: dict[str, Any], config_cli_prefix: str = ""
) -> "PipelineConfig":
"""
Load PipelineConfig from kwargs Dictionary.
kwargs: dictionary of kwargs
config_cli_prefix: prefix of CLI arguments for this PipelineConfig instance
"""
from sglang.multimodal_gen.configs.pipelines.registry import (
get_pipeline_config_cls_from_name,
)
prefix_with_dot = (
f"{config_cli_prefix}." if (config_cli_prefix.strip() != "") else ""
)
model_path: str | None = kwargs.get(
prefix_with_dot + "model_path", None
) or kwargs.get("model_path")
pipeline_config_or_path: str | PipelineConfig | dict[str, Any] | None = (
kwargs.get(prefix_with_dot + "pipeline_config", None)
or kwargs.get("pipeline_config")
)
if model_path is None:
raise ValueError("model_path is required in kwargs")
# 1. Get the pipeline config class from the registry
pipeline_config_cls = get_pipeline_config_cls_from_name(model_path)
# 2. Instantiate PipelineConfig
if pipeline_config_cls is None:
logger.warning(
"Couldn't find pipeline config for %s. Using the default pipeline config.",
model_path,
)
pipeline_config = cls()
else:
pipeline_config = pipeline_config_cls()
# 3. Load PipelineConfig from a json file or a PipelineConfig object if provided
if isinstance(pipeline_config_or_path, str):
pipeline_config.load_from_json(pipeline_config_or_path)
kwargs[prefix_with_dot + "pipeline_config_path"] = pipeline_config_or_path
elif isinstance(pipeline_config_or_path, PipelineConfig):
pipeline_config = pipeline_config_or_path
elif isinstance(pipeline_config_or_path, dict):
pipeline_config.update_pipeline_config(pipeline_config_or_path)
# 4. Update PipelineConfig from CLI arguments if provided
kwargs[prefix_with_dot + "model_path"] = model_path
pipeline_config.update_config_from_dict(kwargs, config_cli_prefix)
return pipeline_config
def check_pipeline_config(self) -> None:
if self.vae_sp and not self.vae_tiling:
raise ValueError(
"Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True."
)
if len(self.text_encoder_configs) != len(self.text_encoder_precisions):
raise ValueError(
f"Length of text encoder configs ({len(self.text_encoder_configs)}) must be equal to length of text encoder precisions ({len(self.text_encoder_precisions)})"
)
if len(self.text_encoder_configs) != len(self.preprocess_text_funcs):
raise ValueError(
f"Length of text encoder configs ({len(self.text_encoder_configs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})"
)
if len(self.preprocess_text_funcs) != len(self.postprocess_text_funcs):
raise ValueError(
f"Length of text postprocess functions ({len(self.postprocess_text_funcs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})"
)
def dump_to_json(self, file_path: str):
output_dict = shallow_asdict(self)
del_keys = []
for key, value in output_dict.items():
if isinstance(value, ModelConfig):
model_dict = asdict(value)
# Model Arch Config should be hidden away from the users
model_dict.pop("arch_config")
output_dict[key] = model_dict
elif isinstance(value, tuple) and all(
isinstance(v, ModelConfig) for v in value
):
model_dicts = []
for v in value:
model_dict = asdict(v)
# Model Arch Config should be hidden away from the users
model_dict.pop("arch_config")
model_dicts.append(model_dict)
output_dict[key] = model_dicts
elif isinstance(value, tuple) and all(callable(f) for f in value):
# Skip dumping functions
del_keys.append(key)
for key in del_keys:
output_dict.pop(key, None)
with open(file_path, "w") as f:
json.dump(output_dict, f, indent=2)
def load_from_json(self, file_path: str):
with open(file_path) as f:
input_pipeline_dict = json.load(f)
self.update_pipeline_config(input_pipeline_dict)
def update_pipeline_config(self, source_pipeline_dict: dict[str, Any]) -> None:
for f in fields(self):
key = f.name
if key in source_pipeline_dict:
current_value = getattr(self, key)
new_value = source_pipeline_dict[key]
# If it's a nested ModelConfig, update it recursively
if isinstance(current_value, ModelConfig):
current_value.update_model_config(new_value)
elif isinstance(current_value, tuple) and all(
isinstance(v, ModelConfig) for v in current_value
):
assert len(current_value) == len(
new_value
), "Users shouldn't delete or add text encoder config objects in your json"
for target_config, source_config in zip(
current_value, new_value, strict=True
):
target_config.update_model_config(source_config)
else:
setattr(self, key, new_value)
if hasattr(self, "__post_init__"):
self.__post_init__()
@dataclass
class SlidingTileAttnConfig(PipelineConfig):
"""Configuration for sliding tile attention."""
# Override any BaseConfig defaults as needed
# Add sliding tile specific parameters
window_size: int = 16
stride: int = 8
# You can provide custom defaults for inherited fields
height: int = 576
width: int = 1024
# Additional configuration specific to sliding tile attention
pad_to_square: bool = False
use_overlap_optimization: bool = True
def parse_int_list(value: str) -> list[int]:
"""Parse a comma-separated string of integers into a list."""
if not value:
return []
return [int(x.strip()) for x in value.split(",")]
@@ -0,0 +1,174 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
from dataclasses import dataclass, field
from typing import Callable
import torch
from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig
from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig
from sglang.multimodal_gen.configs.models.encoders import (
BaseEncoderOutput,
CLIPTextConfig,
T5Config,
)
from sglang.multimodal_gen.configs.models.vaes.flux import FluxVAEConfig
from sglang.multimodal_gen.configs.pipelines.base import PipelineConfig, preprocess_text
from sglang.multimodal_gen.configs.pipelines.hunyuan import (
clip_postprocess_text,
clip_preprocess_text,
)
def t5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor:
return outputs.last_hidden_state
@dataclass
class FluxPipelineConfig(PipelineConfig):
# FIXME: duplicate with SamplingParams.guidance_scale?
embedded_cfg_scale: float = 3.5
is_image_gen: bool = True
vae_tiling: bool = False
vae_sp: bool = False
dit_config: DiTConfig = field(default_factory=FluxConfig)
# VAE
vae_config: VAEConfig = field(default_factory=FluxVAEConfig)
# Text encoding stage
text_encoder_configs: tuple[EncoderConfig, ...] = field(
default_factory=lambda: (CLIPTextConfig(), T5Config())
)
text_encoder_precisions: tuple[str, ...] = field(
default_factory=lambda: ("bf16", "bf16")
)
preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
default_factory=lambda: (clip_preprocess_text, preprocess_text),
)
postprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
default_factory=lambda: (clip_postprocess_text, t5_postprocess_text)
)
text_encoder_extra_args: list[dict] = field(
default_factory=lambda: [
dict(
max_length=77,
padding="max_length",
truncation=True,
return_overflowing_tokens=False,
return_length=False,
),
None,
]
)
def prepare_latent_shape(self, batch, batch_size, num_frames):
height = 2 * (
batch.height // (self.vae_config.arch_config.vae_scale_factor * 2)
)
width = 2 * (batch.width // (self.vae_config.arch_config.vae_scale_factor * 2))
num_channels_latents = self.dit_config.arch_config.in_channels // 4
shape = (batch_size, num_channels_latents, height, width)
return shape
def pack_latents(self, latents, batch_size, batch):
height = 2 * (
batch.height // (self.vae_config.arch_config.vae_scale_factor * 2)
)
width = 2 * (batch.width // (self.vae_config.arch_config.vae_scale_factor * 2))
num_channels_latents = self.dit_config.arch_config.in_channels // 4
# pack latents
latents = latents.view(
batch_size, num_channels_latents, height // 2, 2, width // 2, 2
)
latents = latents.permute(0, 2, 4, 1, 3, 5)
latents = latents.reshape(
batch_size, (height // 2) * (width // 2), num_channels_latents * 4
)
return latents
def get_pos_prompt_embeds(self, batch):
return batch.prompt_embeds[1]
def get_neg_prompt_embeds(self, batch):
return batch.negative_prompt_embeds[1]
def _prepare_latent_image_ids(self, original_height, original_width, device):
vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
height = int(original_height) // (vae_scale_factor * 2)
width = int(original_width) // (vae_scale_factor * 2)
latent_image_ids = torch.zeros(height, width, 3, device=device)
latent_image_ids[..., 1] = (
latent_image_ids[..., 1] + torch.arange(height, device=device)[:, None]
)
latent_image_ids[..., 2] = (
latent_image_ids[..., 2] + torch.arange(width, device=device)[None, :]
)
latent_image_id_height, latent_image_id_width, latent_image_id_channels = (
latent_image_ids.shape
)
latent_image_ids = latent_image_ids.reshape(
latent_image_id_height * latent_image_id_width, latent_image_id_channels
)
return latent_image_ids
def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb):
txt_ids = torch.zeros(prompt_embeds.shape[1], 3, device=device)
img_ids = self._prepare_latent_image_ids(
original_height=height,
original_width=width,
device=device,
)
ids = torch.cat([txt_ids, img_ids], dim=0).to(device=device)
# NOTE(mick): prepare it here, to avoid unnecessary computations
freqs_cis = rotary_emb.forward(ids)
return freqs_cis
def post_denoising_loop(self, latents, batch):
# unpack latents for flux
# VAE applies 8x compression on images but we must also account for packing which requires
# latent height and width to be divisible by 2.
batch_size = latents.shape[0]
channels = latents.shape[-1]
vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
height = 2 * (int(batch.height) // (vae_scale_factor * 2))
width = 2 * (int(batch.width) // (vae_scale_factor * 2))
latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
latents = latents.permute(0, 3, 1, 4, 2, 5)
latents = latents.reshape(batch_size, channels // (2 * 2), height, width)
return latents
def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
return {
"freqs_cis": self.get_freqs_cis(
batch.prompt_embeds[1], batch.width, batch.height, device, rotary_emb
),
"pooled_projections": (
batch.pooled_embeds[0] if batch.pooled_embeds else None
),
}
def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
return {
"freqs_cis": self.get_freqs_cis(
batch.negative_prompt_embeds[1],
batch.width,
batch.height,
device,
rotary_emb,
),
"pooled_projections": (
batch.neg_pooled_embeds[0] if batch.neg_pooled_embeds else None
),
}
@@ -0,0 +1,109 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import TypedDict
import torch
from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig
from sglang.multimodal_gen.configs.models.dits import HunyuanVideoConfig
from sglang.multimodal_gen.configs.models.encoders import (
BaseEncoderOutput,
CLIPTextConfig,
LlamaConfig,
)
from sglang.multimodal_gen.configs.models.vaes import HunyuanVAEConfig
from sglang.multimodal_gen.configs.pipelines.base import PipelineConfig
PROMPT_TEMPLATE_ENCODE_VIDEO = (
"<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: "
"1. The main content and theme of the video."
"2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects."
"3. Actions, events, behaviors temporal relationships, physical movement changes of the objects."
"4. background environment, light, style and atmosphere."
"5. camera angles, movements, and transitions used in the video:<|eot_id|>"
"<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>"
)
class PromptTemplate(TypedDict):
template: str
crop_start: int
prompt_template_video: PromptTemplate = {
"template": PROMPT_TEMPLATE_ENCODE_VIDEO,
"crop_start": 95,
}
def llama_preprocess_text(prompt: str) -> str:
return prompt_template_video["template"].format(prompt)
def llama_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.tensor:
hidden_state_skip_layer = 2
assert outputs.hidden_states is not None
hidden_states: tuple[torch.Tensor, ...] = outputs.hidden_states
last_hidden_state: torch.tensor = hidden_states[-(hidden_state_skip_layer + 1)]
crop_start = prompt_template_video.get("crop_start", -1)
last_hidden_state = last_hidden_state[:, crop_start:]
return last_hidden_state
def clip_preprocess_text(prompt: str) -> str:
return prompt
def clip_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.tensor:
pooler_output: torch.tensor = outputs.pooler_output
return pooler_output
@dataclass
class HunyuanConfig(PipelineConfig):
"""Base configuration for HunYuan pipeline architecture."""
# HunyuanConfig-specific parameters with defaults
# DiT
dit_config: DiTConfig = field(default_factory=HunyuanVideoConfig)
# VAE
vae_config: VAEConfig = field(default_factory=HunyuanVAEConfig)
# Denoising stage
embedded_cfg_scale: int = 6
flow_shift: int = 7
# Text encoding stage
text_encoder_configs: tuple[EncoderConfig, ...] = field(
default_factory=lambda: (LlamaConfig(), CLIPTextConfig())
)
preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
default_factory=lambda: (llama_preprocess_text, clip_preprocess_text)
)
postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.tensor], ...] = (
field(default_factory=lambda: (llama_postprocess_text, clip_postprocess_text))
)
# Precision for each component
dit_precision: str = "bf16"
vae_precision: str = "fp16"
text_encoder_precisions: tuple[str, ...] = field(
default_factory=lambda: ("fp16", "fp16")
)
def __post_init__(self):
self.vae_config.load_encoder = False
self.vae_config.load_decoder = True
@dataclass
class FastHunyuanConfig(HunyuanConfig):
"""Configuration specifically optimized for FastHunyuan weights."""
# Override HunyuanConfig defaults
flow_shift: int = 17
# No need to re-specify guidance_scale or embedded_cfg_scale as they
# already have the desired values from HunyuanConfig
@@ -0,0 +1,299 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
from dataclasses import dataclass, field
from typing import Callable
import torch
from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit import calculate_dimensions
from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig
from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
from sglang.multimodal_gen.configs.models.encoders.qwen_image import Qwen2_5VLConfig
from sglang.multimodal_gen.configs.models.vaes.qwenimage import QwenImageVAEConfig
from sglang.multimodal_gen.configs.pipelines.base import PipelineConfig
def _extract_masked_hidden(hidden_states: torch.Tensor, mask: torch.Tensor):
bool_mask = mask.bool()
valid_lengths = bool_mask.sum(dim=1)
selected = hidden_states[bool_mask]
split_result = torch.split(selected, valid_lengths.tolist(), dim=0)
return split_result
def qwen_image_preprocess_text(prompt):
prompt_template_encode = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
template = prompt_template_encode
txt = template.format(prompt)
return txt
def qwen_image_postprocess_text(outputs, _text_inputs, drop_idx=34):
# squeeze the batch dim
hidden_states = outputs.hidden_states[-1]
split_hidden_states = _extract_masked_hidden(
hidden_states, _text_inputs.attention_mask
)
split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
max_seq_len = max([e.size(0) for e in split_hidden_states])
prompt_embeds = torch.stack(
[
torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))])
for u in split_hidden_states
]
)
return prompt_embeds
# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._pack_latents
def _pack_latents(latents, batch_size, num_channels_latents, height, width):
latents = latents.view(
batch_size, num_channels_latents, height // 2, 2, width // 2, 2
)
latents = latents.permute(0, 2, 4, 1, 3, 5)
latents = latents.reshape(
batch_size, (height // 2) * (width // 2), num_channels_latents * 4
)
return latents
@dataclass
class QwenImagePipelineConfig(PipelineConfig):
should_use_guidance: bool = False
is_image_gen: bool = True
vae_tiling: bool = False
vae_sp: bool = False
dit_config: DiTConfig = field(default_factory=QwenImageDitConfig)
# VAE
vae_config: VAEConfig = field(default_factory=QwenImageVAEConfig)
# Text encoding stage
text_encoder_configs: tuple[EncoderConfig, ...] = field(
default_factory=lambda: (Qwen2_5VLConfig(),)
)
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))
preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
default_factory=lambda: (qwen_image_preprocess_text,)
)
postprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
default_factory=lambda: (qwen_image_postprocess_text,)
)
text_encoder_extra_args: list[dict] = field(
default_factory=lambda: [
dict(
padding=True,
truncation=True,
),
None,
]
)
def get_vae_scale_factor(self):
return self.vae_config.arch_config.vae_scale_factor
def prepare_latent_shape(self, batch, batch_size, num_frames):
height = 2 * (
batch.height // (self.vae_config.arch_config.vae_scale_factor * 2)
)
width = 2 * (batch.width // (self.vae_config.arch_config.vae_scale_factor * 2))
num_channels_latents = self.dit_config.arch_config.in_channels // 4
shape = (batch_size, num_channels_latents, height, width)
return shape
def pack_latents(self, latents, batch_size, batch):
height = 2 * (
batch.height // (self.vae_config.arch_config.vae_scale_factor * 2)
)
width = 2 * (batch.width // (self.vae_config.arch_config.vae_scale_factor * 2))
num_channels_latents = self.dit_config.arch_config.in_channels // 4
# pack latents
# _pack_latents(latents, batch_size, num_channels_latents, height, width)
latents = latents.view(
batch_size, num_channels_latents, height // 2, 2, width // 2, 2
)
latents = latents.permute(0, 2, 4, 1, 3, 5)
latents = latents.reshape(
batch_size, (height // 2) * (width // 2), num_channels_latents * 4
)
return latents
@staticmethod
def get_freqs_cis(img_shapes, txt_seq_lens, rotary_emb, device, dtype):
img_freqs, txt_freqs = rotary_emb(img_shapes, txt_seq_lens, device=device)
img_cos, img_sin = (
img_freqs.real.to(dtype=dtype),
img_freqs.imag.to(dtype=dtype),
)
txt_cos, txt_sin = (
txt_freqs.real.to(dtype=dtype),
txt_freqs.imag.to(dtype=dtype),
)
return (img_cos, img_sin), (txt_cos, txt_sin)
def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
batch_size = batch.latents.shape[0]
vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
img_shapes = [
[
(
1,
batch.height // vae_scale_factor // 2,
batch.width // vae_scale_factor // 2,
)
]
] * batch_size
txt_seq_lens = [batch.prompt_embeds[0].shape[1]]
return {
"img_shapes": img_shapes,
"txt_seq_lens": txt_seq_lens,
"freqs_cis": QwenImagePipelineConfig.get_freqs_cis(
img_shapes, txt_seq_lens, rotary_emb, device, dtype
),
}
def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
batch_size = batch.latents.shape[0]
vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
img_shapes = [
[
(
1,
batch.height // vae_scale_factor // 2,
batch.width // vae_scale_factor // 2,
)
]
] * batch_size
txt_seq_lens = [batch.negative_prompt_embeds[0].shape[1]]
return {
"img_shapes": img_shapes,
"txt_seq_lens": txt_seq_lens,
"freqs_cis": QwenImagePipelineConfig.get_freqs_cis(
img_shapes, txt_seq_lens, rotary_emb, device, dtype
),
}
def post_denoising_loop(self, latents, batch):
# VAE applies 8x compression on images but we must also account for packing which requires
# latent height and width to be divisible by 2.
batch_size = latents.shape[0]
channels = latents.shape[-1]
vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
height = 2 * (int(batch.height) // (vae_scale_factor * 2))
width = 2 * (int(batch.width) // (vae_scale_factor * 2))
latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
latents = latents.permute(0, 3, 1, 4, 2, 5)
latents = latents.reshape(batch_size, channels // (2 * 2), 1, height, width)
return latents
class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
ti2i_task = True
def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
# TODO: lots of duplications here
batch_size = batch.latents.shape[0]
height = batch.height
width = batch.width
image = batch.pil_image
image_size = image[0].size if isinstance(image, list) else image.size
calculated_width, calculated_height, _ = calculate_dimensions(
1024 * 1024, image_size[0] / image_size[1]
)
vae_scale_factor = self.get_vae_scale_factor()
img_shapes = [
[
(1, height // vae_scale_factor // 2, width // vae_scale_factor // 2),
(
1,
calculated_height // vae_scale_factor // 2,
calculated_width // vae_scale_factor // 2,
),
]
] * batch_size
txt_seq_lens = [batch.prompt_embeds[0].shape[1]]
return {
"img_shapes": img_shapes,
"txt_seq_lens": txt_seq_lens,
"freqs_cis": QwenImagePipelineConfig.get_freqs_cis(
img_shapes, txt_seq_lens, rotary_emb, device, dtype
),
}
def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
batch_size = batch.latents.shape[0]
height = batch.height
width = batch.width
image = batch.pil_image
image_size = image[0].size if isinstance(image, list) else image.size
calculated_width, calculated_height, _ = calculate_dimensions(
1024 * 1024, image_size[0] / image_size[1]
)
vae_scale_factor = self.get_vae_scale_factor()
img_shapes = [
[
(1, height // vae_scale_factor // 2, width // vae_scale_factor // 2),
(
1,
calculated_height // vae_scale_factor // 2,
calculated_width // vae_scale_factor // 2,
),
]
] * batch_size
txt_seq_lens = [batch.negative_prompt_embeds[0].shape[1]]
return {
"img_shapes": img_shapes,
"txt_seq_lens": txt_seq_lens,
"freqs_cis": QwenImagePipelineConfig.get_freqs_cis(
img_shapes, txt_seq_lens, rotary_emb, device, dtype
),
}
def prepare_latent_shape(self, batch, batch_size, num_frames):
vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
height = 2 * (batch.height // (vae_scale_factor * 2))
width = 2 * (batch.width // (vae_scale_factor * 2))
num_channels_latents = self.dit_config.arch_config.in_channels // 4
shape = (batch_size, 1, num_channels_latents, height, width)
return shape
def preprocess_image(self, image, image_processor):
image_size = image[0].size if isinstance(image, list) else image.size
calculated_width, calculated_height, _ = calculate_dimensions(
1024 * 1024, image_size[0] / image_size[1]
)
image = image_processor.resize(image, calculated_height, calculated_width)
return image
def set_width_and_height(self, width, height, image):
image_size = image[0].size if isinstance(image, list) else image.size
calculated_width, calculated_height, _ = calculate_dimensions(
1024 * 1024, image_size[0] / image_size[1]
)
height = height or calculated_height
width = width or calculated_width
multiple_of = self.get_vae_scale_factor() * 2
width = width // multiple_of * multiple_of
height = height // multiple_of * multiple_of
return width, height
def slice_noise_pred(self, noise, latents):
noise = noise[:, : latents.size(1)]
return noise
@@ -0,0 +1,168 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
"""Registry for pipeline weight-specific configurations."""
import os
from collections.abc import Callable
from sglang.multimodal_gen.configs.pipelines.base import PipelineConfig
from sglang.multimodal_gen.configs.pipelines.flux import FluxPipelineConfig
from sglang.multimodal_gen.configs.pipelines.hunyuan import (
FastHunyuanConfig,
HunyuanConfig,
)
from sglang.multimodal_gen.configs.pipelines.qwen_image import (
QwenImageEditPipelineConfig,
QwenImagePipelineConfig,
)
from sglang.multimodal_gen.configs.pipelines.stepvideo import StepVideoT2VConfig
# isort: off
from sglang.multimodal_gen.configs.pipelines.wan import (
FastWan2_1_T2V_480P_Config,
FastWan2_2_TI2V_5B_Config,
Wan2_2_I2V_A14B_Config,
Wan2_2_T2V_A14B_Config,
Wan2_2_TI2V_5B_Config,
WanI2V480PConfig,
WanI2V720PConfig,
WanT2V480PConfig,
WanT2V720PConfig,
SelfForcingWanT2V480PConfig,
)
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
verify_model_config_and_directory,
maybe_download_model_index,
)
# isort: on
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
# Registry maps specific model weights to their config classes
PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
"FastVideo/FastHunyuan-diffusers": FastHunyuanConfig,
"hunyuanvideo-community/HunyuanVideo": HunyuanConfig,
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig,
"weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": WanI2V480PConfig,
"Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig,
"Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V720PConfig,
"Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V720PConfig,
"FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWan2_1_T2V_480P_Config,
"FastVideo/FastWan2.1-T2V-14B-480P-Diffusers": FastWan2_1_T2V_480P_Config,
"FastVideo/FastWan2.2-TI2V-5B-Diffusers": FastWan2_2_TI2V_5B_Config,
"FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
"FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig,
"wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig,
"Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_Config,
"Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_Config,
"Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_Config,
# Add other specific weight variants
"black-forest-labs/FLUX.1-dev": FluxPipelineConfig,
"Qwen/Qwen-Image": QwenImagePipelineConfig,
"Qwen/Qwen-Image-Edit": QwenImageEditPipelineConfig,
}
# For determining pipeline type from model ID
PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
"hunyuan": lambda id: "hunyuan" in id.lower(),
"wanpipeline": lambda id: "wanpipeline" in id.lower(),
"wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
"wandmdpipeline": lambda id: "wandmdpipeline" in id.lower(),
"wancausaldmdpipeline": lambda id: "wancausaldmdpipeline" in id.lower(),
"stepvideo": lambda id: "stepvideo" in id.lower(),
"qwenimage": lambda id: "qwen-image" in id.lower() and "edit" not in id.lower(),
"qwenimageedit": lambda id: "qwen-image-edit" in id.lower(),
# Add other pipeline architecture detectors
}
# Fallback configs when exact match isn't found but architecture is detected
PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
"hunyuan": HunyuanConfig, # Base Hunyuan config as fallback for any Hunyuan variant
"wanpipeline": WanT2V480PConfig, # Base Wan config as fallback for any Wan variant
"wanimagetovideo": WanI2V480PConfig,
"wandmdpipeline": FastWan2_1_T2V_480P_Config,
"wancausaldmdpipeline": SelfForcingWanT2V480PConfig,
"stepvideo": StepVideoT2VConfig,
"qwenimage": QwenImagePipelineConfig,
"qwenimageedit": QwenImageEditPipelineConfig,
# Other fallbacks by architecture
}
def get_pipeline_config_cls_from_name(
pipeline_name_or_path: str,
) -> type[PipelineConfig]:
"""Get the appropriate configuration class for a given pipeline name or path.
This function implements a multi-step lookup process to find the most suitable
configuration class for a given pipeline. It follows this order:
1. Exact match in the PIPE_NAME_TO_CONFIG
2. Partial match in the PIPE_NAME_TO_CONFIG
3. Fallback to class name in the model_index.json
4. else raise an error
Args:
pipeline_name_or_path (str): The name or path of the pipeline. This can be:
- A registered model ID (e.g., "FastVideo/FastHunyuan-diffusers")
- A local path to a model directory
- A model ID that will be downloaded
Returns:
Type[PipelineConfig]: The configuration class that best matches the pipeline.
This will be one of:
- A specific weight configuration class if an exact match is found
- A fallback configuration class based on the pipeline architecture
- The base PipelineConfig class if no matches are found
Note:
- For local paths, the function will verify the model configuration
- For remote models, it will attempt to download the model index
- Warning messages are logged when falling back to less specific configurations
"""
pipeline_config_cls: type[PipelineConfig] | None = None
# First try exact match for specific weights
if pipeline_name_or_path in PIPE_NAME_TO_CONFIG:
pipeline_config_cls = PIPE_NAME_TO_CONFIG[pipeline_name_or_path]
if pipeline_config_cls is None:
# Try partial matches (for local paths that might include the weight ID)
for registered_id, config_class in PIPE_NAME_TO_CONFIG.items():
if registered_id in pipeline_name_or_path:
pipeline_config_cls = config_class
break
# If no match, try to use the fallback config
if pipeline_config_cls is None:
if os.path.exists(pipeline_name_or_path):
config = verify_model_config_and_directory(pipeline_name_or_path)
else:
config = maybe_download_model_index(pipeline_name_or_path)
logger.warning(
"Trying to use the config from the model_index.json. sgl-diffusion may not correctly identify the optimal config for this model in this situation."
)
pipeline_name = config["_class_name"]
# Try to determine pipeline architecture for fallback
for pipeline_type, detector in PIPELINE_DETECTOR.items():
if detector(pipeline_name.lower()):
pipeline_config_cls = PIPELINE_FALLBACK_CONFIG.get(pipeline_type)
break
if pipeline_config_cls is not None:
logger.warning(
"No match found for pipeline %s, using fallback config %s.",
pipeline_name_or_path,
pipeline_config_cls,
)
if pipeline_config_cls is None:
raise ValueError(
f"No match found for pipeline {pipeline_name_or_path}, please check the pipeline name or path."
)
return pipeline_config_cls
@@ -0,0 +1,36 @@
# 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.pipelines.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"
@@ -0,0 +1,190 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from collections.abc import Callable
from dataclasses import dataclass, field
import torch
from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig
from sglang.multimodal_gen.configs.models.dits import WanVideoConfig
from sglang.multimodal_gen.configs.models.encoders import (
BaseEncoderOutput,
CLIPVisionConfig,
T5Config,
)
from sglang.multimodal_gen.configs.models.vaes import WanVAEConfig
from sglang.multimodal_gen.configs.pipelines.base import PipelineConfig
def t5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor:
mask: torch.Tensor = outputs.attention_mask
hidden_state: torch.Tensor = outputs.last_hidden_state
seq_lens = mask.gt(0).sum(dim=1).long()
assert torch.isnan(hidden_state).sum() == 0
prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)]
prompt_embeds_tensor: torch.Tensor = torch.stack(
[
torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))])
for u in prompt_embeds
],
dim=0,
)
return prompt_embeds_tensor
@dataclass
class WanT2V480PConfig(PipelineConfig):
"""Base configuration for Wan T2V 1.3B pipeline architecture."""
# WanConfig-specific parameters with defaults
# DiT
dit_config: DiTConfig = field(default_factory=WanVideoConfig)
# VAE
vae_config: VAEConfig = field(default_factory=WanVAEConfig)
vae_tiling: bool = False
vae_sp: bool = False
# Denoising stage
flow_shift: float | None = 3.0
# Text encoding stage
text_encoder_configs: tuple[EncoderConfig, ...] = field(
default_factory=lambda: (T5Config(),)
)
postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = (
field(default_factory=lambda: (t5_postprocess_text,))
)
# Precision for each component
precision: str = "bf16"
vae_precision: str = "fp32"
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32",))
# WanConfig-specific added parameters
def __post_init__(self):
self.vae_config.load_encoder = False
self.vae_config.load_decoder = True
@dataclass
class WanT2V720PConfig(WanT2V480PConfig):
"""Base configuration for Wan T2V 14B 720P pipeline architecture."""
# WanConfig-specific parameters with defaults
# Denoising stage
flow_shift: float | None = 5.0
@dataclass
class WanI2V480PConfig(WanT2V480PConfig):
"""Base configuration for Wan I2V 14B 480P pipeline architecture."""
# WanConfig-specific parameters with defaults
i2v_task: bool = True
# Precision for each component
image_encoder_config: EncoderConfig = field(default_factory=CLIPVisionConfig)
image_encoder_precision: str = "fp32"
image_encoder_extra_args: dict = field(
default_factory=lambda: dict(
output_hidden_states=True,
)
)
def postprocess_image(self, image):
return image.hidden_states[-2]
def __post_init__(self) -> None:
self.vae_config.load_encoder = True
self.vae_config.load_decoder = True
@dataclass
class WanI2V720PConfig(WanI2V480PConfig):
"""Base configuration for Wan I2V 14B 720P pipeline architecture."""
# WanConfig-specific parameters with defaults
# Denoising stage
flow_shift: float | None = 5.0
@dataclass
class FastWan2_1_T2V_480P_Config(WanT2V480PConfig):
"""Base configuration for FastWan T2V 1.3B 480P pipeline architecture with DMD"""
# WanConfig-specific parameters with defaults
# Denoising stage
flow_shift: float | None = 8.0
dmd_denoising_steps: list[int] | None = field(
default_factory=lambda: [1000, 757, 522]
)
@dataclass
class Wan2_2_TI2V_5B_Config(WanT2V480PConfig):
flow_shift: float | None = 5.0
ti2v_task: bool = True
expand_timesteps: bool = True
# ti2v, 5B
vae_stride = (4, 16, 16)
def prepare_latent_shape(self, batch, batch_size, num_frames):
F = num_frames
z_dim = self.vae_config.arch_config.z_dim
vae_stride = self.vae_stride
oh = batch.height
ow = batch.width
shape = (z_dim, F, oh // vae_stride[1], ow // vae_stride[2])
return shape
def __post_init__(self) -> None:
self.vae_config.load_encoder = True
self.vae_config.load_decoder = True
self.dit_config.expand_timesteps = self.expand_timesteps
@dataclass
class FastWan2_2_TI2V_5B_Config(Wan2_2_TI2V_5B_Config):
flow_shift: float | None = 5.0
dmd_denoising_steps: list[int] | None = field(
default_factory=lambda: [1000, 757, 522]
)
@dataclass
class Wan2_2_T2V_A14B_Config(WanT2V480PConfig):
flow_shift: float | None = 12.0
boundary_ratio: float | None = 0.875
def __post_init__(self) -> None:
self.dit_config.boundary_ratio = self.boundary_ratio
@dataclass
class Wan2_2_I2V_A14B_Config(WanI2V480PConfig):
flow_shift: float | None = 5.0
boundary_ratio: float | None = 0.900
def __post_init__(self) -> None:
super().__post_init__()
self.dit_config.boundary_ratio = self.boundary_ratio
# =============================================
# ============= Causal Self-Forcing =============
# =============================================
@dataclass
class SelfForcingWanT2V480PConfig(WanT2V480PConfig):
is_causal: bool = True
flow_shift: float | None = 5.0
dmd_denoising_steps: list[int] | None = field(
default_factory=lambda: [1000, 750, 500, 250]
)
warp_denoising_step: bool = True
@@ -0,0 +1,5 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
from sglang.multimodal_gen.configs.sample.base import SamplingParams
__all__ = ["SamplingParams"]
@@ -0,0 +1,494 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
import argparse
import dataclasses
import hashlib
import json
import os.path
import re
import time
import unicodedata
import uuid
from copy import deepcopy
from dataclasses import dataclass
from enum import Enum, auto
from typing import Any
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.utils import align_to
logger = init_logger(__name__)
def _json_safe(obj: Any):
"""
Recursively convert objects to JSON-serializable forms.
- Enums -> their name
- Sets/Tuples -> lists
- Dicts/Lists -> recursively processed
"""
if isinstance(obj, Enum):
return obj.name
if isinstance(obj, dict):
return {k: _json_safe(v) for k, v in obj.items()}
if isinstance(obj, (list, tuple, set)):
return [_json_safe(v) for v in obj]
return obj
def generate_request_id() -> str:
return str(uuid.uuid4())
def _sanitize_filename(name: str, replacement: str = "_", max_length: int = 150) -> str:
"""Create a filesystem- and ffmpeg-friendly filename.
- Normalize to ASCII (drop accents and unsupported chars)
- Replace spaces with underscores
- Replace any char not in [A-Za-z0-9_.-] with replacement
- Collapse multiple underscores
- Trim leading/trailing dots/underscores and limit length
"""
normalized = unicodedata.normalize("NFKD", name)
ascii_name = normalized.encode("ascii", "ignore").decode("ascii")
ascii_name = ascii_name.replace(" ", "_")
ascii_name = re.sub(r"[^A-Za-z0-9._-]", replacement, ascii_name)
ascii_name = re.sub(r"_+", "_", ascii_name).strip("._")
if not ascii_name:
ascii_name = "output"
if max_length and len(ascii_name) > max_length:
ascii_name = ascii_name[:max_length]
return ascii_name
class DataType(Enum):
IMAGE = auto()
VIDEO = auto()
def get_default_extension(self) -> str:
if self == DataType.IMAGE:
return "jpg"
else:
return "mp4"
@dataclass
class SamplingParams:
"""
Sampling parameters for generation.
"""
data_type: DataType = DataType.VIDEO
request_id: str | None = None
# All fields below are copied from ForwardBatch
# Image inputs
image_path: str | None = None
# Text inputs
prompt: str | list[str] | None = None
negative_prompt: str = (
"Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
)
prompt_path: str | None = None
output_path: str = "outputs/"
output_file_name: str | None = None
# Batch info
num_outputs_per_prompt: int = 1
seed: int = 1024
# Original dimensions (before VAE scaling)
num_frames: int = 125
num_frames_round_down: bool = (
False # Whether to round down num_frames if it's not divisible by num_gpus
)
height: int | None = None
width: int | None = None
# NOTE: this is temporary, we need a way to know if width or height is not provided, or do the image resize earlier
height_not_provided: bool = False
width_not_provided: bool = False
fps: int = 24
# Denoising parameters
num_inference_steps: int = 50
guidance_scale: float = 1.0
guidance_rescale: float = 0.0
boundary_ratio: float | None = None
# TeaCache parameters
enable_teacache: bool = False
# Profiling
profile: bool = False
num_profiled_timesteps: int = 2
# Debugging
debug: bool = False
# Misc
save_output: bool = True
return_frames: bool = False
return_trajectory_latents: bool = False # returns all latents for each timestep
return_trajectory_decoded: bool = False # returns decoded latents for each timestep
def set_output_file_ext(self):
# add extension if needed
if not any(
self.output_file_name.endswith(ext)
for ext in [".mp4", ".jpg", ".png", ".webp"]
):
self.output_file_name = (
f"{self.output_file_name}.{self.data_type.get_default_extension()}"
)
def set_output_file_name(self):
# settle output_file_name
if (
self.output_file_name is None
and self.prompt
and isinstance(self.prompt, str)
):
# generate a random filename
# get a hash of current params
params_dict = dataclasses.asdict(self)
# Avoid recursion
params_dict["output_file_name"] = ""
# Convert to a stable JSON string
params_str = json.dumps(_json_safe(params_dict), sort_keys=True)
# Create a hash
hasher = hashlib.sha256()
hasher.update(params_str.encode("utf-8"))
param_hash = hasher.hexdigest()[:8]
timestamp = time.strftime("%Y%m%d-%H%M%S")
base = f"{self.prompt[:100]}_{timestamp}_{param_hash}"
self.output_file_name = base
if self.output_file_name is None:
timestamp = time.strftime("%Y%m%d-%H%M%S")
self.output_file_name = f"output_{timestamp}"
self.output_file_name = _sanitize_filename(self.output_file_name)
# Ensure a proper extension is present
self.set_output_file_ext()
def __post_init__(self) -> None:
assert self.num_frames >= 1
self.data_type = DataType.VIDEO if self.num_frames > 1 else DataType.IMAGE
if self.width is None:
self.width_not_provided = True
self.width = 1280
if self.height is None:
self.height_not_provided = True
self.height = 720
def check_sampling_param(self):
if self.prompt_path and not self.prompt_path.endswith(".txt"):
raise ValueError("prompt_path must be a txt file")
def update(self, source_dict: dict[str, Any]) -> None:
for key, value in source_dict.items():
if hasattr(self, key):
setattr(self, key, value)
else:
logger.exception("%s has no attribute %s", type(self).__name__, key)
self.__post_init__()
@classmethod
def from_pretrained(cls, model_path: str, **kwargs) -> "SamplingParams":
from sglang.multimodal_gen.configs.sample.registry import (
get_sampling_param_cls_for_name,
)
sampling_cls = get_sampling_param_cls_for_name(model_path)
logger.debug(f"Using pretrained SamplingParam: {sampling_cls}")
if sampling_cls is not None:
sampling_params: SamplingParams = sampling_cls(**kwargs)
else:
logger.warning(
"Couldn't find an optimal sampling param for %s. Using the default sampling param.",
model_path,
)
sampling_params = cls(**kwargs)
return sampling_params
def from_user_sampling_params(self, user_params):
sampling_params = deepcopy(self)
sampling_params._merge_with_user_params(user_params)
return sampling_params
@staticmethod
def add_cli_args(parser: Any) -> Any:
"""Add CLI arguments for SamplingParam fields"""
parser.add_argument("--data-type", type=str, nargs="+", default=DataType.VIDEO)
parser.add_argument(
"--num-frames-round-down",
action="store_true",
default=SamplingParams.num_frames_round_down,
)
parser.add_argument(
"--enable-teacache",
action="store_true",
default=SamplingParams.enable_teacache,
)
parser.add_argument(
"--profile",
action="store_true",
default=SamplingParams.profile,
help="Enable torch profiler for denoising stage",
)
parser.add_argument(
"--debug",
action="store_true",
default=SamplingParams.debug,
help="",
)
parser.add_argument(
"--num-profiled-timesteps",
type=int,
default=SamplingParams.num_profiled_timesteps,
help="Number of timesteps to profile after warmup",
)
parser.add_argument(
"--prompt",
type=str,
default=SamplingParams.prompt,
help="Text prompt for generation",
)
parser.add_argument(
"--negative-prompt",
type=str,
default=SamplingParams.negative_prompt,
help="Negative text prompt for generation",
)
parser.add_argument(
"--prompt-path",
type=str,
default=SamplingParams.prompt_path,
help="Path to a text file containing the prompt",
)
parser.add_argument(
"--output-path",
type=str,
default=SamplingParams.output_path,
help="Path to save the generated image/video",
)
parser.add_argument(
"--output-file-name",
type=str,
default=SamplingParams.output_file_name,
help="Name of the output file",
)
parser.add_argument(
"--num-outputs-per-prompt",
type=int,
default=SamplingParams.num_outputs_per_prompt,
help="Number of outputs to generate per prompt",
)
parser.add_argument(
"--seed",
type=int,
default=SamplingParams.seed,
help="Random seed for generation",
)
parser.add_argument(
"--num-frames",
type=int,
default=SamplingParams.num_frames,
help="Number of frames to generate",
)
parser.add_argument(
"--height",
type=int,
default=SamplingParams.height,
help="Height of generated output",
)
parser.add_argument(
"--width",
type=int,
default=SamplingParams.width,
help="Width of generated output",
)
parser.add_argument(
"--fps",
type=int,
default=SamplingParams.fps,
help="Frames per second for saved output",
)
parser.add_argument(
"--num-inference-steps",
type=int,
default=SamplingParams.num_inference_steps,
help="Number of denoising steps",
)
parser.add_argument(
"--guidance-scale",
type=float,
default=SamplingParams.guidance_scale,
help="Classifier-free guidance scale",
)
parser.add_argument(
"--guidance-rescale",
type=float,
default=SamplingParams.guidance_rescale,
help="Guidance rescale factor",
)
parser.add_argument(
"--boundary-ratio",
type=float,
default=SamplingParams.boundary_ratio,
help="Boundary timestep ratio",
)
parser.add_argument(
"--save-output",
action="store_true",
default=SamplingParams.save_output,
help="Whether to save the output to disk",
)
parser.add_argument(
"--no-save-output",
action="store_false",
dest="save_output",
help="Don't save the output to disk",
)
parser.add_argument(
"--return-frames",
action="store_true",
default=SamplingParams.return_frames,
help="Whether to return the raw frames",
)
parser.add_argument(
"--image-path",
type=str,
default=SamplingParams.image_path,
help="Path to input image for image-to-video generation",
)
parser.add_argument(
"--moba-config-path",
type=str,
default=None,
help="Path to a JSON file containing V-MoBA specific configurations.",
)
parser.add_argument(
"--return-trajectory-latents",
action="store_true",
default=SamplingParams.return_trajectory_latents,
help="Whether to return the trajectory",
)
parser.add_argument(
"--return-trajectory-decoded",
action="store_true",
default=SamplingParams.return_trajectory_decoded,
help="Whether to return the decoded trajectory",
)
return parser
@classmethod
def from_cli_args(cls, args: argparse.Namespace):
attrs = [attr.name for attr in dataclasses.fields(cls)]
args.height_not_provided = False
args.width_not_provided = False
return cls(**{attr: getattr(args, attr) for attr in attrs})
def output_file_path(self):
return os.path.join(self.output_path, self.output_file_name)
def _merge_with_user_params(self, user_params):
"""
Merges parameters from a user-provided SamplingParams object.
This method updates the current object with values from `user_params`,
but skips any fields that are explicitly defined in the current object's
subclass. This is to preserve model-specific optimal parameters.
It also skips fields that the user has not changed from the default
in `user_params`.
"""
if user_params is None:
return
# Get fields defined directly in the subclass (not inherited)
subclass_defined_fields = set(type(self).__annotations__.keys())
# Compare against current instance to avoid constructing a default instance
default_params = SamplingParams()
for field in dataclasses.fields(user_params):
field_name = field.name
user_value = getattr(user_params, field_name)
default_value = getattr(default_params, field_name)
# A field is considered user-modified if its value is different from
# the default, with an exception for `output_file_name` which is
# auto-generated with a random component.
is_user_modified = (
user_value != default_value
if field_name != "output_file_name"
else user_params.output_file_path is not None
)
if is_user_modified and field_name not in subclass_defined_fields:
if hasattr(self, field_name):
setattr(self, field_name, user_value)
self.__post_init__()
@property
def n_tokens(self) -> int:
# Calculate latent sizes
if self.height and self.width:
latents_size = [
(self.num_frames - 1) // 4 + 1,
self.height // 8,
self.width // 8,
]
n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
else:
n_tokens = -1
return n_tokens
def output_file_path(self):
return os.path.join(self.output_path, self.output_file_name)
def log(self, server_args: ServerArgs):
# TODO: in some cases (e.g., TI2I), height and weight might be undecided at this moment
if self.height:
target_height = align_to(self.height, 16)
else:
target_height = -1
if self.width:
target_width = align_to(self.width, 16)
else:
target_width = -1
# Log sampling parameters
debug_str = f"""Sampling params:
height: {target_height}
width: {target_width}
num_frames: {self.num_frames}
prompt: {self.prompt}
neg_prompt: {self.negative_prompt}
seed: {self.seed}
infer_steps: {self.num_inference_steps}
num_outputs_per_prompt: {self.num_outputs_per_prompt}
guidance_scale: {self.guidance_scale}
embedded_guidance_scale: {server_args.pipeline_config.embedded_cfg_scale}
n_tokens: {self.n_tokens}
flow_shift: {server_args.pipeline_config.flow_shift}
image_path: {self.image_path}
save_output: {self.save_output}
output_file_path: {self.output_file_path()}
""" # type: ignore[attr-defined]
logger.info(debug_str)
@dataclass
class CacheParams:
cache_type: str = "none"
@@ -0,0 +1,18 @@
# 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.base import SamplingParams
@dataclass
class FluxSamplingParams(SamplingParams):
# Video parameters
# height: int = 1024
# width: int = 1024
num_frames: int = 1
# Denoising stage
guidance_scale: float = 1.0
negative_prompt: str = None
num_inference_steps: int = 50
@@ -0,0 +1,37 @@
# 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.sample.base import SamplingParams
from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams
@dataclass
class HunyuanSamplingParams(SamplingParams):
num_inference_steps: int = 50
num_frames: int = 125
height: int = 720
width: int = 1280
fps: int = 24
guidance_scale: float = 1.0
teacache_params: TeaCacheParams = field(
default_factory=lambda: TeaCacheParams(
teacache_thresh=0.15,
coefficients=[
7.33226126e02,
-4.01131952e02,
6.75869174e01,
-3.14987800e00,
9.61237896e-02,
],
)
)
@dataclass
class FastHunyuanSamplingParam(HunyuanSamplingParams):
num_inference_steps: int = 6
@@ -0,0 +1,18 @@
# 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.base import SamplingParams
@dataclass
class QwenImageSamplingParams(SamplingParams):
# Video parameters
# height: int = 1024
# width: int = 1024
negative_prompt: str = " "
num_frames: int = 1
# Denoising stage
guidance_scale: float = 4.0
num_inference_steps: int = 50
@@ -0,0 +1,122 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
import os
from collections.abc import Callable
from typing import Any
from sglang.multimodal_gen.configs.sample.flux import FluxSamplingParams
from sglang.multimodal_gen.configs.sample.hunyuan import (
FastHunyuanSamplingParam,
HunyuanSamplingParams,
)
from sglang.multimodal_gen.configs.sample.qwenimage import QwenImageSamplingParams
from sglang.multimodal_gen.configs.sample.stepvideo import StepVideoT2VSamplingParams
# isort: off
from sglang.multimodal_gen.configs.sample.wan import (
FastWanT2V480PConfig,
Wan2_1_Fun_1_3B_InP_SamplingParams,
Wan2_2_I2V_A14B_SamplingParam,
Wan2_2_T2V_A14B_SamplingParam,
Wan2_2_TI2V_5B_SamplingParam,
WanI2V_14B_480P_SamplingParam,
WanI2V_14B_720P_SamplingParam,
WanT2V_1_3B_SamplingParams,
WanT2V_14B_SamplingParams,
SelfForcingWanT2V480PConfig,
)
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
maybe_download_model_index,
verify_model_config_and_directory,
)
# isort: on
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
# Registry maps specific model weights to their config classes
SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
"FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam,
"hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParams,
"FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParams,
# Wan2.1
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V_1_3B_SamplingParams,
"Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParams,
"Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam,
"Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam,
"weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": Wan2_1_Fun_1_3B_InP_SamplingParams,
# Wan2.2
"Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
"FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
"Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam,
"Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam,
# FastWan2.1
"FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig,
# FastWan2.2
"FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
# Causal Self-Forcing Wan2.1
"wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig,
# Add other specific weight variants
"black-forest-labs/FLUX.1-dev": FluxSamplingParams,
"Qwen/Qwen-Image": QwenImageSamplingParams,
"Qwen/Qwen-Image-Edit": QwenImageSamplingParams,
}
# For determining pipeline type from model ID
SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = {
"hunyuan": lambda id: "hunyuan" in id.lower(),
"wanpipeline": lambda id: "wanpipeline" in id.lower(),
"wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
"stepvideo": lambda id: "stepvideo" in id.lower(),
# Add other pipeline architecture detectors
"flux": lambda id: "flux" in id.lower(),
}
# Fallback configs when exact match isn't found but architecture is detected
SAMPLING_FALLBACK_PARAM: dict[str, Any] = {
"hunyuan": HunyuanSamplingParams, # Base Hunyuan config as fallback for any Hunyuan variant
"wanpipeline": WanT2V_1_3B_SamplingParams, # Base Wan config as fallback for any Wan variant
"wanimagetovideo": WanI2V_14B_480P_SamplingParam,
"stepvideo": StepVideoT2VSamplingParams,
# Other fallbacks by architecture
"flux": FluxSamplingParams,
}
def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None:
"""Get the appropriate sampling param for specific pretrained weights."""
if os.path.exists(pipeline_name_or_path):
config = verify_model_config_and_directory(pipeline_name_or_path)
logger.warning(
"sgl-diffusion may not correctly identify the optimal sampling param for this model, as the local directory may have been renamed."
)
else:
config = maybe_download_model_index(pipeline_name_or_path)
pipeline_name = config["_class_name"]
# First try exact match for specific weights
if pipeline_name_or_path in SAMPLING_PARAM_REGISTRY:
return SAMPLING_PARAM_REGISTRY[pipeline_name_or_path]
# Try partial matches (for local paths that might include the weight ID)
for registered_id, config_class in SAMPLING_PARAM_REGISTRY.items():
if registered_id in pipeline_name_or_path:
return config_class
# If no match, try to use the fallback config
fallback_config = None
# Try to determine pipeline architecture for fallback
for pipeline_type, detector in SAMPLING_PARAM_DETECTOR.items():
if detector(pipeline_name.lower()):
fallback_config = SAMPLING_FALLBACK_PARAM.get(pipeline_type)
break
logger.warning(
"No match found for pipeline %s, using fallback sampling param %s.",
pipeline_name_or_path,
fallback_config,
)
return fallback_config
@@ -0,0 +1,22 @@
# 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.base 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 = "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。"
@@ -0,0 +1,43 @@
# 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.sample.base import CacheParams
@dataclass
class TeaCacheParams(CacheParams):
cache_type: str = "teacache"
teacache_thresh: float = 0.0
coefficients: list[float] = field(default_factory=list)
@dataclass
class WanTeaCacheParams(CacheParams):
# Unfortunately, TeaCache is very different for Wan than other models
cache_type: str = "teacache"
teacache_thresh: float = 0.0
use_ret_steps: bool = True
ret_steps_coeffs: list[float] = field(default_factory=list)
non_ret_steps_coeffs: list[float] = field(default_factory=list)
@property
def coefficients(self) -> list[float]:
if self.use_ret_steps:
return self.ret_steps_coeffs
else:
return self.non_ret_steps_coeffs
@property
def ret_steps(self) -> int:
if self.use_ret_steps:
return 5 * 2
else:
return 1 * 2
def get_cutoff_steps(self, num_inference_steps: int) -> int:
if self.use_ret_steps:
return num_inference_steps * 2
else:
return num_inference_steps * 2 - 2
@@ -0,0 +1,217 @@
# 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.sample.base import SamplingParams
from sglang.multimodal_gen.configs.sample.teacache import WanTeaCacheParams
@dataclass
class WanT2V_1_3B_SamplingParams(SamplingParams):
# Video parameters
height: int = 480
width: int = 832
num_frames: int = 81
fps: int = 16
# Denoising stage
guidance_scale: float = 3.0
negative_prompt: str = (
"Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
)
num_inference_steps: int = 50
teacache_params: WanTeaCacheParams = field(
default_factory=lambda: WanTeaCacheParams(
teacache_thresh=0.08,
ret_steps_coeffs=[
-5.21862437e04,
9.23041404e03,
-5.28275948e02,
1.36987616e01,
-4.99875664e-02,
],
non_ret_steps_coeffs=[
2.39676752e03,
-1.31110545e03,
2.01331979e02,
-8.29855975e00,
1.37887774e-01,
],
)
)
@dataclass
class WanT2V_14B_SamplingParams(SamplingParams):
# Video parameters
height: int = 720
width: int = 1280
num_frames: int = 81
fps: int = 16
# Denoising stage
guidance_scale: float = 5.0
negative_prompt: str = (
"Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
)
num_inference_steps: int = 50
teacache_params: WanTeaCacheParams = field(
default_factory=lambda: WanTeaCacheParams(
teacache_thresh=0.20,
use_ret_steps=False,
ret_steps_coeffs=[
-3.03318725e05,
4.90537029e04,
-2.65530556e03,
5.87365115e01,
-3.15583525e-01,
],
non_ret_steps_coeffs=[
-5784.54975374,
5449.50911966,
-1811.16591783,
256.27178429,
-13.02252404,
],
)
)
@dataclass
class WanI2V_14B_480P_SamplingParam(WanT2V_1_3B_SamplingParams):
# Denoising stage
guidance_scale: float = 5.0
num_inference_steps: int = 50
# num_inference_steps: int = 40
teacache_params: WanTeaCacheParams = field(
default_factory=lambda: WanTeaCacheParams(
teacache_thresh=0.26,
ret_steps_coeffs=[
-3.03318725e05,
4.90537029e04,
-2.65530556e03,
5.87365115e01,
-3.15583525e-01,
],
non_ret_steps_coeffs=[
-5784.54975374,
5449.50911966,
-1811.16591783,
256.27178429,
-13.02252404,
],
)
)
@dataclass
class WanI2V_14B_720P_SamplingParam(WanT2V_14B_SamplingParams):
# Denoising stage
guidance_scale: float = 5.0
num_inference_steps: int = 50
# num_inference_steps: int = 40
teacache_params: WanTeaCacheParams = field(
default_factory=lambda: WanTeaCacheParams(
teacache_thresh=0.3,
ret_steps_coeffs=[
-3.03318725e05,
4.90537029e04,
-2.65530556e03,
5.87365115e01,
-3.15583525e-01,
],
non_ret_steps_coeffs=[
-5784.54975374,
5449.50911966,
-1811.16591783,
256.27178429,
-13.02252404,
],
)
)
@dataclass
class FastWanT2V480PConfig(WanT2V_1_3B_SamplingParams):
# DMD parameters
# dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 757, 522])
num_inference_steps: int = 3
num_frames: int = 61
height: int = 448
width: int = 832
fps: int = 16
# =============================================
# ============= Wan2.1 Fun Models =============
# =============================================
@dataclass
class Wan2_1_Fun_1_3B_InP_SamplingParams(SamplingParams):
"""Sampling parameters for Wan2.1 Fun 1.3B InP model."""
height: int = 480
width: int = 832
num_frames: int = 81
fps: int = 16
negative_prompt: str | None = (
"色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走"
)
guidance_scale: float = 6.0
num_inference_steps: int = 50
# =============================================
# ============= Wan2.2 TI2V Models =============
# =============================================
@dataclass
class Wan2_2_Base_SamplingParams(SamplingParams):
"""Sampling parameters for Wan2.2 TI2V 5B model."""
negative_prompt: str | None = (
"色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走"
)
@dataclass
class Wan2_2_TI2V_5B_SamplingParam(Wan2_2_Base_SamplingParams):
"""Sampling parameters for Wan2.2 TI2V 5B model."""
height: int = 704
width: int = 1280
num_frames: int = 121
fps: int = 24
guidance_scale: float = 5.0
num_inference_steps: int = 50
@dataclass
class Wan2_2_T2V_A14B_SamplingParam(Wan2_2_Base_SamplingParams):
guidance_scale: float = 4.0 # high_noise
guidance_scale_2: float = 3.0 # low_noise
num_inference_steps: int = 40
fps: int = 16
# NOTE(will): default boundary timestep is tracked by PipelineConfig, but
# can be overridden during sampling
@dataclass
class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParams):
guidance_scale: float = 3.5 # high_noise
guidance_scale_2: float = 3.5 # low_noise
num_inference_steps: int = 40
fps: int = 16
# NOTE(will): default boundary timestep is tracked by PipelineConfig, but
# can be overridden during sampling
# =============================================
# ============= Causal Self-Forcing =============
# =============================================
@dataclass
class SelfForcingWanT2V480PConfig(WanT2V_1_3B_SamplingParams):
pass
@@ -0,0 +1,61 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import argparse
from typing import Any
def update_config_from_args(
config: Any, args_dict: dict[str, Any], prefix: str = "", pop_args: bool = False
) -> bool:
"""
Update configuration object from arguments dictionary.
Args:
config: The configuration object to update
args_dict: Dictionary containing arguments
prefix: Prefix for the configuration parameters in the args_dict.
If None, assumes direct attribute mapping without prefix.
"""
# Handle top-level attributes (no prefix)
args_not_to_remove = [
"model_path",
]
args_to_remove = []
if prefix.strip() == "":
for key, value in args_dict.items():
if hasattr(config, key) and value is not None:
if key == "text_encoder_precisions" and isinstance(value, list):
setattr(config, key, tuple(value))
else:
setattr(config, key, value)
if pop_args:
args_to_remove.append(key)
else:
# Handle nested attributes with prefix
prefix_with_dot = f"{prefix}."
for key, value in args_dict.items():
if key.startswith(prefix_with_dot) and value is not None:
attr_name = key[len(prefix_with_dot) :]
if hasattr(config, attr_name):
setattr(config, attr_name, value)
if pop_args:
args_to_remove.append(key)
if pop_args:
for key in args_to_remove:
if key not in args_not_to_remove:
args_dict.pop(key)
return len(args_to_remove) > 0
def clean_cli_args(args: argparse.Namespace) -> dict[str, Any]:
"""
Clean the arguments by removing the ones that not explicitly provided by the user.
"""
provided_args = {}
for k, v in vars(args).items():
if v is not None and hasattr(args, "_provided") and k in args._provided:
provided_args[k] = v
return provided_args
@@ -0,0 +1,41 @@
{
"embedded_cfg_scale": 6.0,
"flow_shift": 3,
"dit_cpu_offload": true,
"disable_autocast": false,
"precision": "bf16",
"vae_precision": "fp32",
"vae_tiling": false,
"vae_sp": false,
"vae_config": {
"load_encoder": false,
"load_decoder": true,
"tile_sample_min_height": 256,
"tile_sample_min_width": 256,
"tile_sample_min_num_frames": 16,
"tile_sample_stride_height": 192,
"tile_sample_stride_width": 192,
"tile_sample_stride_num_frames": 12,
"blend_num_frames": 8,
"use_tiling": false,
"use_temporal_tiling": false,
"use_parallel_tiling": false,
"use_feature_cache": true
},
"dit_config": {
"prefix": "Wan",
"quant_config": null
},
"text_encoder_precisions": [
"fp32"
],
"text_encoder_configs": [
{
"prefix": "t5",
"quant_config": null,
"lora_config": null
}
],
"mask_strategy_file_path": null,
"enable_torch_compile": false
}
@@ -0,0 +1,49 @@
{
"embedded_cfg_scale": 6.0,
"flow_shift": 3,
"dit_cpu_offload": true,
"disable_autocast": false,
"precision": "bf16",
"vae_precision": "fp32",
"vae_tiling": false,
"vae_sp": false,
"vae_config": {
"load_encoder": true,
"load_decoder": true,
"tile_sample_min_height": 256,
"tile_sample_min_width": 256,
"tile_sample_min_num_frames": 16,
"tile_sample_stride_height": 192,
"tile_sample_stride_width": 192,
"tile_sample_stride_num_frames": 12,
"blend_num_frames": 8,
"use_tiling": false,
"use_temporal_tiling": false,
"use_parallel_tiling": false,
"use_feature_cache": true
},
"dit_config": {
"prefix": "Wan",
"quant_config": null
},
"text_encoder_precisions": [
"fp32"
],
"text_encoder_configs": [
{
"prefix": "t5",
"quant_config": null,
"lora_config": null
}
],
"mask_strategy_file_path": null,
"enable_torch_compile": false,
"image_encoder_config": {
"prefix": "clip",
"quant_config": null,
"lora_config": null,
"num_hidden_layers_override": null,
"require_post_norm": null
},
"image_encoder_precision": "fp32"
}