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:
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
@@ -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"
|
||||
}
|
||||
Reference in New Issue
Block a user