WIP: initial multimodal-gen support (#12484)
Co-authored-by: yhyang201 <yhyang201@gmail.com> Co-authored-by: yizhang2077 <1109276519@qq.com> Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com> Co-authored-by: ispobock <ispobaoke@gmail.com> Co-authored-by: JiLi <leege233@gmail.com> Co-authored-by: CHEN Xi <78632976+RubiaCx@users.noreply.github.com> Co-authored-by: laixin <xielx@shanghaitech.edu.cn> Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com> Co-authored-by: jzhang38 <a1286225768@gmail.com> Co-authored-by: BrianChen1129 <yongqichcd@gmail.com> Co-authored-by: Kevin Lin <42618777+kevin314@users.noreply.github.com> Co-authored-by: Edenzzzz <wtan45@wisc.edu> Co-authored-by: rlsu9 <r3su@ucsd.edu> Co-authored-by: Jinzhe Pan <48981407+eigensystem@users.noreply.github.com> Co-authored-by: foreverpiano <pianoqwz@qq.com> Co-authored-by: RandNMR73 <notomatthew31@gmail.com> Co-authored-by: PorridgeSwim <yz3883@columbia.edu> Co-authored-by: Jiali Chen <90408393+gary-chenjl@users.noreply.github.com>
This commit is contained in:
Mick
@@ -79,6 +79,25 @@ dependencies = [
|
||||
|
||||
[project.optional-dependencies]
|
||||
checkpoint-engine = ["checkpoint-engine==0.1.2"]
|
||||
diffusion = [
|
||||
"diffusers==0.35.2",
|
||||
"yunchang==0.6.3.post1",
|
||||
"opencv-python==4.10.0.84",
|
||||
"imageio==2.36.0",
|
||||
"imageio-ffmpeg==0.5.1",
|
||||
"PyYAML==6.0.1",
|
||||
"moviepy>=2.0.0",
|
||||
"cloudpickle",
|
||||
"remote-pdb",
|
||||
"torchcodec==0.5.0",
|
||||
"st_attn ==0.0.7",
|
||||
"vsa==0.0.4",
|
||||
]
|
||||
|
||||
[tool.uv.extra-build-dependencies]
|
||||
st-attn = ["torch", "setuptools"]
|
||||
vsa = ["torch", "setuptools"]
|
||||
|
||||
test = [
|
||||
"accelerate",
|
||||
"expecttest",
|
||||
@@ -102,6 +121,9 @@ tracing = [
|
||||
"Homepage" = "https://github.com/sgl-project/sglang"
|
||||
"Bug Tracker" = "https://github.com/sgl-project/sglang/issues"
|
||||
|
||||
[project.scripts]
|
||||
sglang = "sglang.cli.main:main"
|
||||
|
||||
[tool.setuptools.package-data]
|
||||
"sglang" = [
|
||||
"srt/layers/moe/fused_moe_triton/configs/*/*.json",
|
||||
|
||||
0
python/sglang/cli/__init__.py
Normal file
0
python/sglang/cli/__init__.py
Normal file
21
python/sglang/cli/generate.py
Normal file
21
python/sglang/cli/generate.py
Normal file
@@ -0,0 +1,21 @@
|
||||
import argparse
|
||||
|
||||
from sglang.cli.main import get_is_diffusion_model, get_model_path
|
||||
from sglang.multimodal_gen.runtime.entrypoints.cli.generate import (
|
||||
add_multimodal_gen_generate_args,
|
||||
generate_cmd,
|
||||
)
|
||||
|
||||
|
||||
def generate(args, extra_argv):
|
||||
model_path = get_model_path(extra_argv)
|
||||
is_diffusion_model = get_is_diffusion_model(model_path)
|
||||
if is_diffusion_model:
|
||||
parser = argparse.ArgumentParser(description="SGLang Multimodal Generation")
|
||||
add_multimodal_gen_generate_args(parser)
|
||||
parsed_args = parser.parse_args(extra_argv)
|
||||
generate_cmd(parsed_args)
|
||||
else:
|
||||
raise Exception(
|
||||
f"Generate subcommand is not yet supported for model: {model_path}"
|
||||
)
|
||||
178
python/sglang/cli/main.py
Normal file
178
python/sglang/cli/main.py
Normal file
@@ -0,0 +1,178 @@
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
from typing import Optional
|
||||
|
||||
import filelock
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
from sglang.cli.generate import generate
|
||||
from sglang.cli.serve import serve
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
temp_dir = tempfile.gettempdir()
|
||||
|
||||
|
||||
def _get_lock(model_name_or_path: str, cache_dir: Optional[str] = None):
|
||||
lock_dir = cache_dir or temp_dir
|
||||
os.makedirs(os.path.dirname(lock_dir), exist_ok=True)
|
||||
model_name = model_name_or_path.replace("/", "-")
|
||||
hash_name = hashlib.sha256(model_name.encode()).hexdigest()
|
||||
lock_file_name = hash_name + model_name + ".lock"
|
||||
lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666)
|
||||
return lock
|
||||
|
||||
|
||||
# Copied and adapted from hf_diffusers_utils.py
|
||||
def _maybe_download_model(
|
||||
model_name_or_path: str, local_dir: str | None = None, download: bool = True
|
||||
) -> str:
|
||||
"""
|
||||
Resolve a model path. If it's a local directory, return it.
|
||||
If it's a Hugging Face Hub ID, download only the config file
|
||||
(`model_index.json` or `config.json`) and return its directory.
|
||||
|
||||
Args:
|
||||
model_name_or_path: Local path or Hugging Face Hub model ID
|
||||
local_dir: Local directory to save the downloaded file (if any)
|
||||
download: Whether to download from Hugging Face Hub when needed
|
||||
|
||||
Returns:
|
||||
Local directory path that contains the downloaded config file, or the original local directory.
|
||||
"""
|
||||
|
||||
if os.path.exists(model_name_or_path):
|
||||
logger.info("Model already exists locally")
|
||||
return model_name_or_path
|
||||
|
||||
if not download:
|
||||
return model_name_or_path
|
||||
|
||||
with _get_lock(model_name_or_path):
|
||||
# Try `model_index.json` first (diffusers models)
|
||||
try:
|
||||
logger.info(
|
||||
"Downloading model_index.json from HF Hub for %s...",
|
||||
model_name_or_path,
|
||||
)
|
||||
file_path = hf_hub_download(
|
||||
repo_id=model_name_or_path,
|
||||
filename="model_index.json",
|
||||
local_dir=local_dir,
|
||||
)
|
||||
logger.info("Downloaded to %s", file_path)
|
||||
return os.path.dirname(file_path)
|
||||
except Exception as e_index:
|
||||
logger.debug("model_index.json not found or failed: %s", e_index)
|
||||
|
||||
# Fallback to `config.json`
|
||||
try:
|
||||
logger.info(
|
||||
"Downloading config.json from HF Hub for %s...", model_name_or_path
|
||||
)
|
||||
file_path = hf_hub_download(
|
||||
repo_id=model_name_or_path,
|
||||
filename="config.json",
|
||||
local_dir=local_dir,
|
||||
)
|
||||
logger.info("Downloaded to %s", file_path)
|
||||
return os.path.dirname(file_path)
|
||||
except Exception as e_config:
|
||||
raise ValueError(
|
||||
(
|
||||
"Could not find model locally at %s and failed to download "
|
||||
"model_index.json/config.json from HF Hub: %s"
|
||||
)
|
||||
% (model_name_or_path, e_config)
|
||||
) from e_config
|
||||
|
||||
|
||||
# Copied and adapted from hf_diffusers_utils.py
|
||||
def is_diffusers_model_path(model_path: str) -> True:
|
||||
"""
|
||||
Verify if the model directory contains a valid diffusers configuration.
|
||||
|
||||
Args:
|
||||
model_path: Path to the model directory
|
||||
|
||||
Returns:
|
||||
The loaded model configuration as a dictionary if the model is a diffusers model
|
||||
None if the model is not a diffusers model
|
||||
"""
|
||||
|
||||
# Prefer model_index.json which indicates a diffusers pipeline
|
||||
config_path = os.path.join(model_path, "model_index.json")
|
||||
if not os.path.exists(config_path):
|
||||
return False
|
||||
|
||||
# Load the config
|
||||
with open(config_path) as f:
|
||||
config = json.load(f)
|
||||
|
||||
# Verify diffusers version exists
|
||||
if "_diffusers_version" not in config:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def get_is_diffusion_model(model_path: str):
|
||||
model_path = _maybe_download_model(model_path)
|
||||
is_diffusion_model = is_diffusers_model_path(model_path)
|
||||
if is_diffusion_model:
|
||||
logger.info("Diffusion model detected")
|
||||
return is_diffusion_model
|
||||
|
||||
|
||||
def get_model_path(extra_argv):
|
||||
# Find the model_path argument
|
||||
model_path = None
|
||||
for i, arg in enumerate(extra_argv):
|
||||
if arg == "--model-path":
|
||||
if i + 1 < len(extra_argv):
|
||||
model_path = extra_argv[i + 1]
|
||||
break
|
||||
elif arg.startswith("--model-path="):
|
||||
model_path = arg.split("=", 1)[1]
|
||||
break
|
||||
|
||||
if model_path is None:
|
||||
# Fallback for --help or other cases where model-path is not provided
|
||||
if any(h in extra_argv for h in ["-h", "--help"]):
|
||||
raise Exception(
|
||||
"Usage: sglang serve --model-path <model-name-or-path> [additional-arguments]\n\n"
|
||||
"This command can launch either a standard language model server or a diffusion model server.\n"
|
||||
"The server type is determined by the model path.\n"
|
||||
"For specific arguments, please provide a model_path."
|
||||
)
|
||||
else:
|
||||
raise Exception(
|
||||
"Error: --model-path is required. "
|
||||
"Please provide the path to the model."
|
||||
)
|
||||
return model_path
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
subparsers = parser.add_subparsers(dest="subcommand", required=True)
|
||||
|
||||
serve_parser = subparsers.add_parser(
|
||||
"serve",
|
||||
help="Launch the SGLang server.",
|
||||
add_help=False, # Defer help to the specific parser
|
||||
)
|
||||
serve_parser.set_defaults(func=serve)
|
||||
|
||||
generate_parser = subparsers.add_parser(
|
||||
"generate",
|
||||
help="Run inference on a multimodal model.",
|
||||
add_help=False, # Defer help to the specific parser
|
||||
)
|
||||
generate_parser.set_defaults(func=generate)
|
||||
|
||||
args, extra_argv = parser.parse_known_args()
|
||||
args.func(args, extra_argv)
|
||||
42
python/sglang/cli/serve.py
Normal file
42
python/sglang/cli/serve.py
Normal file
@@ -0,0 +1,42 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import os
|
||||
|
||||
from sglang.cli.main import get_is_diffusion_model, get_model_path
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def serve(args, extra_argv):
|
||||
model_path = get_model_path(extra_argv)
|
||||
try:
|
||||
is_diffusion_model = get_is_diffusion_model(model_path)
|
||||
if is_diffusion_model:
|
||||
# Logic for Diffusion Models
|
||||
from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (
|
||||
add_multimodal_gen_serve_args,
|
||||
execute_serve_cmd,
|
||||
)
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="SGLang Diffusion Model Serving"
|
||||
)
|
||||
add_multimodal_gen_serve_args(parser)
|
||||
parsed_args, remaining_argv = parser.parse_known_args(extra_argv)
|
||||
|
||||
execute_serve_cmd(parsed_args, remaining_argv)
|
||||
else:
|
||||
# Logic for Standard Language Models
|
||||
from sglang.launch_server import run_server
|
||||
from sglang.srt.server_args import prepare_server_args
|
||||
|
||||
# Add a dummy argument for the program name, expected by prepare_server_args
|
||||
# as it typically processes sys.argv
|
||||
server_args = prepare_server_args(extra_argv)
|
||||
|
||||
run_server(server_args)
|
||||
finally:
|
||||
kill_process_tree(os.getpid(), include_parent=False)
|
||||
@@ -7,19 +7,23 @@ import sys
|
||||
from sglang.srt.server_args import prepare_server_args
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
|
||||
|
||||
def run_server(server_args):
|
||||
"""Run the server based on server_args.grpc_mode."""
|
||||
if server_args.grpc_mode:
|
||||
from sglang.srt.entrypoints.grpc_server import serve_grpc
|
||||
|
||||
asyncio.run(serve_grpc(server_args))
|
||||
else:
|
||||
from sglang.srt.entrypoints.http_server import launch_server
|
||||
|
||||
launch_server(server_args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
server_args = prepare_server_args(sys.argv[1:])
|
||||
|
||||
try:
|
||||
if server_args.grpc_mode:
|
||||
# Handle gRPC server
|
||||
from sglang.srt.entrypoints.grpc_server import serve_grpc
|
||||
|
||||
asyncio.run(serve_grpc(server_args))
|
||||
else:
|
||||
# Handle HTTP server
|
||||
from sglang.srt.entrypoints.http_server import launch_server
|
||||
|
||||
launch_server(server_args)
|
||||
run_server(server_args)
|
||||
finally:
|
||||
kill_process_tree(os.getpid(), include_parent=False)
|
||||
|
||||
83
python/sglang/multimodal_gen/README.md
Normal file
83
python/sglang/multimodal_gen/README.md
Normal file
@@ -0,0 +1,83 @@
|
||||
<div align="center">
|
||||
<img src=assets/logos/logo.svg width="30%"/>
|
||||
</div>
|
||||
|
||||
**sgl-diffusion is an inference framework for accelerated image/video generation.**
|
||||
|
||||
sgl-diffusion features an end-to-end unified pipeline for accelerating diffusion models. It is designed to be modular and extensible, allowing users to easily add new optimizations and techniques.
|
||||
|
||||
## Key Features
|
||||
|
||||
sgl-diffusion has the following features:
|
||||
|
||||
- State-of-the-art performance optimizations for inference
|
||||
- [Video Sparse Attention](https://arxiv.org/pdf/2505.13389)
|
||||
- [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
|
||||
- [TeaCache](https://arxiv.org/pdf/2411.19108)
|
||||
- [Sage Attention](https://arxiv.org/abs/2410.02367)
|
||||
- USP
|
||||
- CFG Parallel
|
||||
- Diverse hardware and OS support
|
||||
- Supported hardware: H100, H200, A100, B200, 4090
|
||||
- Supported OS: Linux, Windows, MacOS
|
||||
|
||||
## Getting Started
|
||||
|
||||
```bash
|
||||
uv pip install sglang[.diffusion] --prerelease=allow
|
||||
```
|
||||
|
||||
For more information, check the [docs](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen/docs/install.md).
|
||||
|
||||
|
||||
## Inference
|
||||
|
||||
Here's a minimal example to generate a video using the default settings:
|
||||
|
||||
```python
|
||||
from sglang.multimodal_gen import DiffGenerator
|
||||
|
||||
def main():
|
||||
# Create a diff generator from a pre-trained model
|
||||
generator = DiffGenerator.from_pretrained(
|
||||
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
num_gpus=1, # Adjust based on your hardware
|
||||
)
|
||||
|
||||
# Provide a prompt for your video
|
||||
prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest."
|
||||
|
||||
# Generate the video
|
||||
video = generator.generate(
|
||||
prompt,
|
||||
return_frames=True, # Also return frames from this call (defaults to False)
|
||||
output_path="my_videos/", # Controls where videos are saved
|
||||
save_output=True
|
||||
)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
```
|
||||
|
||||
Or, more simply, with the CLI:
|
||||
|
||||
```bash
|
||||
sglang generate --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
|
||||
--text-encoder-cpu-offload --pin-cpu-memory \
|
||||
--prompt "A curious raccoon" \
|
||||
--save-output
|
||||
```
|
||||
|
||||
For more information, check the [docs](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen/docs/cli.md).
|
||||
|
||||
## Contributing
|
||||
|
||||
All contributions are welcome.
|
||||
|
||||
## Acknowledgement
|
||||
|
||||
We learnt and reused code from the following projects:
|
||||
|
||||
- [FastVideo](https://github.com/hao-ai-lab/FastVideo.git). The major components of this repo are based on a fork of FastVide on Sept. 24, 2025.
|
||||
- [xDiT](https://github.com/xdit-project/xDiT). We used the parallelism library from it.
|
||||
- [diffusers](https://github.com/huggingface/diffusers) We used the pipeline design from it.
|
||||
7
python/sglang/multimodal_gen/__init__.py
Normal file
7
python/sglang/multimodal_gen/__init__.py
Normal file
@@ -0,0 +1,7 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
from sglang.multimodal_gen.configs.pipelines import PipelineConfig
|
||||
from sglang.multimodal_gen.configs.sample import SamplingParams
|
||||
from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator
|
||||
|
||||
__all__ = ["DiffGenerator", "PipelineConfig", "SamplingParams"]
|
||||
3
python/sglang/multimodal_gen/configs/__init__.py
Normal file
3
python/sglang/multimodal_gen/configs/__init__.py
Normal file
@@ -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
|
||||
}
|
||||
258
python/sglang/multimodal_gen/configs/configs.py
Normal file
258
python/sglang/multimodal_gen/configs/configs.py
Normal file
@@ -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")
|
||||
48
python/sglang/multimodal_gen/configs/fasthunyuan_t2v.json
Normal file
48
python/sglang/multimodal_gen/configs/fasthunyuan_t2v.json
Normal file
@@ -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
|
||||
}
|
||||
8
python/sglang/multimodal_gen/configs/models/__init__.py
Normal file
8
python/sglang/multimodal_gen/configs/models/__init__.py
Normal file
@@ -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"]
|
||||
105
python/sglang/multimodal_gen/configs/models/base.py
Normal file
105
python/sglang/multimodal_gen/configs/models/base.py
Normal file
@@ -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"]
|
||||
69
python/sglang/multimodal_gen/configs/models/dits/base.py
Normal file
69
python/sglang/multimodal_gen/configs/models/dits/base.py
Normal file
@@ -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
|
||||
36
python/sglang/multimodal_gen/configs/models/dits/flux.py
Normal file
36
python/sglang/multimodal_gen/configs/models/dits/flux.py
Normal file
@@ -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"
|
||||
185
python/sglang/multimodal_gen/configs/models/dits/hunyuanvideo.py
Normal file
185
python/sglang/multimodal_gen/configs/models/dits/hunyuanvideo.py
Normal file
@@ -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"
|
||||
103
python/sglang/multimodal_gen/configs/models/dits/wanvideo.py
Normal file
103
python/sglang/multimodal_gen/configs/models/dits/wanvideo.py
Normal file
@@ -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",
|
||||
]
|
||||
85
python/sglang/multimodal_gen/configs/models/encoders/base.py
Normal file
85
python/sglang/multimodal_gen/configs/models/encoders/base.py
Normal file
@@ -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)
|
||||
95
python/sglang/multimodal_gen/configs/models/encoders/clip.py
Normal file
95
python/sglang/multimodal_gen/configs/models/encoders/clip.py
Normal file
@@ -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"
|
||||
86
python/sglang/multimodal_gen/configs/models/encoders/t5.py
Normal file
86
python/sglang/multimodal_gen/configs/models/encoders/t5.py
Normal file
@@ -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"
|
||||
11
python/sglang/multimodal_gen/configs/models/vaes/__init__.py
Normal file
11
python/sglang/multimodal_gen/configs/models/vaes/__init__.py
Normal file
@@ -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",
|
||||
]
|
||||
158
python/sglang/multimodal_gen/configs/models/vaes/base.py
Normal file
158
python/sglang/multimodal_gen/configs/models/vaes/base.py
Normal file
@@ -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)
|
||||
50
python/sglang/multimodal_gen/configs/models/vaes/flux.py
Normal file
50
python/sglang/multimodal_gen/configs/models/vaes/flux.py
Normal file
@@ -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
|
||||
88
python/sglang/multimodal_gen/configs/models/vaes/wanvae.py
Normal file
88
python/sglang/multimodal_gen/configs/models/vaes/wanvae.py
Normal file
@@ -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
|
||||
37
python/sglang/multimodal_gen/configs/pipelines/__init__.py
Normal file
37
python/sglang/multimodal_gen/configs/pipelines/__init__.py
Normal file
@@ -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",
|
||||
]
|
||||
485
python/sglang/multimodal_gen/configs/pipelines/base.py
Normal file
485
python/sglang/multimodal_gen/configs/pipelines/base.py
Normal file
@@ -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(",")]
|
||||
174
python/sglang/multimodal_gen/configs/pipelines/flux.py
Normal file
174
python/sglang/multimodal_gen/configs/pipelines/flux.py
Normal file
@@ -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
|
||||
),
|
||||
}
|
||||
109
python/sglang/multimodal_gen/configs/pipelines/hunyuan.py
Normal file
109
python/sglang/multimodal_gen/configs/pipelines/hunyuan.py
Normal file
@@ -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
|
||||
299
python/sglang/multimodal_gen/configs/pipelines/qwen_image.py
Normal file
299
python/sglang/multimodal_gen/configs/pipelines/qwen_image.py
Normal file
@@ -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
|
||||
168
python/sglang/multimodal_gen/configs/pipelines/registry.py
Normal file
168
python/sglang/multimodal_gen/configs/pipelines/registry.py
Normal file
@@ -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
|
||||
36
python/sglang/multimodal_gen/configs/pipelines/stepvideo.py
Normal file
36
python/sglang/multimodal_gen/configs/pipelines/stepvideo.py
Normal file
@@ -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"
|
||||
190
python/sglang/multimodal_gen/configs/pipelines/wan.py
Normal file
190
python/sglang/multimodal_gen/configs/pipelines/wan.py
Normal file
@@ -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
|
||||
5
python/sglang/multimodal_gen/configs/sample/__init__.py
Normal file
5
python/sglang/multimodal_gen/configs/sample/__init__.py
Normal file
@@ -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"]
|
||||
494
python/sglang/multimodal_gen/configs/sample/base.py
Normal file
494
python/sglang/multimodal_gen/configs/sample/base.py
Normal file
@@ -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"
|
||||
18
python/sglang/multimodal_gen/configs/sample/flux.py
Normal file
18
python/sglang/multimodal_gen/configs/sample/flux.py
Normal file
@@ -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
|
||||
37
python/sglang/multimodal_gen/configs/sample/hunyuan.py
Normal file
37
python/sglang/multimodal_gen/configs/sample/hunyuan.py
Normal file
@@ -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
|
||||
18
python/sglang/multimodal_gen/configs/sample/qwenimage.py
Normal file
18
python/sglang/multimodal_gen/configs/sample/qwenimage.py
Normal file
@@ -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
|
||||
122
python/sglang/multimodal_gen/configs/sample/registry.py
Normal file
122
python/sglang/multimodal_gen/configs/sample/registry.py
Normal file
@@ -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
|
||||
22
python/sglang/multimodal_gen/configs/sample/stepvideo.py
Normal file
22
python/sglang/multimodal_gen/configs/sample/stepvideo.py
Normal file
@@ -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 = "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。"
|
||||
43
python/sglang/multimodal_gen/configs/sample/teacache.py
Normal file
43
python/sglang/multimodal_gen/configs/sample/teacache.py
Normal file
@@ -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
|
||||
217
python/sglang/multimodal_gen/configs/sample/wan.py
Normal file
217
python/sglang/multimodal_gen/configs/sample/wan.py
Normal file
@@ -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
|
||||
61
python/sglang/multimodal_gen/configs/utils.py
Normal file
61
python/sglang/multimodal_gen/configs/utils.py
Normal file
@@ -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"
|
||||
}
|
||||
31
python/sglang/multimodal_gen/csrc/attn/vmoba_attn/README.md
Normal file
31
python/sglang/multimodal_gen/csrc/attn/vmoba_attn/README.md
Normal file
@@ -0,0 +1,31 @@
|
||||
# Attention Kernel Used in sgl-diffusion
|
||||
|
||||
## VMoBA: Mixture-of-Block Attention for Video Diffusion Models (VMoBA)
|
||||
|
||||
### Installation
|
||||
Please ensure that you have installed FlashAttention version **2.7.1 or higher**, as some interfaces have changed in recent releases.
|
||||
|
||||
### Usage
|
||||
|
||||
You can use `moba_attn_varlen` in the following ways:
|
||||
|
||||
**Install from source:**
|
||||
```bash
|
||||
python setup.py install
|
||||
```
|
||||
|
||||
**Import after installation:**
|
||||
```python
|
||||
from vmoba import moba_attn_varlen
|
||||
```
|
||||
|
||||
**Or import directly from the project root:**
|
||||
```python
|
||||
from csrc.attn.vmoba_attn.vmoba import moba_attn_varlen
|
||||
```
|
||||
|
||||
### Verify if you have successfully installed
|
||||
|
||||
```bash
|
||||
python csrc/attn/vmoba_attn/vmoba/vmoba.py
|
||||
```
|
||||
26
python/sglang/multimodal_gen/csrc/attn/vmoba_attn/setup.py
Normal file
26
python/sglang/multimodal_gen/csrc/attn/vmoba_attn/setup.py
Normal file
@@ -0,0 +1,26 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from setuptools import find_packages, setup
|
||||
|
||||
PACKAGE_NAME = "vmoba"
|
||||
VERSION = "0.0.0"
|
||||
AUTHOR = "JianzongWu"
|
||||
DESCRIPTION = "VMoBA: Mixture-of-Block Attention for Video Diffusion Models"
|
||||
URL = "https://github.com/KwaiVGI/VMoBA"
|
||||
|
||||
setup(
|
||||
name=PACKAGE_NAME,
|
||||
version=VERSION,
|
||||
author=AUTHOR,
|
||||
description=DESCRIPTION,
|
||||
url=URL,
|
||||
packages=find_packages(),
|
||||
classifiers=[
|
||||
"Programming Language :: Python :: 3",
|
||||
"License :: OSI Approved :: Apache Software License",
|
||||
],
|
||||
python_requires=">=3.12",
|
||||
install_requires=[
|
||||
"flash-attn >= 2.7.1",
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,137 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from csrc.attn.vmoba_attn.vmoba import moba_attn_varlen
|
||||
|
||||
|
||||
def generate_test_data(
|
||||
batch_size, total_seqlen, num_heads, head_dim, dtype, device="cuda"
|
||||
):
|
||||
"""
|
||||
Generates random data for testing the variable-length attention function.
|
||||
"""
|
||||
torch.manual_seed(42)
|
||||
random.seed(42)
|
||||
torch.cuda.manual_seed_all(42)
|
||||
|
||||
# Generate sequence lengths for each item in the batch
|
||||
if batch_size > 1:
|
||||
# Ensure sequence lengths are reasonably distributed
|
||||
avg_seqlen = total_seqlen // batch_size
|
||||
seqlens = [
|
||||
random.randint(avg_seqlen // 2, avg_seqlen + avg_seqlen // 2)
|
||||
for _ in range(batch_size - 1)
|
||||
]
|
||||
remaining_len = total_seqlen - sum(seqlens)
|
||||
if remaining_len > 0:
|
||||
seqlens.append(remaining_len)
|
||||
else: # Adjust if sum exceeds total_seqlen
|
||||
seqlens.append(avg_seqlen)
|
||||
current_sum = sum(seqlens)
|
||||
seqlens[-1] -= current_sum - total_seqlen
|
||||
# Ensure all lengths are positive
|
||||
seqlens = [max(1, s) for s in seqlens]
|
||||
# Final adjustment to match total_seqlen
|
||||
seqlens[-1] += total_seqlen - sum(seqlens)
|
||||
|
||||
else:
|
||||
seqlens = [total_seqlen]
|
||||
|
||||
cu_seqlens = torch.tensor(
|
||||
[0] + list(torch.cumsum(torch.tensor(seqlens), 0)),
|
||||
device=device,
|
||||
dtype=torch.int32,
|
||||
)
|
||||
max_seqlen = max(seqlens) if seqlens else 0
|
||||
|
||||
q = torch.randn(
|
||||
(total_seqlen, num_heads, head_dim),
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
requires_grad=False,
|
||||
)
|
||||
k = torch.randn(
|
||||
(total_seqlen, num_heads, head_dim),
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
requires_grad=False,
|
||||
)
|
||||
v = torch.randn(
|
||||
(total_seqlen, num_heads, head_dim),
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
return q, k, v, cu_seqlens, max_seqlen
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [1, 2])
|
||||
@pytest.mark.parametrize("total_seqlen", [512, 1024])
|
||||
@pytest.mark.parametrize("num_heads", [8])
|
||||
@pytest.mark.parametrize("head_dim", [64])
|
||||
@pytest.mark.parametrize("moba_chunk_size", [64])
|
||||
@pytest.mark.parametrize("moba_topk", [2, 4])
|
||||
@pytest.mark.parametrize("select_mode", ["topk", "threshold"])
|
||||
@pytest.mark.parametrize("threshold_type", ["query_head", "head_global", "overall"])
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
||||
def test_moba_attn_varlen_forward(
|
||||
batch_size,
|
||||
total_seqlen,
|
||||
num_heads,
|
||||
head_dim,
|
||||
moba_chunk_size,
|
||||
moba_topk,
|
||||
select_mode,
|
||||
threshold_type,
|
||||
dtype,
|
||||
):
|
||||
"""
|
||||
Tests the forward pass of moba_attn_varlen for basic correctness.
|
||||
It checks output shape, dtype, and for the presence of NaNs/Infs.
|
||||
"""
|
||||
if dtype == torch.float32:
|
||||
pytest.skip("float32 is not supported in flash attention")
|
||||
|
||||
q, k, v, cu_seqlens, max_seqlen = generate_test_data(
|
||||
batch_size, total_seqlen, num_heads, head_dim, dtype
|
||||
)
|
||||
|
||||
# Ensure chunk size is not larger than the smallest sequence length
|
||||
min_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).min().item()
|
||||
if moba_chunk_size > min_seqlen:
|
||||
pytest.skip(
|
||||
"moba_chunk_size is larger than the minimum sequence length in the batch"
|
||||
)
|
||||
|
||||
try:
|
||||
output = moba_attn_varlen(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
moba_chunk_size=moba_chunk_size,
|
||||
moba_topk=moba_topk,
|
||||
select_mode=select_mode,
|
||||
threshold_type=threshold_type,
|
||||
simsum_threshold=0.5, # A reasonable default for threshold mode
|
||||
)
|
||||
except Exception as e:
|
||||
pytest.fail(f"moba_attn_varlen forward pass failed with exception: {e}")
|
||||
|
||||
# 1. Check output shape
|
||||
assert (
|
||||
output.shape == q.shape
|
||||
), f"Expected output shape {q.shape}, but got {output.shape}"
|
||||
|
||||
# 2. Check output dtype
|
||||
assert (
|
||||
output.dtype == q.dtype
|
||||
), f"Expected output dtype {q.dtype}, but got {output.dtype}"
|
||||
|
||||
# 3. Check for NaNs or Infs in the output
|
||||
assert torch.all(torch.isfinite(output)), "Output contains NaN or Inf values"
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from .vmoba import moba_attn_varlen, process_moba_input, process_moba_output
|
||||
1086
python/sglang/multimodal_gen/csrc/attn/vmoba_attn/vmoba/vmoba.py
Normal file
1086
python/sglang/multimodal_gen/csrc/attn/vmoba_attn/vmoba/vmoba.py
Normal file
File diff suppressed because it is too large
Load Diff
274
python/sglang/multimodal_gen/docs/cli.md
Normal file
274
python/sglang/multimodal_gen/docs/cli.md
Normal file
@@ -0,0 +1,274 @@
|
||||
# sgl-diffusion CLI Inference
|
||||
|
||||
The sgl-diffusion CLI provides a quick way to access the sgl-diffusion inference pipeline for image and video generation.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- A working sgl-diffusion installation and the `sgl-diffusion` CLI available in `$PATH`.
|
||||
- Python 3.10+ if you plan to use the OpenAI Python SDK.
|
||||
|
||||
|
||||
## Supported Arguments
|
||||
|
||||
### Server Arguments
|
||||
|
||||
- `--model-path {MODEL_PATH}`: Path to the model or model ID
|
||||
- `--num-gpus {NUM_GPUS}`: Number of GPUs to use
|
||||
- `--tp-size {TP_SIZE}`: Tensor parallelism size (only for the encoder; should not be larger than 1 if text encoder offload is enabled, as layer-wise offload plus prefetch is faster)
|
||||
- `--sp-size {SP_SIZE}`: Sequence parallelism size (typically should match the number of GPUs)
|
||||
- `--ulysses-degree {ULYSSES_DEGREE}`: The degree of DeepSpeed-Ulysses-style SP in USP
|
||||
- `--ring-degree {RING_DEGREE}`: The degree of ring attention-style SP in USP
|
||||
|
||||
|
||||
### Sampling Parameters
|
||||
|
||||
- `--prompt {PROMPT}`: Text description for the video you want to generate
|
||||
- `--num-inference-steps {STEPS}`: Number of denoising steps
|
||||
- `--negative-prompt {PROMPT}`: Negative prompt to guide generation away from certain concepts
|
||||
- `--seed {SEED}`: Random seed for reproducible generation
|
||||
|
||||
|
||||
#### Image/Video Configuration
|
||||
|
||||
- `--height {HEIGHT}`: Height of the generated output
|
||||
- `--width {WIDTH}`: Width of the generated output
|
||||
- `--num-frames {NUM_FRAMES}`: Number of frames to generate
|
||||
- `--fps {FPS}`: Frames per second for the saved output, if this is a video-generation task
|
||||
|
||||
|
||||
#### Output Options
|
||||
|
||||
- `--output-path {PATH}`: Directory to save the generated video
|
||||
- `--save-output`: Whether to save the image/video to disk
|
||||
- `--return-frames`: Whether to return the raw frames
|
||||
|
||||
### Using Configuration Files
|
||||
|
||||
Instead of specifying all parameters on the command line, you can use a configuration file:
|
||||
|
||||
```bash
|
||||
sglang generate --config {CONFIG_FILE_PATH}
|
||||
```
|
||||
|
||||
The configuration file should be in JSON or YAML format with the same parameter names as the CLI options. Command-line arguments take precedence over settings in the configuration file, allowing you to override specific values while keeping the rest from the configuration file.
|
||||
|
||||
Example configuration file (config.json):
|
||||
|
||||
```json
|
||||
{
|
||||
"model_path": "FastVideo/FastHunyuan-diffusers",
|
||||
"prompt": "A beautiful woman in a red dress walking down a street",
|
||||
"output_path": "outputs/",
|
||||
"num_gpus": 2,
|
||||
"sp_size": 2,
|
||||
"tp_size": 1,
|
||||
"num_frames": 45,
|
||||
"height": 720,
|
||||
"width": 1280,
|
||||
"num_inference_steps": 6,
|
||||
"seed": 1024,
|
||||
"fps": 24,
|
||||
"precision": "bf16",
|
||||
"vae_precision": "fp16",
|
||||
"vae_tiling": true,
|
||||
"vae_sp": true,
|
||||
"vae_config": {
|
||||
"load_encoder": false,
|
||||
"load_decoder": true,
|
||||
"tile_sample_min_height": 256,
|
||||
"tile_sample_min_width": 256
|
||||
},
|
||||
"text_encoder_precisions": [
|
||||
"fp16",
|
||||
"fp16"
|
||||
],
|
||||
"mask_strategy_file_path": null,
|
||||
"enable_torch_compile": false
|
||||
}
|
||||
```
|
||||
|
||||
Or using YAML format (config.yaml):
|
||||
|
||||
```yaml
|
||||
model_path: "FastVideo/FastHunyuan-diffusers"
|
||||
prompt: "A beautiful woman in a red dress walking down a street"
|
||||
output_path: "outputs/"
|
||||
num_gpus: 2
|
||||
sp_size: 2
|
||||
tp_size: 1
|
||||
num_frames: 45
|
||||
height: 720
|
||||
width: 1280
|
||||
num_inference_steps: 6
|
||||
seed: 1024
|
||||
fps: 24
|
||||
precision: "bf16"
|
||||
vae_precision: "fp16"
|
||||
vae_tiling: true
|
||||
vae_sp: true
|
||||
vae_config:
|
||||
load_encoder: false
|
||||
load_decoder: true
|
||||
tile_sample_min_height: 256
|
||||
tile_sample_min_width: 256
|
||||
text_encoder_precisions:
|
||||
- "fp16"
|
||||
- "fp16"
|
||||
mask_strategy_file_path: null
|
||||
enable_torch_compile: false
|
||||
```
|
||||
|
||||
|
||||
To see all the options, you can use the `--help` flag:
|
||||
|
||||
```bash
|
||||
sglang generate --help
|
||||
```
|
||||
|
||||
## Serve
|
||||
|
||||
Launch the sgl-diffusion HTTP server and interact with it using the OpenAI SDK and curl. The server implements an OpenAI-compatible subset for Videos under the `/v1/videos` namespace.
|
||||
|
||||
### Start the server
|
||||
|
||||
Use the following command to launch the server:
|
||||
|
||||
```bash
|
||||
SERVER_ARGS=(
|
||||
--model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers
|
||||
--text-encoder-cpu-offload
|
||||
--pin-cpu-memory
|
||||
--num-gpus 4
|
||||
--ulysses-degree=2
|
||||
--ring-degree=2
|
||||
)
|
||||
|
||||
sglang serve $SERVER_ARGS
|
||||
```
|
||||
|
||||
- **--model-path**: Which model to load. The example uses `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`.
|
||||
- **--port**: HTTP port to listen on (the default here is `30010`).
|
||||
|
||||
Wait until the port is listening. In CI, the tests probe `127.0.0.1:30010` before sending requests.
|
||||
|
||||
### OpenAI Python SDK usage
|
||||
|
||||
Initialize the client with a dummy API key and point `base_url` to your local server:
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(api_key="sk-proj-1234567890", base_url="http://localhost:30010/v1")
|
||||
```
|
||||
|
||||
- **Create a video**
|
||||
|
||||
```python
|
||||
video = client.videos.create(prompt="A calico cat playing a piano on stage", size="1280x720")
|
||||
print(video.id, video.status)
|
||||
```
|
||||
|
||||
Response example fields include `id`, `status` (e.g., `queued` → `completed`), `size`, and `seconds`.
|
||||
|
||||
- **List videos**
|
||||
|
||||
```python
|
||||
videos = client.videos.list()
|
||||
for item in videos.data:
|
||||
print(item.id, item.status)
|
||||
```
|
||||
|
||||
- **Poll for completion and download content**
|
||||
|
||||
```python
|
||||
import time
|
||||
|
||||
video = client.videos.create(prompt="A calico cat playing a piano on stage", size="1280x720")
|
||||
video_id = video.id
|
||||
|
||||
# Simple polling loop
|
||||
while True:
|
||||
page = client.videos.list()
|
||||
item = next((v for v in page.data if v.id == video_id), None)
|
||||
if item and item.status == "completed":
|
||||
break
|
||||
time.sleep(5)
|
||||
|
||||
# Download binary content (MP4)
|
||||
resp = client.videos.download_content(video_id=video_id)
|
||||
content = resp.read() # bytes
|
||||
with open("output.mp4", "wb") as f:
|
||||
f.write(content)
|
||||
```
|
||||
|
||||
### curl examples
|
||||
|
||||
- **Create a video**
|
||||
|
||||
```bash
|
||||
curl -sS -X POST "http://localhost:30010/v1/videos" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer sk-proj-1234567890" \
|
||||
-d '{
|
||||
"prompt": "A calico cat playing a piano on stage",
|
||||
"size": "1280x720"
|
||||
}'
|
||||
```
|
||||
|
||||
- **List videos**
|
||||
|
||||
```bash
|
||||
curl -sS -X GET "http://localhost:30010/v1/videos" \
|
||||
-H "Authorization: Bearer sk-proj-1234567890"
|
||||
```
|
||||
|
||||
- **Download video content**
|
||||
|
||||
```bash
|
||||
curl -sS -L "http://localhost:30010/v1/videos/<VIDEO_ID>/content" \
|
||||
-H "Authorization: Bearer sk-proj-1234567890" \
|
||||
-o output.mp4
|
||||
```
|
||||
|
||||
### API surface implemented here
|
||||
|
||||
The server exposes these endpoints (OpenAPI tag `videos`):
|
||||
|
||||
- `POST /v1/videos` — Create a generation job and return a queued `video` object.
|
||||
- `GET /v1/videos` — List jobs.
|
||||
- `GET /v1/videos/{video_id}/content` — Download binary content when ready (e.g., MP4).
|
||||
|
||||
### Reference
|
||||
|
||||
- OpenAI Videos API reference: `https://platform.openai.com/docs/api-reference/videos`
|
||||
|
||||
## Generate
|
||||
|
||||
Run a one-off generation task without launching a persistent server.
|
||||
|
||||
To use it, pass both server arguments and sampling parameters in one command, after the `generate` subcommand, for example:
|
||||
|
||||
```bash
|
||||
SERVER_ARGS=(
|
||||
--model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers
|
||||
--text-encoder-cpu-offload
|
||||
--pin-cpu-memory
|
||||
--num-gpus 4
|
||||
--ulysses-degree=2
|
||||
--ring-degree=2
|
||||
)
|
||||
|
||||
SAMPLING_ARGS=(
|
||||
--prompt "A curious raccoon"
|
||||
--save-output
|
||||
--output-path outputs
|
||||
--output-file-name "A curious raccoon.mp4"
|
||||
)
|
||||
|
||||
sglang generate $SERVER_ARGS $SAMPLING_ARGS
|
||||
```
|
||||
|
||||
Once the generation task has finished, the server will shut down automatically.
|
||||
|
||||
> [!NOTE]
|
||||
> The HTTP server-related arguments are ignored in this subcommand.
|
||||
52
python/sglang/multimodal_gen/docs/install.md
Normal file
52
python/sglang/multimodal_gen/docs/install.md
Normal file
@@ -0,0 +1,52 @@
|
||||
# Install sgl-diffusion
|
||||
|
||||
You can install sgl-diffusion using one of the methods below.
|
||||
|
||||
This page primarily applies to common NVIDIA GPU platforms.
|
||||
|
||||
## Method 1: With pip or uv
|
||||
|
||||
It is recommended to use uv for a faster installation:
|
||||
|
||||
```bash
|
||||
pip install --upgrade pip
|
||||
pip install uv
|
||||
uv pip install sglang[.diffusion] --prerelease=allow
|
||||
```
|
||||
|
||||
## Method 2: From source
|
||||
|
||||
```bash
|
||||
# Use the latest release branch
|
||||
git clone https://github.com/sgl-project/sglang.git
|
||||
cd sglang
|
||||
|
||||
# Install the Python packages
|
||||
pip install --upgrade pip
|
||||
pip install -e "python/.[diffusion]"
|
||||
|
||||
# With uv
|
||||
uv pip install --prerelease=allow -e "python/.[diffusion]"
|
||||
```
|
||||
|
||||
**Quick fixes for common problems:**
|
||||
|
||||
- If you want to develop sgl-diffusion, it is recommended to use Docker. The Docker image is `lmsysorg/sgl-diffusion:latest`.
|
||||
|
||||
## Method 3: Using Docker
|
||||
|
||||
The Docker images are available on Docker Hub at [lmsysorg/sgl-diffusion](), built from the [Dockerfile](https://github.com/sgl-project/sgl-diffusion/tree/main/docker).
|
||||
Replace `<secret>` below with your HuggingFace Hub [token](https://huggingface.co/docs/hub/en/security-tokens).
|
||||
|
||||
```bash
|
||||
docker run --gpus all \
|
||||
--shm-size 32g \
|
||||
-p 30000:30000 \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=<secret>" \
|
||||
--ipc=host \
|
||||
lmsysorg/sglang:diffusion \
|
||||
sglang generate --model-path black-forest-labs/FLUX.1-dev \
|
||||
--prompt "A logo With Bold Large text: SGL Diffusion" \
|
||||
--save-output
|
||||
```
|
||||
46
python/sglang/multimodal_gen/docs/support_matrix.md
Normal file
46
python/sglang/multimodal_gen/docs/support_matrix.md
Normal file
@@ -0,0 +1,46 @@
|
||||
# Compatibility Matrix
|
||||
|
||||
The table below shows every supported model and the optimizations supported for them.
|
||||
|
||||
The symbols used have the following meanings:
|
||||
|
||||
- ✅ = Full compatibility
|
||||
- ❌ = No compatibility
|
||||
- ⭕ = Does not apply to this model
|
||||
|
||||
## Models x Optimization
|
||||
|
||||
The `HuggingFace Model ID` can be passed directly to `from_pretrained()` methods, and sgl-diffusion will use the optimal
|
||||
default parameters when initializing and generating videos.
|
||||
|
||||
### Video Generation Models
|
||||
|
||||
| Model Name | Hugging Face Model ID | Resolutions | TeaCache | Sliding Tile Attn | Sage Attn | Video Sparse Attention (VSA) |
|
||||
|:-----------------------------|:--------------------------------------------------|:---------------------------------------------|:--------:|:-----------------:|:---------:|:----------------------------:|
|
||||
| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480p | ⭕ | ⭕ | ⭕ | ✅ |
|
||||
| FastWan2.2 TI2V 5B Full Attn | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720p | ⭕ | ⭕ | ⭕ | ✅ |
|
||||
| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720p | ⭕ | ⭕ | ✅ | ⭕ |
|
||||
| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480p<br>720p | ❌ | ❌ | ✅ | ⭕ |
|
||||
| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480p<br>720p | ❌ | ❌ | ✅ | ⭕ |
|
||||
| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720×1280<br>544×960 | ❌ | ✅ | ✅ | ⭕ |
|
||||
| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720×1280<br>544×960 | ❌ | ✅ | ✅ | ⭕ |
|
||||
| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ |
|
||||
| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480p, 720p | ✅ | ✅ | ✅ | ⭕ |
|
||||
| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ |
|
||||
| Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720p | ✅ | ✅ | ✅ | ⭕ |
|
||||
|
||||
**Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.
|
||||
|
||||
### Image Generation Models
|
||||
|
||||
| Model Name | HuggingFace Model ID | Resolutions | TeaCache | Sage Attn |
|
||||
|:----------------|:-------------------------------|:---------------|:--------:|:---------:|
|
||||
| FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | Any resolution | ❌ | ❌ |
|
||||
| Qwen Image | `Qwen/Qwen-Image` | Any resolution | ❌ | ❌ |
|
||||
| Qwen Image Edit | `Qwen/Qwen-Image-Edit` | Any resolution | ❌ | ❌ |
|
||||
|
||||
## Special requirements
|
||||
|
||||
### Sliding Tile Attention
|
||||
|
||||
- Currently, only Hopper GPUs (H100s) are supported.
|
||||
326
python/sglang/multimodal_gen/envs.py
Normal file
326
python/sglang/multimodal_gen/envs.py
Normal file
@@ -0,0 +1,326 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
import importlib.util
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/envs.py
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import diffusers
|
||||
import torch
|
||||
from packaging import version
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
SGL_DIFFUSION_RINGBUFFER_WARNING_INTERVAL: int = 60
|
||||
SGL_DIFFUSION_NCCL_SO_PATH: str | None = None
|
||||
LD_LIBRARY_PATH: str | None = None
|
||||
LOCAL_RANK: int = 0
|
||||
CUDA_VISIBLE_DEVICES: str | None = None
|
||||
SGL_DIFFUSION_CACHE_ROOT: str = os.path.expanduser("~/.cache/sgl_diffusion")
|
||||
SGL_DIFFUSION_CONFIG_ROOT: str = os.path.expanduser("~/.config/sgl_diffusion")
|
||||
SGL_DIFFUSION_CONFIGURE_LOGGING: int = 1
|
||||
SGL_DIFFUSION_LOGGING_LEVEL: str = "INFO"
|
||||
SGL_DIFFUSION_LOGGING_PREFIX: str = ""
|
||||
SGL_DIFFUSION_LOGGING_CONFIG_PATH: str | None = None
|
||||
SGL_DIFFUSION_TRACE_FUNCTION: int = 0
|
||||
SGL_DIFFUSION_WORKER_MULTIPROC_METHOD: str = "fork"
|
||||
SGL_DIFFUSION_TARGET_DEVICE: str = "cuda"
|
||||
MAX_JOBS: str | None = None
|
||||
NVCC_THREADS: str | None = None
|
||||
CMAKE_BUILD_TYPE: str | None = None
|
||||
VERBOSE: bool = False
|
||||
SGL_DIFFUSION_SERVER_DEV_MODE: bool = False
|
||||
SGL_DIFFUSION_STAGE_LOGGING: bool = False
|
||||
|
||||
|
||||
def _is_hip():
|
||||
has_rocm = torch.version.hip is not None
|
||||
return has_rocm
|
||||
|
||||
|
||||
def _is_cuda():
|
||||
has_cuda = torch.version.cuda is not None
|
||||
return has_cuda
|
||||
|
||||
|
||||
def _is_musa():
|
||||
try:
|
||||
if hasattr(torch, "musa") and torch.musa.is_available():
|
||||
return True
|
||||
except ModuleNotFoundError:
|
||||
return False
|
||||
|
||||
|
||||
def _is_mps():
|
||||
return torch.backends.mps.is_available()
|
||||
|
||||
|
||||
class PackagesEnvChecker:
|
||||
_instance = None
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super(PackagesEnvChecker, cls).__new__(cls)
|
||||
cls._instance.initialize()
|
||||
return cls._instance
|
||||
|
||||
def initialize(self):
|
||||
self.packages_info = {
|
||||
"has_aiter": self.check_aiter(),
|
||||
"diffusers_version": self.check_diffusers_version(),
|
||||
}
|
||||
|
||||
def check_aiter(self):
|
||||
"""
|
||||
Checks whether ROCm AITER library is installed
|
||||
"""
|
||||
try:
|
||||
|
||||
logger.info("Using AITER as the attention library")
|
||||
return True
|
||||
except:
|
||||
if _is_hip():
|
||||
logger.warning(
|
||||
f'Using AMD GPUs, but library "aiter" is not installed, '
|
||||
"defaulting to other attention mechanisms"
|
||||
)
|
||||
return False
|
||||
|
||||
def check_flash_attn(self):
|
||||
if not torch.cuda.is_available():
|
||||
return False
|
||||
if _is_musa():
|
||||
logger.info(
|
||||
"Flash Attention library is not supported on MUSA for the moment."
|
||||
)
|
||||
return False
|
||||
try:
|
||||
return True
|
||||
except ImportError:
|
||||
logger.warning(
|
||||
f'Flash Attention library "flash_attn" not found, '
|
||||
f"using pytorch attention implementation"
|
||||
)
|
||||
return False
|
||||
|
||||
def check_long_ctx_attn(self):
|
||||
if not torch.cuda.is_available():
|
||||
return False
|
||||
try:
|
||||
return importlib.util.find_spec("yunchang") is not None
|
||||
except ImportError:
|
||||
logger.warning(
|
||||
f'Ring Flash Attention library "yunchang" not found, '
|
||||
f"using pytorch attention implementation"
|
||||
)
|
||||
return False
|
||||
|
||||
def check_diffusers_version(self):
|
||||
if version.parse(
|
||||
version.parse(diffusers.__version__).base_version
|
||||
) < version.parse("0.30.0"):
|
||||
raise RuntimeError(
|
||||
f"Diffusers version: {version.parse(version.parse(diffusers.__version__).base_version)} is not supported,"
|
||||
f"please upgrade to version > 0.30.0"
|
||||
)
|
||||
return version.parse(version.parse(diffusers.__version__).base_version)
|
||||
|
||||
def get_packages_info(self):
|
||||
return self.packages_info
|
||||
|
||||
|
||||
PACKAGES_CHECKER = PackagesEnvChecker()
|
||||
|
||||
|
||||
def get_default_cache_root() -> str:
|
||||
return os.getenv(
|
||||
"XDG_CACHE_HOME",
|
||||
os.path.join(os.path.expanduser("~"), ".cache"),
|
||||
)
|
||||
|
||||
|
||||
def get_default_config_root() -> str:
|
||||
return os.getenv(
|
||||
"XDG_CONFIG_HOME",
|
||||
os.path.join(os.path.expanduser("~"), ".config"),
|
||||
)
|
||||
|
||||
|
||||
def maybe_convert_int(value: str | None) -> int | None:
|
||||
if value is None:
|
||||
return None
|
||||
return int(value)
|
||||
|
||||
|
||||
# The begin-* and end* here are used by the documentation generator
|
||||
# to extract the used env vars.
|
||||
|
||||
# begin-env-vars-definition
|
||||
|
||||
environment_variables: dict[str, Callable[[], Any]] = {
|
||||
# ================== Installation Time Env Vars ==================
|
||||
# Target device of sgl-diffusion, supporting [cuda (by default),
|
||||
# rocm, neuron, cpu, openvino]
|
||||
"SGL_DIFFUSION_TARGET_DEVICE": lambda: os.getenv(
|
||||
"SGL_DIFFUSION_TARGET_DEVICE", "cuda"
|
||||
),
|
||||
# Maximum number of compilation jobs to run in parallel.
|
||||
# By default this is the number of CPUs
|
||||
"MAX_JOBS": lambda: os.getenv("MAX_JOBS", None),
|
||||
# Number of threads to use for nvcc
|
||||
# By default this is 1.
|
||||
# If set, `MAX_JOBS` will be reduced to avoid oversubscribing the CPU.
|
||||
"NVCC_THREADS": lambda: os.getenv("NVCC_THREADS", None),
|
||||
# If set, sgl_diffusion will use precompiled binaries (*.so)
|
||||
"SGL_DIFFUSION_USE_PRECOMPILED": lambda: bool(
|
||||
os.environ.get("SGL_DIFFUSION_USE_PRECOMPILED")
|
||||
)
|
||||
or bool(os.environ.get("SGL_DIFFUSION_PRECOMPILED_WHEEL_LOCATION")),
|
||||
# CMake build type
|
||||
# If not set, defaults to "Debug" or "RelWithDebInfo"
|
||||
# Available options: "Debug", "Release", "RelWithDebInfo"
|
||||
"CMAKE_BUILD_TYPE": lambda: os.getenv("CMAKE_BUILD_TYPE"),
|
||||
# If set, sgl_diffusion will print verbose logs during installation
|
||||
"VERBOSE": lambda: bool(int(os.getenv("VERBOSE", "0"))),
|
||||
# Root directory for FASTVIDEO configuration files
|
||||
# Defaults to `~/.config/sgl_diffusion` unless `XDG_CONFIG_HOME` is set
|
||||
# Note that this not only affects how sgl_diffusion finds its configuration files
|
||||
# during runtime, but also affects how sgl_diffusion installs its configuration
|
||||
# files during **installation**.
|
||||
"SGL_DIFFUSION_CONFIG_ROOT": lambda: os.path.expanduser(
|
||||
os.getenv(
|
||||
"SGL_DIFFUSION_CONFIG_ROOT",
|
||||
os.path.join(get_default_config_root(), "sgl_diffusion"),
|
||||
)
|
||||
),
|
||||
# ================== Runtime Env Vars ==================
|
||||
# Root directory for FASTVIDEO cache files
|
||||
# Defaults to `~/.cache/sgl_diffusion` unless `XDG_CACHE_HOME` is set
|
||||
"SGL_DIFFUSION_CACHE_ROOT": lambda: os.path.expanduser(
|
||||
os.getenv(
|
||||
"SGL_DIFFUSION_CACHE_ROOT",
|
||||
os.path.join(get_default_cache_root(), "sgl_diffusion"),
|
||||
)
|
||||
),
|
||||
# Interval in seconds to log a warning message when the ring buffer is full
|
||||
"SGL_DIFFUSION_RINGBUFFER_WARNING_INTERVAL": lambda: int(
|
||||
os.environ.get("SGL_DIFFUSION_RINGBUFFER_WARNING_INTERVAL", "60")
|
||||
),
|
||||
# Path to the NCCL library file. It is needed because nccl>=2.19 brought
|
||||
# by PyTorch contains a bug: https://github.com/NVIDIA/nccl/issues/1234
|
||||
"SGL_DIFFUSION_NCCL_SO_PATH": lambda: os.environ.get(
|
||||
"SGL_DIFFUSION_NCCL_SO_PATH", None
|
||||
),
|
||||
# when `SGL_DIFFUSION_NCCL_SO_PATH` is not set, sgl_diffusion will try to find the nccl
|
||||
# library file in the locations specified by `LD_LIBRARY_PATH`
|
||||
"LD_LIBRARY_PATH": lambda: os.environ.get("LD_LIBRARY_PATH", None),
|
||||
# Internal flag to enable Dynamo fullgraph capture
|
||||
"SGL_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE": lambda: bool(
|
||||
os.environ.get("SGL_DIFFUSION_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1") != "0"
|
||||
),
|
||||
# local rank of the process in the distributed setting, used to determine
|
||||
# the GPU device id
|
||||
"LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")),
|
||||
# used to control the visible devices in the distributed setting
|
||||
"CUDA_VISIBLE_DEVICES": lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None),
|
||||
# timeout for each iteration in the engine
|
||||
"SGL_DIFFUSION_ENGINE_ITERATION_TIMEOUT_S": lambda: int(
|
||||
os.environ.get("SGL_DIFFUSION_ENGINE_ITERATION_TIMEOUT_S", "60")
|
||||
),
|
||||
# Logging configuration
|
||||
# If set to 0, sgl_diffusion will not configure logging
|
||||
# If set to 1, sgl_diffusion will configure logging using the default configuration
|
||||
# or the configuration file specified by SGL_DIFFUSION_LOGGING_CONFIG_PATH
|
||||
"SGL_DIFFUSION_CONFIGURE_LOGGING": lambda: int(
|
||||
os.getenv("SGL_DIFFUSION_CONFIGURE_LOGGING", "1")
|
||||
),
|
||||
"SGL_DIFFUSION_LOGGING_CONFIG_PATH": lambda: os.getenv(
|
||||
"SGL_DIFFUSION_LOGGING_CONFIG_PATH"
|
||||
),
|
||||
# this is used for configuring the default logging level
|
||||
"SGL_DIFFUSION_LOGGING_LEVEL": lambda: os.getenv(
|
||||
"SGL_DIFFUSION_LOGGING_LEVEL", "INFO"
|
||||
),
|
||||
# if set, SGL_DIFFUSION_LOGGING_PREFIX will be prepended to all log messages
|
||||
"SGL_DIFFUSION_LOGGING_PREFIX": lambda: os.getenv(
|
||||
"SGL_DIFFUSION_LOGGING_PREFIX", ""
|
||||
),
|
||||
# Trace function calls
|
||||
# If set to 1, sgl_diffusion will trace function calls
|
||||
# Useful for debugging
|
||||
"SGL_DIFFUSION_TRACE_FUNCTION": lambda: int(
|
||||
os.getenv("SGL_DIFFUSION_TRACE_FUNCTION", "0")
|
||||
),
|
||||
# Path to the attention configuration file. Only used for sliding tile
|
||||
# attention for now.
|
||||
"SGL_DIFFUSION_ATTENTION_CONFIG": lambda: (
|
||||
None
|
||||
if os.getenv("SGL_DIFFUSION_ATTENTION_CONFIG", None) is None
|
||||
else os.path.expanduser(os.getenv("SGL_DIFFUSION_ATTENTION_CONFIG", "."))
|
||||
),
|
||||
# Use dedicated multiprocess context for workers.
|
||||
# Both spawn and fork work
|
||||
"SGL_DIFFUSION_WORKER_MULTIPROC_METHOD": lambda: os.getenv(
|
||||
"SGL_DIFFUSION_WORKER_MULTIPROC_METHOD", "fork"
|
||||
),
|
||||
# Enables torch profiler if set. Path to the directory where torch profiler
|
||||
# traces are saved. Note that it must be an absolute path.
|
||||
"SGL_DIFFUSION_TORCH_PROFILER_DIR": lambda: (
|
||||
None
|
||||
if os.getenv("SGL_DIFFUSION_TORCH_PROFILER_DIR", None) is None
|
||||
else os.path.expanduser(os.getenv("SGL_DIFFUSION_TORCH_PROFILER_DIR", "."))
|
||||
),
|
||||
# If set, sgl_diffusion will run in development mode, which will enable
|
||||
# some additional endpoints for developing and debugging,
|
||||
# e.g. `/reset_prefix_cache`
|
||||
"SGL_DIFFUSION_SERVER_DEV_MODE": lambda: bool(
|
||||
int(os.getenv("SGL_DIFFUSION_SERVER_DEV_MODE", "0"))
|
||||
),
|
||||
# If set, sgl_diffusion will enable stage logging, which will print the time
|
||||
# taken for each stage
|
||||
"SGL_DIFFUSION_STAGE_LOGGING": lambda: bool(
|
||||
int(os.getenv("SGL_DIFFUSION_STAGE_LOGGING", "0"))
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
# end-env-vars-definition
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
# lazy evaluation of environment variables
|
||||
if name in environment_variables:
|
||||
return environment_variables[name]()
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
def __dir__():
|
||||
return list(environment_variables.keys())
|
||||
|
||||
|
||||
def get_torch_distributed_backend() -> str:
|
||||
if torch.cuda.is_available():
|
||||
return "nccl"
|
||||
elif _is_musa():
|
||||
return "mccl"
|
||||
elif _is_mps():
|
||||
return "gloo"
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"No Accelerators(AMD/NV/MTT GPU, AMD MI instinct accelerators) available"
|
||||
)
|
||||
|
||||
|
||||
def get_device(local_rank: int) -> torch.device:
|
||||
if torch.cuda.is_available():
|
||||
return torch.device("cuda", local_rank)
|
||||
elif _is_musa():
|
||||
return torch.device("musa", local_rank)
|
||||
elif _is_mps():
|
||||
return torch.device("mps")
|
||||
else:
|
||||
return torch.device("cpu")
|
||||
@@ -0,0 +1,8 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Basic inference pipelines for sglang.multimodal_gen.
|
||||
|
||||
This package contains basic pipelines for video and image generation.
|
||||
"""
|
||||
@@ -0,0 +1 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
@@ -0,0 +1,126 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Hunyuan video diffusion pipeline implementation.
|
||||
|
||||
This module contains an implementation of the Hunyuan video diffusion pipeline
|
||||
using the modular pipeline architecture.
|
||||
"""
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines import ComposedPipelineBase, Req
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages import (
|
||||
ConditioningStage,
|
||||
DecodingStage,
|
||||
DenoisingStage,
|
||||
InputValidationStage,
|
||||
LatentPreparationStage,
|
||||
TextEncodingStage,
|
||||
TimestepPreparationStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
# TODO(will): move PRECISION_TO_TYPE to better place
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def calculate_shift(
|
||||
image_seq_len,
|
||||
base_seq_len: int = 256,
|
||||
max_seq_len: int = 4096,
|
||||
base_shift: float = 0.5,
|
||||
max_shift: float = 1.15,
|
||||
):
|
||||
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
|
||||
b = base_shift - m * base_seq_len
|
||||
mu = image_seq_len * m + b
|
||||
return mu
|
||||
|
||||
|
||||
def prepare_mu(batch: Req, server_args: ServerArgs):
|
||||
height = batch.height
|
||||
width = batch.width
|
||||
vae_scale_factor = (
|
||||
server_args.pipeline_config.vae_config.arch_config.vae_scale_factor
|
||||
)
|
||||
image_seq_len = (int(height) // vae_scale_factor) * (int(width) // vae_scale_factor)
|
||||
|
||||
mu = calculate_shift(
|
||||
image_seq_len,
|
||||
# hard code, since scheduler_config is not in PipelineConfig now
|
||||
256,
|
||||
4096,
|
||||
0.5,
|
||||
1.15,
|
||||
)
|
||||
return "mu", mu
|
||||
|
||||
|
||||
class FluxPipeline(ComposedPipelineBase):
|
||||
pipeline_name = "FluxPipeline"
|
||||
|
||||
_required_config_modules = [
|
||||
"text_encoder",
|
||||
"text_encoder_2",
|
||||
"tokenizer",
|
||||
"tokenizer_2",
|
||||
"vae",
|
||||
"transformer",
|
||||
"scheduler",
|
||||
]
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||
"""Set up pipeline stages with proper dependency injection."""
|
||||
|
||||
self.add_stage(
|
||||
stage_name="input_validation_stage", stage=InputValidationStage()
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage_primary",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[
|
||||
self.get_module("text_encoder"),
|
||||
self.get_module("text_encoder_2"),
|
||||
],
|
||||
tokenizers=[
|
||||
self.get_module("tokenizer"),
|
||||
self.get_module("tokenizer_2"),
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
|
||||
|
||||
self.add_stage(
|
||||
stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
prepare_extra_set_timesteps_kwargs=[prepare_mu],
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="latent_preparation_stage",
|
||||
stage=LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="denoising_stage",
|
||||
stage=DenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
|
||||
)
|
||||
|
||||
|
||||
EntryClass = FluxPipeline
|
||||
@@ -0,0 +1 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
@@ -0,0 +1,93 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Hunyuan video diffusion pipeline implementation.
|
||||
|
||||
This module contains an implementation of the Hunyuan video diffusion pipeline
|
||||
using the modular pipeline architecture.
|
||||
"""
|
||||
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines import ComposedPipelineBase
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages import (
|
||||
ConditioningStage,
|
||||
DecodingStage,
|
||||
DenoisingStage,
|
||||
InputValidationStage,
|
||||
LatentPreparationStage,
|
||||
TextEncodingStage,
|
||||
TimestepPreparationStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
# TODO(will): move PRECISION_TO_TYPE to better place
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class HunyuanVideoPipeline(ComposedPipelineBase):
|
||||
|
||||
pipeline_name = "HunyuanVideoPipeline"
|
||||
|
||||
_required_config_modules = [
|
||||
"text_encoder",
|
||||
"text_encoder_2",
|
||||
"tokenizer",
|
||||
"tokenizer_2",
|
||||
"vae",
|
||||
"transformer",
|
||||
"scheduler",
|
||||
]
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||
"""Set up pipeline stages with proper dependency injection."""
|
||||
|
||||
self.add_stage(
|
||||
stage_name="input_validation_stage", stage=InputValidationStage()
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage_primary",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[
|
||||
self.get_module("text_encoder"),
|
||||
self.get_module("text_encoder_2"),
|
||||
],
|
||||
tokenizers=[
|
||||
self.get_module("tokenizer"),
|
||||
self.get_module("tokenizer_2"),
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
|
||||
|
||||
self.add_stage(
|
||||
stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="latent_preparation_stage",
|
||||
stage=LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="denoising_stage",
|
||||
stage=DenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
|
||||
)
|
||||
|
||||
|
||||
EntryClass = HunyuanVideoPipeline
|
||||
@@ -0,0 +1 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
@@ -0,0 +1,196 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Hunyuan video diffusion pipeline implementation.
|
||||
|
||||
This module contains an implementation of the Hunyuan video diffusion pipeline
|
||||
using the modular pipeline architecture.
|
||||
"""
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines import ComposedPipelineBase, Req
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages import (
|
||||
ConditioningStage,
|
||||
DecodingStage,
|
||||
DenoisingStage,
|
||||
ImageEncodingStage,
|
||||
ImageVAEEncodingStage,
|
||||
InputValidationStage,
|
||||
LatentPreparationStage,
|
||||
TextEncodingStage,
|
||||
TimestepPreparationStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
# TODO(will): move PRECISION_TO_TYPE to better place
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def calculate_shift(
|
||||
image_seq_len,
|
||||
base_seq_len: int = 256,
|
||||
max_seq_len: int = 4096,
|
||||
base_shift: float = 0.5,
|
||||
max_shift: float = 1.15,
|
||||
):
|
||||
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
|
||||
b = base_shift - m * base_seq_len
|
||||
mu = image_seq_len * m + b
|
||||
return mu
|
||||
|
||||
|
||||
def prepare_mu(batch: Req, server_args: ServerArgs):
|
||||
height = batch.height
|
||||
width = batch.width
|
||||
vae_scale_factor = server_args.pipeline_config.vae_config.vae_scale_factor
|
||||
image_seq_len = (int(height) // vae_scale_factor) * (int(width) // vae_scale_factor)
|
||||
|
||||
mu = calculate_shift(
|
||||
image_seq_len,
|
||||
# hard code, since scheduler_config is not in PipelineConfig now
|
||||
256,
|
||||
4096,
|
||||
0.5,
|
||||
1.15,
|
||||
)
|
||||
return "mu", mu
|
||||
|
||||
|
||||
class QwenImagePipeline(ComposedPipelineBase):
|
||||
pipeline_name = "QwenImagePipeline"
|
||||
|
||||
_required_config_modules = [
|
||||
"text_encoder",
|
||||
"tokenizer",
|
||||
"vae",
|
||||
"transformer",
|
||||
"scheduler",
|
||||
]
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||
"""Set up pipeline stages with proper dependency injection."""
|
||||
|
||||
self.add_stage(
|
||||
stage_name="input_validation_stage", stage=InputValidationStage()
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage_primary",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[
|
||||
self.get_module("text_encoder"),
|
||||
],
|
||||
tokenizers=[
|
||||
self.get_module("tokenizer"),
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
|
||||
|
||||
self.add_stage(
|
||||
stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
prepare_extra_set_timesteps_kwargs=[prepare_mu],
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="latent_preparation_stage",
|
||||
stage=LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="denoising_stage",
|
||||
stage=DenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
|
||||
)
|
||||
|
||||
|
||||
class QwenImageEditPipeline(ComposedPipelineBase):
|
||||
pipeline_name = "QwenImageEditPipeline"
|
||||
|
||||
_required_config_modules = [
|
||||
"processor",
|
||||
"scheduler",
|
||||
"text_encoder",
|
||||
"tokenizer",
|
||||
"transformer",
|
||||
"vae",
|
||||
]
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||
"""Set up pipeline stages with proper dependency injection."""
|
||||
|
||||
self.add_stage(
|
||||
stage_name="input_validation_stage", stage=InputValidationStage()
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage_primary",
|
||||
stage=ImageEncodingStage(
|
||||
image_processor=self.get_module("processor"),
|
||||
text_encoder=self.get_module("text_encoder"),
|
||||
vae_image_processor=VaeImageProcessor(
|
||||
vae_scale_factor=server_args.pipeline_config.vae_config.arch_config.vae_scale_factor
|
||||
* 2
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="image_encoding_stage_primary",
|
||||
stage=ImageVAEEncodingStage(
|
||||
vae_image_processor=VaeImageProcessor(
|
||||
vae_scale_factor=server_args.pipeline_config.vae_config.arch_config.vae_scale_factor
|
||||
* 2
|
||||
),
|
||||
vae=self.get_module("vae"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
prepare_extra_set_timesteps_kwargs=[prepare_mu],
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="latent_preparation_stage",
|
||||
stage=LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
|
||||
|
||||
self.add_stage(
|
||||
stage_name="denoising_stage",
|
||||
stage=DenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
|
||||
)
|
||||
|
||||
|
||||
EntryClass = [QwenImagePipeline, QwenImageEditPipeline]
|
||||
@@ -0,0 +1 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
@@ -0,0 +1,182 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# type: ignore
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Hunyuan video diffusion pipeline implementation.
|
||||
|
||||
This module contains an implementation of the Hunyuan video diffusion pipeline
|
||||
using the modular pipeline architecture.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.loader.component_loader import (
|
||||
PipelineComponentLoader,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.encoders.bert import (
|
||||
HunyuanClip, # type: ignore
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.encoders.stepllm import STEP1TextEncoder
|
||||
from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import (
|
||||
ComposedPipelineBase,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines.lora_pipeline import LoRAPipeline
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages import (
|
||||
DecodingStage,
|
||||
DenoisingStage,
|
||||
InputValidationStage,
|
||||
LatentPreparationStage,
|
||||
StepvideoPromptEncodingStage,
|
||||
TimestepPreparationStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class StepVideoPipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
pipeline_name = "StepVideoPipeline"
|
||||
|
||||
_required_config_modules = ["transformer", "scheduler", "vae"]
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||
"""Set up pipeline stages with proper dependency injection."""
|
||||
|
||||
self.add_stage(
|
||||
stage_name="input_validation_stage", stage=InputValidationStage()
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage",
|
||||
stage=StepvideoPromptEncodingStage(
|
||||
stepllm=self.get_module("text_encoder"),
|
||||
clip=self.get_module("text_encoder_2"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="latent_preparation_stage",
|
||||
stage=LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="denoising_stage",
|
||||
stage=DenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
|
||||
)
|
||||
|
||||
def build_llm(self, model_dir, device) -> torch.nn.Module:
|
||||
text_encoder = (
|
||||
STEP1TextEncoder(model_dir, max_length=320).to(torch.bfloat16).eval()
|
||||
)
|
||||
return text_encoder
|
||||
|
||||
def build_clip(self, model_dir, device) -> HunyuanClip:
|
||||
clip = HunyuanClip(model_dir, max_length=77).eval()
|
||||
return clip
|
||||
|
||||
def initialize_pipeline(self, server_args: ServerArgs):
|
||||
"""
|
||||
Initialize the pipeline.
|
||||
"""
|
||||
target_device = get_local_torch_device()
|
||||
llm_dir = os.path.join(self.model_path, "step_llm")
|
||||
clip_dir = os.path.join(self.model_path, "hunyuan_clip")
|
||||
text_enc = self.build_llm(llm_dir, target_device)
|
||||
clip_enc = self.build_clip(clip_dir, target_device)
|
||||
self.add_module("text_encoder", text_enc)
|
||||
self.add_module("text_encoder_2", clip_enc)
|
||||
lib_path = (
|
||||
os.path.join(
|
||||
server_args.model_path,
|
||||
"lib/liboptimus_ths-torch2.5-cu124.cpython-310-x86_64-linux-gnu.so",
|
||||
)
|
||||
if os.path.isdir(server_args.model_path) # local checkout
|
||||
else hf_hub_download(
|
||||
repo_id=server_args.model_path,
|
||||
filename="lib/liboptimus_ths-torch2.5-cu124.cpython-310-x86_64-linux-gnu.so",
|
||||
)
|
||||
)
|
||||
torch.ops.load_library(lib_path)
|
||||
|
||||
def load_modules(
|
||||
self,
|
||||
server_args: ServerArgs,
|
||||
loaded_modules: dict[str, torch.nn.Module] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Load the modules from the config.
|
||||
"""
|
||||
model_index = self._load_config()
|
||||
logger.info("Loading pipeline modules from config: %s", model_index)
|
||||
|
||||
# remove keys that are not pipeline modules
|
||||
model_index.pop("_class_name")
|
||||
model_index.pop("_diffusers_version")
|
||||
|
||||
# some sanity checks
|
||||
assert (
|
||||
len(model_index) > 1
|
||||
), "model_index.json must contain at least one pipeline module"
|
||||
|
||||
required_modules = ["transformer", "scheduler", "vae"]
|
||||
for module_name in required_modules:
|
||||
if module_name not in model_index:
|
||||
raise ValueError(
|
||||
f"model_index.json must contain a {module_name} module"
|
||||
)
|
||||
logger.info("Diffusers config passed sanity checks")
|
||||
|
||||
# all the component models used by the pipeline
|
||||
modules = {}
|
||||
for module_name, (
|
||||
transformers_or_diffusers,
|
||||
architecture,
|
||||
) in model_index.items():
|
||||
component_model_path = os.path.join(self.model_path, module_name)
|
||||
module = PipelineComponentLoader.load_module(
|
||||
module_name=module_name,
|
||||
component_model_path=component_model_path,
|
||||
transformers_or_diffusers=transformers_or_diffusers,
|
||||
server_args=server_args,
|
||||
)
|
||||
logger.info("Loaded module %s from %s", module_name, component_model_path)
|
||||
|
||||
if module_name in modules:
|
||||
logger.warning("Overwriting module %s", module_name)
|
||||
modules[module_name] = module
|
||||
|
||||
required_modules = self.required_config_modules
|
||||
# Check if all required modules were loaded
|
||||
for module_name in required_modules:
|
||||
if module_name not in modules or modules[module_name] is None:
|
||||
raise ValueError(
|
||||
f"Required module {module_name} was not loaded properly"
|
||||
)
|
||||
|
||||
return modules
|
||||
|
||||
|
||||
EntryClass = StepVideoPipeline
|
||||
@@ -0,0 +1 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
@@ -0,0 +1,78 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Wan causal DMD pipeline implementation.
|
||||
|
||||
This module wires the causal DMD denoising stage into the modular pipeline.
|
||||
"""
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines import ComposedPipelineBase, LoRAPipeline
|
||||
|
||||
# isort: off
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages import (
|
||||
ConditioningStage,
|
||||
DecodingStage,
|
||||
CausalDMDDenoisingStage,
|
||||
InputValidationStage,
|
||||
LatentPreparationStage,
|
||||
TextEncodingStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
# isort: on
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
pipeline_name = "WanPipeline"
|
||||
|
||||
_required_config_modules = [
|
||||
"text_encoder",
|
||||
"tokenizer",
|
||||
"vae",
|
||||
"transformer",
|
||||
"scheduler",
|
||||
]
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs) -> None:
|
||||
"""Set up pipeline stages with proper dependency injection."""
|
||||
|
||||
self.add_stage(
|
||||
stage_name="input_validation_stage", stage=InputValidationStage()
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[self.get_module("text_encoder")],
|
||||
tokenizers=[self.get_module("tokenizer")],
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
|
||||
|
||||
self.add_stage(
|
||||
stage_name="latent_preparation_stage",
|
||||
stage=LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer", None),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="denoising_stage",
|
||||
stage=CausalDMDDenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
|
||||
)
|
||||
|
||||
|
||||
EntryClass = WanCausalDMDPipeline
|
||||
@@ -0,0 +1,98 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Wan video diffusion pipeline implementation.
|
||||
|
||||
This module contains an implementation of the Wan video diffusion pipeline
|
||||
using the modular pipeline architecture.
|
||||
"""
|
||||
|
||||
from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import (
|
||||
FlowMatchEulerDiscreteScheduler,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines import ComposedPipelineBase, LoRAPipeline
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
# isort: off
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages import (
|
||||
ConditioningStage,
|
||||
DecodingStage,
|
||||
DmdDenoisingStage,
|
||||
InputValidationStage,
|
||||
LatentPreparationStage,
|
||||
TextEncodingStage,
|
||||
TimestepPreparationStage,
|
||||
)
|
||||
|
||||
# isort: on
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class WanDMDPipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
"""
|
||||
Wan video diffusion pipeline with LoRA support.
|
||||
"""
|
||||
|
||||
pipeline_name = "WanDMDPipeline"
|
||||
|
||||
_required_config_modules = [
|
||||
"text_encoder",
|
||||
"tokenizer",
|
||||
"vae",
|
||||
"transformer",
|
||||
"scheduler",
|
||||
]
|
||||
|
||||
def initialize_pipeline(self, server_args: ServerArgs):
|
||||
|
||||
self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
|
||||
shift=server_args.pipeline_config.flow_shift
|
||||
)
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs) -> None:
|
||||
"""Set up pipeline stages with proper dependency injection."""
|
||||
|
||||
self.add_stage(
|
||||
stage_name="input_validation_stage", stage=InputValidationStage()
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[self.get_module("text_encoder")],
|
||||
tokenizers=[self.get_module("tokenizer")],
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
|
||||
|
||||
self.add_stage(
|
||||
stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="latent_preparation_stage",
|
||||
stage=LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer", None),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="denoising_stage",
|
||||
stage=DmdDenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
|
||||
)
|
||||
|
||||
|
||||
EntryClass = WanDMDPipeline
|
||||
@@ -0,0 +1,113 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Wan video diffusion pipeline implementation.
|
||||
|
||||
This module contains an implementation of the Wan video diffusion pipeline
|
||||
using the modular pipeline architecture.
|
||||
"""
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import (
|
||||
ComposedPipelineBase,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines.lora_pipeline import LoRAPipeline
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
# isort: off
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages import (
|
||||
ImageEncodingStage,
|
||||
ConditioningStage,
|
||||
DecodingStage,
|
||||
DmdDenoisingStage,
|
||||
ImageVAEEncodingStage,
|
||||
InputValidationStage,
|
||||
LatentPreparationStage,
|
||||
TextEncodingStage,
|
||||
TimestepPreparationStage,
|
||||
)
|
||||
|
||||
# isort: on
|
||||
from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import (
|
||||
FlowMatchEulerDiscreteScheduler,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class WanImageToVideoDmdPipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
pipeline_name = "WanCausalDMDPipeline"
|
||||
|
||||
_required_config_modules = [
|
||||
"text_encoder",
|
||||
"tokenizer",
|
||||
"vae",
|
||||
"transformer",
|
||||
"scheduler",
|
||||
"image_encoder",
|
||||
"image_processor",
|
||||
]
|
||||
|
||||
def initialize_pipeline(self, server_args: ServerArgs):
|
||||
self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
|
||||
shift=server_args.pipeline_config.flow_shift
|
||||
)
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||
"""Set up pipeline stages with proper dependency injection."""
|
||||
|
||||
self.add_stage(
|
||||
stage_name="input_validation_stage", stage=InputValidationStage()
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[self.get_module("text_encoder")],
|
||||
tokenizers=[self.get_module("tokenizer")],
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="image_encoding_stage",
|
||||
stage=ImageEncodingStage(
|
||||
image_encoder=self.get_module("image_encoder"),
|
||||
image_processor=self.get_module("image_processor"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
|
||||
|
||||
self.add_stage(
|
||||
stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="latent_preparation_stage",
|
||||
stage=LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="image_latent_preparation_stage",
|
||||
stage=ImageVAEEncodingStage(vae=self.get_module("vae")),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="denoising_stage",
|
||||
stage=DmdDenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
|
||||
)
|
||||
|
||||
|
||||
EntryClass = WanImageToVideoDmdPipeline
|
||||
@@ -0,0 +1,118 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Wan video diffusion pipeline implementation.
|
||||
|
||||
This module contains an implementation of the Wan video diffusion pipeline
|
||||
using the modular pipeline architecture.
|
||||
"""
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import (
|
||||
ComposedPipelineBase,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines.lora_pipeline import LoRAPipeline
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
# isort: off
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages import (
|
||||
ImageEncodingStage,
|
||||
ConditioningStage,
|
||||
DecodingStage,
|
||||
DenoisingStage,
|
||||
ImageVAEEncodingStage,
|
||||
InputValidationStage,
|
||||
LatentPreparationStage,
|
||||
TextEncodingStage,
|
||||
TimestepPreparationStage,
|
||||
)
|
||||
|
||||
# isort: on
|
||||
from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_unipc_multistep import (
|
||||
FlowUniPCMultistepScheduler,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class WanImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
pipeline_name = "WanImageToVideoPipeline"
|
||||
|
||||
_required_config_modules = [
|
||||
"text_encoder",
|
||||
"tokenizer",
|
||||
"vae",
|
||||
"transformer",
|
||||
"scheduler",
|
||||
"image_encoder",
|
||||
"image_processor",
|
||||
]
|
||||
|
||||
def initialize_pipeline(self, server_args: ServerArgs):
|
||||
self.modules["scheduler"] = FlowUniPCMultistepScheduler(
|
||||
shift=server_args.pipeline_config.flow_shift
|
||||
)
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||
"""Set up pipeline stages with proper dependency injection."""
|
||||
|
||||
self.add_stage(
|
||||
stage_name="input_validation_stage", stage=InputValidationStage()
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[self.get_module("text_encoder")],
|
||||
tokenizers=[self.get_module("tokenizer")],
|
||||
),
|
||||
)
|
||||
|
||||
if (
|
||||
self.get_module("image_encoder") is not None
|
||||
and self.get_module("image_processor") is not None
|
||||
):
|
||||
self.add_stage(
|
||||
stage_name="image_encoding_stage",
|
||||
stage=ImageEncodingStage(
|
||||
image_encoder=self.get_module("image_encoder"),
|
||||
image_processor=self.get_module("image_processor"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
|
||||
|
||||
self.add_stage(
|
||||
stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="latent_preparation_stage",
|
||||
stage=LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="image_latent_preparation_stage",
|
||||
stage=ImageVAEEncodingStage(vae=self.get_module("vae")),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="denoising_stage",
|
||||
stage=DenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
transformer_2=self.get_module("transformer_2"),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
|
||||
)
|
||||
|
||||
|
||||
EntryClass = WanImageToVideoPipeline
|
||||
@@ -0,0 +1,98 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Wan video diffusion pipeline implementation.
|
||||
|
||||
This module contains an implementation of the Wan video diffusion pipeline
|
||||
using the modular pipeline architecture.
|
||||
"""
|
||||
|
||||
from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_unipc_multistep import (
|
||||
FlowUniPCMultistepScheduler,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines import ComposedPipelineBase, LoRAPipeline
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages import (
|
||||
ConditioningStage,
|
||||
DecodingStage,
|
||||
DenoisingStage,
|
||||
InputValidationStage,
|
||||
LatentPreparationStage,
|
||||
TextEncodingStage,
|
||||
TimestepPreparationStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class WanPipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
"""
|
||||
Wan video diffusion pipeline with LoRA support.
|
||||
"""
|
||||
|
||||
pipeline_name = "WanImageToVideoPipeline"
|
||||
|
||||
_required_config_modules = [
|
||||
"text_encoder",
|
||||
"tokenizer",
|
||||
"vae",
|
||||
"transformer",
|
||||
"scheduler",
|
||||
]
|
||||
|
||||
def initialize_pipeline(self, server_args: ServerArgs):
|
||||
# We use UniPCMScheduler from Wan2.1 official repo, not the one in diffusers.
|
||||
self.modules["scheduler"] = FlowUniPCMultistepScheduler(
|
||||
shift=server_args.pipeline_config.flow_shift
|
||||
)
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs) -> None:
|
||||
"""Set up pipeline stages with proper dependency injection."""
|
||||
|
||||
self.add_stage(
|
||||
stage_name="input_validation_stage", stage=InputValidationStage()
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[self.get_module("text_encoder")],
|
||||
tokenizers=[self.get_module("tokenizer")],
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
|
||||
|
||||
self.add_stage(
|
||||
stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="latent_preparation_stage",
|
||||
stage=LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer", None),
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="denoising_stage",
|
||||
stage=DenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
transformer_2=self.get_module("transformer_2", None),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
vae=self.get_module("vae"),
|
||||
pipeline=self,
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="decoding_stage",
|
||||
stage=DecodingStage(vae=self.get_module("vae"), pipeline=self),
|
||||
)
|
||||
|
||||
|
||||
EntryClass = WanPipeline
|
||||
@@ -0,0 +1 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
@@ -0,0 +1,433 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm import tqdm
|
||||
|
||||
from sglang.multimodal_gen.dataset import getdataset
|
||||
from sglang.multimodal_gen.dataset.dataloader.parquet_io import (
|
||||
ParquetDatasetWriter,
|
||||
records_to_table,
|
||||
)
|
||||
from sglang.multimodal_gen.dataset.preprocessing_datasets import PreprocessBatch
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import (
|
||||
ComposedPipelineBase,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages import TextEncodingStage
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class BasePreprocessPipeline(ComposedPipelineBase):
|
||||
"""Base class for preprocessing pipelines that handles common functionality."""
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||
"""Set up pipeline stages with proper dependency injection."""
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[self.get_module("text_encoder")],
|
||||
tokenizers=[self.get_module("tokenizer")],
|
||||
),
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(
|
||||
self,
|
||||
batch: Req,
|
||||
server_args: ServerArgs,
|
||||
args,
|
||||
):
|
||||
if not self.post_init_called:
|
||||
self.post_init()
|
||||
|
||||
# Initialize class variables for data sharing
|
||||
self.video_data: dict[str, Any] = {} # Store video metadata and paths
|
||||
self.latent_data: dict[str, Any] = {} # Store latent tensors
|
||||
self.preprocess_video_and_text(server_args, args)
|
||||
|
||||
def get_extra_features(
|
||||
self, valid_data: dict[str, Any], server_args: ServerArgs
|
||||
) -> dict[str, Any]:
|
||||
"""Get additional features specific to the pipeline type. Override in subclasses."""
|
||||
return {}
|
||||
|
||||
def get_pyarrow_schema(self) -> pa.Schema:
|
||||
"""Return the PyArrow schema for this pipeline. Must be overridden."""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_schema_fields(self) -> list[str]:
|
||||
"""Get the schema fields for the pipeline type."""
|
||||
return [f.name for f in self.get_pyarrow_schema()]
|
||||
|
||||
def create_record_for_schema(
|
||||
self, preprocess_batch: PreprocessBatch, schema: pa.Schema, strict: bool = False
|
||||
) -> dict[str, Any]:
|
||||
"""Create a record for the Parquet dataset using a generic schema-based approach.
|
||||
|
||||
Args:
|
||||
preprocess_batch: The batch containing the data to extract
|
||||
schema: PyArrow schema defining the expected fields
|
||||
strict: If True, raises an exception when required fields are missing or unfilled
|
||||
|
||||
Returns:
|
||||
Dictionary record matching the schema
|
||||
|
||||
Raises:
|
||||
ValueError: If strict=True and required fields are missing or unfilled
|
||||
"""
|
||||
record = {}
|
||||
unfilled_fields = []
|
||||
|
||||
for field in schema.names:
|
||||
field_filled = False
|
||||
|
||||
if field.endswith("_bytes"):
|
||||
# Handle binary tensor data - convert numpy array or tensor to bytes
|
||||
tensor_name = field.replace("_bytes", "")
|
||||
tensor_data = getattr(preprocess_batch, tensor_name, None)
|
||||
if tensor_data is not None:
|
||||
try:
|
||||
if hasattr(tensor_data, "numpy"): # torch tensor
|
||||
record[field] = tensor_data.cpu().numpy().tobytes()
|
||||
field_filled = True
|
||||
elif hasattr(tensor_data, "tobytes"): # numpy array
|
||||
record[field] = tensor_data.tobytes()
|
||||
field_filled = True
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported tensor type for field {field}: {type(tensor_data)}"
|
||||
)
|
||||
except Exception as e:
|
||||
if strict:
|
||||
raise ValueError(
|
||||
f"Failed to convert tensor {tensor_name} to bytes: {e}"
|
||||
) from e
|
||||
record[field] = b"" # Empty bytes for missing data
|
||||
else:
|
||||
record[field] = b"" # Empty bytes for missing data
|
||||
|
||||
elif field.endswith("_shape"):
|
||||
# Handle tensor shape info
|
||||
tensor_name = field.replace("_shape", "")
|
||||
tensor_data = getattr(preprocess_batch, tensor_name, None)
|
||||
if tensor_data is not None and hasattr(tensor_data, "shape"):
|
||||
record[field] = list(tensor_data.shape)
|
||||
field_filled = True
|
||||
else:
|
||||
record[field] = []
|
||||
|
||||
elif field.endswith("_dtype"):
|
||||
# Handle tensor dtype info
|
||||
tensor_name = field.replace("_dtype", "")
|
||||
tensor_data = getattr(preprocess_batch, tensor_name, None)
|
||||
if tensor_data is not None and hasattr(tensor_data, "dtype"):
|
||||
record[field] = str(tensor_data.dtype)
|
||||
field_filled = True
|
||||
else:
|
||||
record[field] = "unknown"
|
||||
|
||||
elif field in ["width", "height", "num_frames"]:
|
||||
# Handle integer metadata fields
|
||||
value = getattr(preprocess_batch, field, None)
|
||||
if value is not None:
|
||||
try:
|
||||
record[field] = int(value)
|
||||
field_filled = True
|
||||
except (ValueError, TypeError) as e:
|
||||
if strict:
|
||||
raise ValueError(
|
||||
f"Failed to convert field {field} to int: {e}"
|
||||
) from e
|
||||
record[field] = 0
|
||||
else:
|
||||
record[field] = 0
|
||||
|
||||
elif field in ["duration_sec", "fps"]:
|
||||
# Handle float metadata fields
|
||||
# Map schema field names to batch attribute names
|
||||
attr_name = "duration" if field == "duration_sec" else field
|
||||
value = getattr(preprocess_batch, attr_name, None)
|
||||
if value is not None:
|
||||
try:
|
||||
record[field] = float(value)
|
||||
field_filled = True
|
||||
except (ValueError, TypeError) as e:
|
||||
if strict:
|
||||
raise ValueError(
|
||||
f"Failed to convert field {field} to float: {e}"
|
||||
) from e
|
||||
record[field] = 0.0
|
||||
else:
|
||||
record[field] = 0.0
|
||||
|
||||
else:
|
||||
# Handle string fields (id, file_name, caption, media_type, etc.)
|
||||
# Map common schema field names to batch attribute names
|
||||
attr_name = field
|
||||
if field == "caption":
|
||||
attr_name = "text"
|
||||
elif field == "file_name":
|
||||
attr_name = "path"
|
||||
elif field == "id":
|
||||
# Generate ID from path if available
|
||||
path_value = getattr(preprocess_batch, "path", None)
|
||||
if path_value:
|
||||
import os
|
||||
|
||||
record[field] = os.path.basename(path_value).split(".")[0]
|
||||
field_filled = True
|
||||
else:
|
||||
record[field] = ""
|
||||
continue
|
||||
elif field == "media_type":
|
||||
# Determine media type from path
|
||||
path_value = getattr(preprocess_batch, "path", None)
|
||||
if path_value:
|
||||
record[field] = (
|
||||
"video" if path_value.endswith(".mp4") else "image"
|
||||
)
|
||||
field_filled = True
|
||||
else:
|
||||
record[field] = ""
|
||||
continue
|
||||
|
||||
value = getattr(preprocess_batch, attr_name, None)
|
||||
if value is not None:
|
||||
record[field] = str(value)
|
||||
field_filled = True
|
||||
else:
|
||||
record[field] = ""
|
||||
|
||||
# Track unfilled fields
|
||||
if not field_filled:
|
||||
unfilled_fields.append(field)
|
||||
|
||||
# Handle strict mode
|
||||
if strict and unfilled_fields:
|
||||
raise ValueError(f"Required fields were not filled: {unfilled_fields}")
|
||||
|
||||
# Log unfilled fields as warning if not in strict mode
|
||||
if unfilled_fields:
|
||||
logger.warning(
|
||||
"Some fields were not filled and got default values: %s",
|
||||
unfilled_fields,
|
||||
)
|
||||
|
||||
return record
|
||||
|
||||
def create_record(
|
||||
self,
|
||||
video_name: str,
|
||||
vae_latent: np.ndarray,
|
||||
text_embedding: np.ndarray,
|
||||
valid_data: dict[str, Any],
|
||||
idx: int,
|
||||
extra_features: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Create a record for the Parquet dataset."""
|
||||
record = {
|
||||
"id": video_name,
|
||||
"vae_latent_bytes": vae_latent.tobytes(),
|
||||
"vae_latent_shape": list(vae_latent.shape),
|
||||
"vae_latent_dtype": str(vae_latent.dtype),
|
||||
"text_embedding_bytes": text_embedding.tobytes(),
|
||||
"text_embedding_shape": list(text_embedding.shape),
|
||||
"text_embedding_dtype": str(text_embedding.dtype),
|
||||
"file_name": video_name,
|
||||
"caption": valid_data["text"][idx] if len(valid_data["text"]) > 0 else "",
|
||||
"media_type": "video",
|
||||
"width": (
|
||||
valid_data["pixel_values"][idx].shape[-2]
|
||||
if len(valid_data["pixel_values"]) > 0
|
||||
else 0
|
||||
),
|
||||
"height": (
|
||||
valid_data["pixel_values"][idx].shape[-1]
|
||||
if len(valid_data["pixel_values"]) > 0
|
||||
else 0
|
||||
),
|
||||
"num_frames": vae_latent.shape[1] if len(vae_latent.shape) > 1 else 0,
|
||||
"duration_sec": (
|
||||
float(valid_data["duration"][idx])
|
||||
if len(valid_data["duration"]) > 0
|
||||
else 0.0
|
||||
),
|
||||
"fps": float(valid_data["fps"][idx]) if len(valid_data["fps"]) > 0 else 0.0,
|
||||
}
|
||||
if extra_features:
|
||||
record.update(extra_features)
|
||||
return record
|
||||
|
||||
def preprocess_video_and_text(self, server_args: ServerArgs, args):
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
# Create directory for combined data
|
||||
combined_parquet_dir = os.path.join(args.output_dir, "combined_parquet_dataset")
|
||||
os.makedirs(combined_parquet_dir, exist_ok=True)
|
||||
local_rank = int(os.getenv("RANK", 0))
|
||||
|
||||
# Get how many samples have already been processed
|
||||
start_idx = 0
|
||||
for root, _, files in os.walk(combined_parquet_dir):
|
||||
for file in files:
|
||||
if file.endswith(".parquet"):
|
||||
table = pq.read_table(os.path.join(root, file))
|
||||
start_idx += table.num_rows
|
||||
|
||||
# Loading dataset
|
||||
train_dataset = getdataset(args)
|
||||
|
||||
train_dataloader = DataLoader(
|
||||
train_dataset,
|
||||
batch_size=args.preprocess_video_batch_size,
|
||||
num_workers=args.dataloader_num_workers,
|
||||
)
|
||||
|
||||
num_processed_samples = 0
|
||||
# Add progress bar for video preprocessing
|
||||
pbar = tqdm(
|
||||
train_dataloader,
|
||||
desc="Processing videos",
|
||||
unit="batch",
|
||||
disable=local_rank != 0,
|
||||
)
|
||||
|
||||
for batch_idx, data in enumerate(pbar):
|
||||
if data is None:
|
||||
continue
|
||||
|
||||
with torch.inference_mode():
|
||||
# Filter out invalid samples (those with all zeros)
|
||||
valid_indices = []
|
||||
for i, pixel_values in enumerate(data["pixel_values"]):
|
||||
if not torch.all(pixel_values == 0): # Check if all values are zero
|
||||
valid_indices.append(i)
|
||||
num_processed_samples += len(valid_indices)
|
||||
|
||||
if not valid_indices:
|
||||
continue
|
||||
|
||||
# Create new batch with only valid samples
|
||||
valid_data = {
|
||||
"pixel_values": torch.stack(
|
||||
[data["pixel_values"][i] for i in valid_indices]
|
||||
),
|
||||
"text": [data["text"][i] for i in valid_indices],
|
||||
"path": [data["path"][i] for i in valid_indices],
|
||||
"fps": [data["fps"][i] for i in valid_indices],
|
||||
"duration": [data["duration"][i] for i in valid_indices],
|
||||
}
|
||||
|
||||
# VAE
|
||||
with torch.autocast("cuda", dtype=torch.float32):
|
||||
latents = (
|
||||
self.get_module("vae")
|
||||
.encode(valid_data["pixel_values"].to(get_local_torch_device()))
|
||||
.mean
|
||||
)
|
||||
|
||||
# Get extra features if needed
|
||||
extra_features = self.get_extra_features(valid_data, server_args)
|
||||
|
||||
batch_captions = valid_data["text"]
|
||||
batch = Req(
|
||||
data_type="video",
|
||||
prompt=batch_captions,
|
||||
prompt_embeds=[],
|
||||
prompt_attention_mask=[],
|
||||
)
|
||||
assert hasattr(self, "prompt_encoding_stage")
|
||||
result_batch = self.prompt_encoding_stage(batch, server_args)
|
||||
prompt_embeds, prompt_attention_mask = (
|
||||
result_batch.prompt_embeds[0],
|
||||
result_batch.prompt_attention_mask[0],
|
||||
)
|
||||
assert prompt_embeds.shape[0] == prompt_attention_mask.shape[0]
|
||||
|
||||
# Get sequence lengths from attention masks (number of 1s)
|
||||
seq_lens = prompt_attention_mask.sum(dim=1)
|
||||
|
||||
non_padded_embeds = []
|
||||
non_padded_masks = []
|
||||
|
||||
# Process each item in the batch
|
||||
for i in range(prompt_embeds.size(0)):
|
||||
seq_len = seq_lens[i].item()
|
||||
# Slice the embeddings and masks to keep only non-padding parts
|
||||
non_padded_embeds.append(prompt_embeds[i, :seq_len])
|
||||
non_padded_masks.append(prompt_attention_mask[i, :seq_len])
|
||||
|
||||
# Update the tensors with non-padded versions
|
||||
prompt_embeds = non_padded_embeds
|
||||
prompt_attention_mask = non_padded_masks
|
||||
|
||||
# Prepare batch data for Parquet dataset
|
||||
batch_data = []
|
||||
|
||||
# Add progress bar for saving outputs
|
||||
save_pbar = tqdm(
|
||||
enumerate(valid_data["path"]),
|
||||
desc="Saving outputs",
|
||||
unit="item",
|
||||
leave=False,
|
||||
)
|
||||
for idx, video_path in save_pbar:
|
||||
# Get the corresponding latent and info using video name
|
||||
latent = latents[idx].cpu()
|
||||
video_name = os.path.basename(video_path).split(".")[0]
|
||||
|
||||
# Convert tensors to numpy arrays
|
||||
vae_latent = latent.cpu().numpy()
|
||||
text_embedding = prompt_embeds[idx].cpu().numpy()
|
||||
|
||||
# Get extra features for this sample if needed
|
||||
sample_extra_features = {}
|
||||
if extra_features:
|
||||
for key, value in extra_features.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
sample_extra_features[key] = value[idx].cpu().numpy()
|
||||
else:
|
||||
sample_extra_features[key] = value[idx]
|
||||
|
||||
# Create record for Parquet dataset
|
||||
record = self.create_record(
|
||||
video_name=video_name,
|
||||
vae_latent=vae_latent,
|
||||
text_embedding=text_embedding,
|
||||
valid_data=valid_data,
|
||||
idx=idx,
|
||||
extra_features=sample_extra_features,
|
||||
)
|
||||
batch_data.append(record)
|
||||
|
||||
if batch_data:
|
||||
write_pbar = tqdm(
|
||||
total=1, desc="Writing to Parquet dataset", unit="batch"
|
||||
)
|
||||
table = records_to_table(batch_data, self.get_pyarrow_schema())
|
||||
write_pbar.update(1)
|
||||
write_pbar.close()
|
||||
|
||||
if not hasattr(self, "dataset_writer"):
|
||||
self.dataset_writer = ParquetDatasetWriter(
|
||||
out_dir=combined_parquet_dir,
|
||||
samples_per_file=args.samples_per_file,
|
||||
)
|
||||
self.dataset_writer.append_table(table)
|
||||
logger.info("Collected batch with %s samples", len(table))
|
||||
|
||||
if num_processed_samples >= args.flush_frequency:
|
||||
written = self.dataset_writer.flush()
|
||||
logger.info("Flushed %s samples to parquet", written)
|
||||
num_processed_samples = 0
|
||||
@@ -0,0 +1,247 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
I2V Data Preprocessing pipeline implementation.
|
||||
|
||||
This module contains an implementation of the I2V Data Preprocessing pipeline
|
||||
using the modular pipeline architecture.
|
||||
"""
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from sglang.multimodal_gen.dataset.dataloader.schema import pyarrow_schema_i2v
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
|
||||
from sglang.multimodal_gen.runtime.pipelines.preprocess.preprocess_pipeline_base import (
|
||||
BasePreprocessPipeline,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages import (
|
||||
ImageEncodingStage,
|
||||
TextEncodingStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
|
||||
|
||||
class PreprocessPipeline_I2V(BasePreprocessPipeline):
|
||||
"""I2V preprocessing pipeline implementation."""
|
||||
|
||||
_required_config_modules = [
|
||||
"text_encoder",
|
||||
"tokenizer",
|
||||
"vae",
|
||||
"image_encoder",
|
||||
"image_processor",
|
||||
]
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[self.get_module("text_encoder")],
|
||||
tokenizers=[self.get_module("tokenizer")],
|
||||
),
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="image_encoding_stage",
|
||||
stage=ImageEncodingStage(
|
||||
image_encoder=self.get_module("image_encoder"),
|
||||
image_processor=self.get_module("image_processor"),
|
||||
),
|
||||
)
|
||||
|
||||
def get_pyarrow_schema(self):
|
||||
"""Return the PyArrow schema for I2V pipeline."""
|
||||
return pyarrow_schema_i2v
|
||||
|
||||
def get_extra_features(
|
||||
self, valid_data: dict[str, Any], server_args: ServerArgs
|
||||
) -> dict[str, Any]:
|
||||
|
||||
# TODO(will): move these to cpu at some point
|
||||
self.get_module("image_encoder").to(get_local_torch_device())
|
||||
self.get_module("vae").to(get_local_torch_device())
|
||||
|
||||
features = {}
|
||||
"""Get CLIP features from the first frame of each video."""
|
||||
first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute(
|
||||
0, 2, 3, 1
|
||||
) # (B, C, T, H, W) -> (B, H, W, C)
|
||||
_, _, num_frames, height, width = valid_data["pixel_values"].shape
|
||||
# latent_height = height // self.get_module(
|
||||
# "vae").spatial_compression_ratio
|
||||
# latent_width = width // self.get_module("vae").spatial_compression_ratio
|
||||
|
||||
processed_images = []
|
||||
# Frame has values between -1 and 1
|
||||
for frame in first_frame:
|
||||
frame = (frame + 1) * 127.5
|
||||
frame_pil = Image.fromarray(frame.cpu().numpy().astype(np.uint8))
|
||||
processed_img = self.get_module("image_processor")(
|
||||
images=frame_pil, return_tensors="pt"
|
||||
)
|
||||
processed_images.append(processed_img)
|
||||
|
||||
# Get CLIP features
|
||||
pixel_values = torch.cat(
|
||||
[img["pixel_values"] for img in processed_images], dim=0
|
||||
).to(get_local_torch_device())
|
||||
with torch.no_grad():
|
||||
image_inputs = {"pixel_values": pixel_values}
|
||||
with set_forward_context(current_timestep=0, attn_metadata=None):
|
||||
clip_features = self.get_module("image_encoder")(**image_inputs)
|
||||
clip_features = clip_features.last_hidden_state
|
||||
|
||||
features["clip_feature"] = clip_features
|
||||
"""Get VAE features from the first frame of each video"""
|
||||
video_conditions = []
|
||||
for frame in first_frame:
|
||||
processed_img = frame.to(device="cpu", dtype=torch.float32)
|
||||
processed_img = processed_img.unsqueeze(0).permute(0, 3, 1, 2).unsqueeze(2)
|
||||
# (B, H, W, C) -> (B, C, 1, H, W)
|
||||
video_condition = torch.cat(
|
||||
[
|
||||
processed_img,
|
||||
processed_img.new_zeros(
|
||||
processed_img.shape[0],
|
||||
processed_img.shape[1],
|
||||
num_frames - 1,
|
||||
height,
|
||||
width,
|
||||
),
|
||||
],
|
||||
dim=2,
|
||||
)
|
||||
video_condition = video_condition.to(
|
||||
device=get_local_torch_device(), dtype=torch.float32
|
||||
)
|
||||
video_conditions.append(video_condition)
|
||||
|
||||
video_conditions = torch.cat(video_conditions, dim=0)
|
||||
|
||||
with torch.autocast(device_type="cuda", dtype=torch.float32, enabled=True):
|
||||
encoder_outputs = self.get_module("vae").encode(video_conditions)
|
||||
|
||||
latent_condition = encoder_outputs.mean
|
||||
if (
|
||||
hasattr(self.get_module("vae"), "shift_factor")
|
||||
and self.get_module("vae").shift_factor is not None
|
||||
):
|
||||
if isinstance(self.get_module("vae").shift_factor, torch.Tensor):
|
||||
latent_condition -= self.get_module("vae").shift_factor.to(
|
||||
latent_condition.device, latent_condition.dtype
|
||||
)
|
||||
else:
|
||||
latent_condition -= self.get_module("vae").shift_factor
|
||||
|
||||
if isinstance(self.get_module("vae").scaling_factor, torch.Tensor):
|
||||
latent_condition = latent_condition * self.get_module(
|
||||
"vae"
|
||||
).scaling_factor.to(latent_condition.device, latent_condition.dtype)
|
||||
else:
|
||||
latent_condition = latent_condition * self.get_module("vae").scaling_factor
|
||||
|
||||
# mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height,
|
||||
# latent_width)
|
||||
# mask_lat_size[:, :, list(range(1, num_frames))] = 0
|
||||
# first_frame_mask = mask_lat_size[:, :, 0:1]
|
||||
# first_frame_mask = torch.repeat_interleave(
|
||||
# first_frame_mask,
|
||||
# dim=2,
|
||||
# repeats=self.get_module("vae").temporal_compression_ratio)
|
||||
# mask_lat_size = torch.concat(
|
||||
# [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
|
||||
# mask_lat_size = mask_lat_size.view(
|
||||
# batch_size, -1,
|
||||
# self.get_module("vae").temporal_compression_ratio, latent_height,
|
||||
# latent_width)
|
||||
# mask_lat_size = mask_lat_size.transpose(1, 2)
|
||||
# mask_lat_size = mask_lat_size.to(latent_condition.device)
|
||||
|
||||
# image_latent = torch.concat([mask_lat_size, latent_condition], dim=1)
|
||||
|
||||
features["first_frame_latent"] = latent_condition
|
||||
|
||||
return features
|
||||
|
||||
def create_record(
|
||||
self,
|
||||
video_name: str,
|
||||
vae_latent: np.ndarray,
|
||||
text_embedding: np.ndarray,
|
||||
valid_data: dict[str, Any],
|
||||
idx: int,
|
||||
extra_features: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Create a record for the Parquet dataset with CLIP features."""
|
||||
record = super().create_record(
|
||||
video_name=video_name,
|
||||
vae_latent=vae_latent,
|
||||
text_embedding=text_embedding,
|
||||
valid_data=valid_data,
|
||||
idx=idx,
|
||||
extra_features=extra_features,
|
||||
)
|
||||
|
||||
if extra_features and "clip_feature" in extra_features:
|
||||
clip_feature = extra_features["clip_feature"]
|
||||
record.update(
|
||||
{
|
||||
"clip_feature_bytes": clip_feature.tobytes(),
|
||||
"clip_feature_shape": list(clip_feature.shape),
|
||||
"clip_feature_dtype": str(clip_feature.dtype),
|
||||
}
|
||||
)
|
||||
else:
|
||||
record.update(
|
||||
{
|
||||
"clip_feature_bytes": b"",
|
||||
"clip_feature_shape": [],
|
||||
"clip_feature_dtype": "",
|
||||
}
|
||||
)
|
||||
|
||||
if extra_features and "first_frame_latent" in extra_features:
|
||||
first_frame_latent = extra_features["first_frame_latent"]
|
||||
record.update(
|
||||
{
|
||||
"first_frame_latent_bytes": first_frame_latent.tobytes(),
|
||||
"first_frame_latent_shape": list(first_frame_latent.shape),
|
||||
"first_frame_latent_dtype": str(first_frame_latent.dtype),
|
||||
}
|
||||
)
|
||||
else:
|
||||
record.update(
|
||||
{
|
||||
"first_frame_latent_bytes": b"",
|
||||
"first_frame_latent_shape": [],
|
||||
"first_frame_latent_dtype": "",
|
||||
}
|
||||
)
|
||||
|
||||
if extra_features and "pil_image" in extra_features:
|
||||
pil_image = extra_features["pil_image"]
|
||||
record.update(
|
||||
{
|
||||
"pil_image_bytes": pil_image.tobytes(),
|
||||
"pil_image_shape": list(pil_image.shape),
|
||||
"pil_image_dtype": str(pil_image.dtype),
|
||||
}
|
||||
)
|
||||
else:
|
||||
record.update(
|
||||
{
|
||||
"pil_image_bytes": b"",
|
||||
"pil_image_shape": [],
|
||||
"pil_image_dtype": "",
|
||||
}
|
||||
)
|
||||
|
||||
return record
|
||||
|
||||
|
||||
EntryClass = PreprocessPipeline_I2V
|
||||
@@ -0,0 +1,355 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
ODE Trajectory Data Preprocessing pipeline implementation.
|
||||
|
||||
This module contains an implementation of the ODE Trajectory Data Preprocessing pipeline
|
||||
using the modular pipeline architecture.
|
||||
|
||||
Sec 4.3 of CausVid paper: https://arxiv.org/pdf/2412.07772
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
|
||||
import pyarrow as pa
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from torchdata.stateful_dataloader import StatefulDataLoader
|
||||
from tqdm import tqdm
|
||||
|
||||
from sglang.multimodal_gen.configs.sample import SamplingParams
|
||||
from sglang.multimodal_gen.dataset import gettextdataset
|
||||
from sglang.multimodal_gen.dataset.dataloader.parquet_io import (
|
||||
ParquetDatasetWriter,
|
||||
records_to_table,
|
||||
)
|
||||
from sglang.multimodal_gen.dataset.dataloader.record_schema import (
|
||||
ode_text_only_record_creator,
|
||||
)
|
||||
from sglang.multimodal_gen.dataset.dataloader.schema import (
|
||||
pyarrow_schema_ode_trajectory_text_only,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.schedulers.scheduling_self_forcing_flow_match import (
|
||||
SelfForcingFlowMatchScheduler,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req
|
||||
from sglang.multimodal_gen.runtime.pipelines.preprocess.preprocess_pipeline_base import (
|
||||
BasePreprocessPipeline,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages import (
|
||||
DecodingStage,
|
||||
DenoisingStage,
|
||||
InputValidationStage,
|
||||
LatentPreparationStage,
|
||||
TextEncodingStage,
|
||||
TimestepPreparationStage,
|
||||
)
|
||||
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 save_decoded_latents_as_video, shallow_asdict
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
|
||||
"""ODE Trajectory preprocessing pipeline implementation."""
|
||||
|
||||
_required_config_modules = [
|
||||
"text_encoder",
|
||||
"tokenizer",
|
||||
"vae",
|
||||
"transformer",
|
||||
"scheduler",
|
||||
]
|
||||
preprocess_dataloader: StatefulDataLoader
|
||||
preprocess_loader_iter: Iterator[dict[str, Any]]
|
||||
pbar: Any
|
||||
num_processed_samples: int
|
||||
|
||||
def get_pyarrow_schema(self) -> pa.Schema:
|
||||
"""Return the PyArrow schema for ODE Trajectory pipeline."""
|
||||
return pyarrow_schema_ode_trajectory_text_only
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||
"""Set up pipeline stages with proper dependency injection."""
|
||||
assert server_args.pipeline_config.flow_shift == 5
|
||||
self.modules["scheduler"] = SelfForcingFlowMatchScheduler(
|
||||
shift=server_args.pipeline_config.flow_shift,
|
||||
sigma_min=0.0,
|
||||
extra_one_step=True,
|
||||
)
|
||||
self.modules["scheduler"].set_timesteps(
|
||||
num_inference_steps=48, denoising_strength=1.0
|
||||
)
|
||||
|
||||
self.add_stage(
|
||||
stage_name="input_validation_stage", stage=InputValidationStage()
|
||||
)
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[self.get_module("text_encoder")],
|
||||
tokenizers=[self.get_module("tokenizer")],
|
||||
),
|
||||
)
|
||||
self.add_stage(
|
||||
stage_name="timestep_preparation_stage",
|
||||
stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
|
||||
)
|
||||
self.add_stage(
|
||||
stage_name="latent_preparation_stage",
|
||||
stage=LatentPreparationStage(
|
||||
scheduler=self.get_module("scheduler"),
|
||||
transformer=self.get_module("transformer", None),
|
||||
),
|
||||
)
|
||||
self.add_stage(
|
||||
stage_name="denoising_stage",
|
||||
stage=DenoisingStage(
|
||||
transformer=self.get_module("transformer"),
|
||||
scheduler=self.get_module("scheduler"),
|
||||
pipeline=self,
|
||||
),
|
||||
)
|
||||
self.add_stage(
|
||||
stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
|
||||
)
|
||||
|
||||
def preprocess_text_and_trajectory(self, server_args: ServerArgs, args):
|
||||
"""Preprocess text-only data and generate trajectory information."""
|
||||
|
||||
for batch_idx, data in enumerate(self.pbar):
|
||||
if data is None:
|
||||
continue
|
||||
|
||||
with torch.inference_mode():
|
||||
# For text-only processing, we only need text data
|
||||
# Filter out samples without text
|
||||
valid_indices = []
|
||||
for i, text in enumerate(data["text"]):
|
||||
if text and text.strip(): # Check if text is not empty
|
||||
valid_indices.append(i)
|
||||
self.num_processed_samples += len(valid_indices)
|
||||
|
||||
if not valid_indices:
|
||||
continue
|
||||
|
||||
# Create new batch with only valid samples (text-only)
|
||||
valid_data = {
|
||||
"text": [data["text"][i] for i in valid_indices],
|
||||
"path": [data["path"][i] for i in valid_indices],
|
||||
}
|
||||
|
||||
# Add fps and duration if available in data
|
||||
if "fps" in data:
|
||||
valid_data["fps"] = [data["fps"][i] for i in valid_indices]
|
||||
if "duration" in data:
|
||||
valid_data["duration"] = [
|
||||
data["duration"][i] for i in valid_indices
|
||||
]
|
||||
|
||||
batch_captions = valid_data["text"]
|
||||
# Encode text using the standalone TextEncodingStage API
|
||||
prompt_embeds_list, prompt_masks_list = (
|
||||
self.prompt_encoding_stage.encode_text(
|
||||
batch_captions,
|
||||
server_args,
|
||||
encoder_index=[0],
|
||||
return_attention_mask=True,
|
||||
)
|
||||
)
|
||||
prompt_embeds = prompt_embeds_list[0]
|
||||
prompt_attention_masks = prompt_masks_list[0]
|
||||
assert prompt_embeds.shape[0] == prompt_attention_masks.shape[0]
|
||||
|
||||
sampling_params = SamplingParams.from_pretrained(args.model_path)
|
||||
|
||||
# encode negative prompt for trajectory collection
|
||||
if (
|
||||
sampling_params.guidance_scale > 1
|
||||
and sampling_params.negative_prompt is not None
|
||||
):
|
||||
negative_prompt_embeds_list, negative_prompt_masks_list = (
|
||||
self.prompt_encoding_stage.encode_text(
|
||||
sampling_params.negative_prompt,
|
||||
server_args,
|
||||
encoder_index=[0],
|
||||
return_attention_mask=True,
|
||||
)
|
||||
)
|
||||
negative_prompt_embed = negative_prompt_embeds_list[0][0]
|
||||
negative_prompt_attention_mask = negative_prompt_masks_list[0][0]
|
||||
else:
|
||||
negative_prompt_embed = None
|
||||
negative_prompt_attention_mask = None
|
||||
|
||||
trajectory_latents = []
|
||||
trajectory_timesteps = []
|
||||
trajectory_decoded = []
|
||||
|
||||
for i, (prompt_embed, prompt_attention_mask) in enumerate(
|
||||
zip(prompt_embeds, prompt_attention_masks, strict=False)
|
||||
):
|
||||
prompt_embed = prompt_embed.unsqueeze(0)
|
||||
prompt_attention_mask = prompt_attention_mask.unsqueeze(0)
|
||||
|
||||
# Collect the trajectory data (text-to-video generation)
|
||||
batch = Req(
|
||||
**shallow_asdict(sampling_params),
|
||||
)
|
||||
batch.prompt_embeds = [prompt_embed]
|
||||
batch.prompt_attention_mask = [prompt_attention_mask]
|
||||
batch.negative_prompt_embeds = [negative_prompt_embed]
|
||||
batch.negative_attention_mask = [negative_prompt_attention_mask]
|
||||
batch.num_inference_steps = 48
|
||||
batch.return_trajectory_latents = True
|
||||
# Enabling this will save the decoded trajectory videos.
|
||||
# Used for debugging.
|
||||
batch.return_trajectory_decoded = False
|
||||
batch.height = args.max_height
|
||||
batch.width = args.max_width
|
||||
batch.fps = args.train_fps
|
||||
batch.guidance_scale = 6.0
|
||||
batch.do_classifier_free_guidance = True
|
||||
|
||||
result_batch = self.input_validation_stage(batch, server_args)
|
||||
result_batch = self.timestep_preparation_stage(batch, server_args)
|
||||
result_batch = self.latent_preparation_stage(
|
||||
result_batch, server_args
|
||||
)
|
||||
result_batch = self.denoising_stage(result_batch, server_args)
|
||||
result_batch = self.decoding_stage(result_batch, server_args)
|
||||
|
||||
trajectory_latents.append(result_batch.trajectory_latents.cpu())
|
||||
trajectory_timesteps.append(result_batch.trajectory_timesteps.cpu())
|
||||
trajectory_decoded.append(result_batch.trajectory_decoded)
|
||||
|
||||
# Prepare extra features for text-only processing
|
||||
extra_features = {
|
||||
"trajectory_latents": trajectory_latents,
|
||||
"trajectory_timesteps": trajectory_timesteps,
|
||||
}
|
||||
|
||||
if batch.return_trajectory_decoded:
|
||||
for i, decoded_frames in enumerate(trajectory_decoded):
|
||||
for j, decoded_frame in enumerate(decoded_frames):
|
||||
save_decoded_latents_as_video(
|
||||
decoded_frame,
|
||||
f"decoded_videos/trajectory_decoded_{i}_{j}.mp4",
|
||||
args.train_fps,
|
||||
)
|
||||
|
||||
# Prepare batch data for Parquet dataset
|
||||
batch_data: list[dict[str, Any]] = []
|
||||
|
||||
# Add progress bar for saving outputs
|
||||
save_pbar = tqdm(
|
||||
enumerate(valid_data["path"]),
|
||||
desc="Saving outputs",
|
||||
unit="item",
|
||||
leave=False,
|
||||
)
|
||||
|
||||
for idx, video_path in save_pbar:
|
||||
video_name = os.path.basename(video_path).split(".")[0]
|
||||
|
||||
# Convert tensors to numpy arrays
|
||||
text_embedding = prompt_embeds[idx].cpu().numpy()
|
||||
|
||||
# Get extra features for this sample
|
||||
sample_extra_features = {}
|
||||
if extra_features:
|
||||
for key, value in extra_features.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
sample_extra_features[key] = value[idx].cpu().numpy()
|
||||
else:
|
||||
assert isinstance(value, list)
|
||||
if isinstance(value[idx], torch.Tensor):
|
||||
sample_extra_features[key] = (
|
||||
value[idx].cpu().float().numpy()
|
||||
)
|
||||
else:
|
||||
sample_extra_features[key] = value[idx]
|
||||
|
||||
# Create record for Parquet dataset (text-only ODE schema)
|
||||
record: dict[str, Any] = ode_text_only_record_creator(
|
||||
video_name=video_name,
|
||||
text_embedding=text_embedding,
|
||||
caption=valid_data["text"][idx],
|
||||
trajectory_latents=sample_extra_features["trajectory_latents"],
|
||||
trajectory_timesteps=sample_extra_features[
|
||||
"trajectory_timesteps"
|
||||
],
|
||||
)
|
||||
batch_data.append(record)
|
||||
|
||||
if batch_data:
|
||||
write_pbar = tqdm(
|
||||
total=1, desc="Writing to Parquet dataset", unit="batch"
|
||||
)
|
||||
table = records_to_table(batch_data, self.get_pyarrow_schema())
|
||||
write_pbar.update(1)
|
||||
write_pbar.close()
|
||||
|
||||
if not hasattr(self, "dataset_writer"):
|
||||
self.dataset_writer = ParquetDatasetWriter(
|
||||
out_dir=self.combined_parquet_dir,
|
||||
samples_per_file=args.samples_per_file,
|
||||
)
|
||||
self.dataset_writer.append_table(table)
|
||||
|
||||
logger.info("Collected batch with %s samples", len(table))
|
||||
|
||||
if self.num_processed_samples >= args.flush_frequency:
|
||||
written = self.dataset_writer.flush()
|
||||
logger.info("Flushed %s samples to parquet", written)
|
||||
self.num_processed_samples = 0
|
||||
|
||||
# Final flush for any remaining samples
|
||||
if hasattr(self, "dataset_writer"):
|
||||
written = self.dataset_writer.flush(write_remainder=True)
|
||||
if written:
|
||||
logger.info("Final flush wrote %s samples", written)
|
||||
|
||||
def forward(self, batch: Req, server_args: ServerArgs, args):
|
||||
if not self.post_init_called:
|
||||
self.post_init()
|
||||
|
||||
self.local_rank = int(os.getenv("RANK", 0))
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
# Create directory for combined data
|
||||
self.combined_parquet_dir = os.path.join(
|
||||
args.output_dir, "combined_parquet_dataset"
|
||||
)
|
||||
os.makedirs(self.combined_parquet_dir, exist_ok=True)
|
||||
|
||||
# Loading dataset
|
||||
train_dataset = gettextdataset(args)
|
||||
|
||||
self.preprocess_dataloader = DataLoader(
|
||||
train_dataset,
|
||||
batch_size=args.preprocess_video_batch_size,
|
||||
num_workers=args.dataloader_num_workers,
|
||||
)
|
||||
|
||||
self.preprocess_loader_iter = iter(self.preprocess_dataloader)
|
||||
|
||||
self.num_processed_samples = 0
|
||||
# Add progress bar for video preprocessing
|
||||
self.pbar = tqdm(
|
||||
self.preprocess_loader_iter,
|
||||
desc="Processing videos",
|
||||
unit="batch",
|
||||
disable=self.local_rank != 0,
|
||||
)
|
||||
|
||||
# Initialize class variables for data sharing
|
||||
self.video_data: dict[str, Any] = {} # Store video metadata and paths
|
||||
self.latent_data: dict[str, Any] = {} # Store latent tensors
|
||||
self.preprocess_text_and_trajectory(server_args, args)
|
||||
|
||||
|
||||
EntryClass = PreprocessPipeline_ODE_Trajectory
|
||||
@@ -0,0 +1,26 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
T2V Data Preprocessing pipeline implementation.
|
||||
|
||||
This module contains an implementation of the T2V Data Preprocessing pipeline
|
||||
using the modular pipeline architecture.
|
||||
"""
|
||||
from sglang.multimodal_gen.dataset.dataloader.schema import pyarrow_schema_t2v
|
||||
from sglang.multimodal_gen.runtime.pipelines.preprocess.preprocess_pipeline_base import (
|
||||
BasePreprocessPipeline,
|
||||
)
|
||||
|
||||
|
||||
class PreprocessPipeline_T2V(BasePreprocessPipeline):
|
||||
"""T2V preprocessing pipeline implementation."""
|
||||
|
||||
_required_config_modules = ["text_encoder", "tokenizer", "vae"]
|
||||
|
||||
def get_pyarrow_schema(self):
|
||||
"""Return the PyArrow schema for T2V pipeline."""
|
||||
return pyarrow_schema_t2v
|
||||
|
||||
|
||||
EntryClass = PreprocessPipeline_T2V
|
||||
@@ -0,0 +1,200 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Text-only Data Preprocessing pipeline implementation.
|
||||
|
||||
This module contains an implementation of the Text-only Data Preprocessing pipeline
|
||||
using the modular pipeline architecture, based on the ODE Trajectory preprocessing.
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from torchdata.stateful_dataloader import StatefulDataLoader
|
||||
from tqdm import tqdm
|
||||
|
||||
from sglang.multimodal_gen.dataset import gettextdataset
|
||||
from sglang.multimodal_gen.dataset.dataloader.parquet_io import (
|
||||
ParquetDatasetWriter,
|
||||
records_to_table,
|
||||
)
|
||||
from sglang.multimodal_gen.dataset.dataloader.record_schema import (
|
||||
text_only_record_creator,
|
||||
)
|
||||
from sglang.multimodal_gen.dataset.dataloader.schema import pyarrow_schema_text_only
|
||||
from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req
|
||||
from sglang.multimodal_gen.runtime.pipelines.preprocess.preprocess_pipeline_base import (
|
||||
BasePreprocessPipeline,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages import TextEncodingStage
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class PreprocessPipeline_Text(BasePreprocessPipeline):
|
||||
"""Text-only preprocessing pipeline implementation."""
|
||||
|
||||
_required_config_modules = ["text_encoder", "tokenizer"]
|
||||
preprocess_dataloader: StatefulDataLoader
|
||||
preprocess_loader_iter: Iterator[dict[str, Any]]
|
||||
pbar: Any
|
||||
num_processed_samples: int = 0
|
||||
|
||||
def get_pyarrow_schema(self):
|
||||
"""Return the PyArrow schema for text-only pipeline."""
|
||||
return pyarrow_schema_text_only
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||
"""Set up pipeline stages with proper dependency injection."""
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[self.get_module("text_encoder")],
|
||||
tokenizers=[self.get_module("tokenizer")],
|
||||
),
|
||||
)
|
||||
|
||||
def preprocess_text_only(self, server_args: ServerArgs, args):
|
||||
"""Preprocess text-only data."""
|
||||
|
||||
for batch_idx, data in enumerate(self.pbar):
|
||||
if data is None:
|
||||
continue
|
||||
|
||||
with torch.inference_mode():
|
||||
# For text-only processing, we only need text data
|
||||
# Filter out samples without text
|
||||
valid_indices = []
|
||||
for i, text in enumerate(data["text"]):
|
||||
if text and text.strip(): # Check if text is not empty
|
||||
valid_indices.append(i)
|
||||
self.num_processed_samples += len(valid_indices)
|
||||
|
||||
if not valid_indices:
|
||||
continue
|
||||
|
||||
# Create new batch with only valid samples (text-only)
|
||||
valid_data = {
|
||||
"text": [data["text"][i] for i in valid_indices],
|
||||
"path": [data["path"][i] for i in valid_indices],
|
||||
}
|
||||
|
||||
batch_captions = valid_data["text"]
|
||||
# Encode text using the standalone TextEncodingStage API
|
||||
prompt_embeds_list, prompt_masks_list = (
|
||||
self.prompt_encoding_stage.encode_text(
|
||||
batch_captions,
|
||||
server_args,
|
||||
encoder_index=[0],
|
||||
return_attention_mask=True,
|
||||
)
|
||||
)
|
||||
prompt_embeds = prompt_embeds_list[0]
|
||||
prompt_attention_masks = prompt_masks_list[0]
|
||||
assert prompt_embeds.shape[0] == prompt_attention_masks.shape[0]
|
||||
|
||||
logger.info("===== prompt_embeds: %s", prompt_embeds.shape)
|
||||
logger.info(
|
||||
"===== prompt_attention_masks: %s", prompt_attention_masks.shape
|
||||
)
|
||||
|
||||
# Prepare batch data for Parquet dataset
|
||||
batch_data = []
|
||||
|
||||
# Add progress bar for saving outputs
|
||||
save_pbar = tqdm(
|
||||
enumerate(valid_data["path"]),
|
||||
desc="Saving outputs",
|
||||
unit="item",
|
||||
leave=False,
|
||||
)
|
||||
|
||||
for idx, text_path in save_pbar:
|
||||
text_name = os.path.basename(text_path).split(".")[0]
|
||||
|
||||
# Convert tensors to numpy arrays
|
||||
text_embedding = prompt_embeds[idx].cpu().numpy()
|
||||
|
||||
# Create record for Parquet dataset (text-only schema)
|
||||
record = text_only_record_creator(
|
||||
text_name=text_name,
|
||||
text_embedding=text_embedding,
|
||||
caption=valid_data["text"][idx],
|
||||
)
|
||||
batch_data.append(record)
|
||||
|
||||
if batch_data:
|
||||
write_pbar = tqdm(
|
||||
total=1, desc="Writing to Parquet dataset", unit="batch"
|
||||
)
|
||||
table = records_to_table(batch_data, pyarrow_schema_text_only)
|
||||
write_pbar.update(1)
|
||||
write_pbar.close()
|
||||
|
||||
if not hasattr(self, "dataset_writer"):
|
||||
self.dataset_writer = ParquetDatasetWriter(
|
||||
out_dir=self.combined_parquet_dir,
|
||||
samples_per_file=args.samples_per_file,
|
||||
)
|
||||
self.dataset_writer.append_table(table)
|
||||
|
||||
logger.info("Collected batch with %s samples", len(table))
|
||||
|
||||
if self.num_processed_samples >= args.flush_frequency:
|
||||
written = self.dataset_writer.flush()
|
||||
logger.info("Flushed %s samples to parquet", written)
|
||||
self.num_processed_samples = 0
|
||||
|
||||
# Final flush for any remaining samples
|
||||
if hasattr(self, "dataset_writer"):
|
||||
written = self.dataset_writer.flush(write_remainder=True)
|
||||
if written:
|
||||
logger.info("Final flush wrote %s samples", written)
|
||||
|
||||
# Text-only record creation moved to sglang.multimodal_gen.dataset.dataloader.record_schema
|
||||
|
||||
def forward(self, batch: Req, server_args: ServerArgs, args):
|
||||
if not self.post_init_called:
|
||||
self.post_init()
|
||||
|
||||
self.local_rank = int(os.getenv("RANK", 0))
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
# Create directory for combined data
|
||||
self.combined_parquet_dir = os.path.join(
|
||||
args.output_dir, "combined_parquet_dataset"
|
||||
)
|
||||
os.makedirs(self.combined_parquet_dir, exist_ok=True)
|
||||
|
||||
# Loading text dataset
|
||||
train_dataset = gettextdataset(args)
|
||||
|
||||
self.preprocess_dataloader = DataLoader(
|
||||
train_dataset,
|
||||
batch_size=args.preprocess_video_batch_size,
|
||||
num_workers=args.dataloader_num_workers,
|
||||
)
|
||||
|
||||
self.preprocess_loader_iter = iter(self.preprocess_dataloader)
|
||||
|
||||
self.num_processed_samples = 0
|
||||
# Add progress bar for text preprocessing
|
||||
self.pbar = tqdm(
|
||||
self.preprocess_loader_iter,
|
||||
desc="Processing text",
|
||||
unit="batch",
|
||||
disable=self.local_rank != 0,
|
||||
)
|
||||
|
||||
# Initialize class variables for data sharing
|
||||
self.text_data: dict[str, Any] = {} # Store text metadata and paths
|
||||
|
||||
self.preprocess_text_only(server_args, args)
|
||||
|
||||
|
||||
EntryClass = PreprocessPipeline_Text
|
||||
@@ -0,0 +1,134 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
import random
|
||||
from collections.abc import Callable
|
||||
from typing import cast
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision
|
||||
from einops import rearrange
|
||||
from torchvision import transforms
|
||||
|
||||
from sglang.multimodal_gen.configs.configs import VideoLoaderType
|
||||
from sglang.multimodal_gen.dataset.transform import (
|
||||
CenterCropResizeVideo,
|
||||
TemporalRandomCrop,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import (
|
||||
PreprocessBatch,
|
||||
Req,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages.base import PipelineStage
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs, WorkloadType
|
||||
|
||||
|
||||
class VideoTransformStage(PipelineStage):
|
||||
"""
|
||||
Crop a video in temporal dimension.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_fps: int,
|
||||
num_frames: int,
|
||||
max_height: int,
|
||||
max_width: int,
|
||||
do_temporal_sample: bool,
|
||||
) -> None:
|
||||
self.train_fps = train_fps
|
||||
self.num_frames = num_frames
|
||||
if do_temporal_sample:
|
||||
self.temporal_sample_fn: Callable | None = TemporalRandomCrop(num_frames)
|
||||
else:
|
||||
self.temporal_sample_fn = None
|
||||
|
||||
self.video_transform = transforms.Compose(
|
||||
[
|
||||
CenterCropResizeVideo((max_height, max_width)),
|
||||
]
|
||||
)
|
||||
|
||||
def forward(self, batch: Req, server_args: ServerArgs) -> Req:
|
||||
batch = cast(PreprocessBatch, batch)
|
||||
assert isinstance(batch.fps, list)
|
||||
assert isinstance(batch.num_frames, list)
|
||||
|
||||
if batch.data_type != "video":
|
||||
return batch
|
||||
|
||||
if len(batch.video_loader) == 0:
|
||||
raise ValueError("Video loader is not set")
|
||||
|
||||
video_pixel_batch = []
|
||||
|
||||
for i in range(len(batch.video_loader)):
|
||||
frame_interval = batch.fps[i] / self.train_fps
|
||||
start_frame_idx = 0
|
||||
frame_indices = np.arange(
|
||||
start_frame_idx, batch.num_frames[i], frame_interval
|
||||
).astype(int)
|
||||
if len(frame_indices) > self.num_frames:
|
||||
if self.temporal_sample_fn is not None:
|
||||
begin_index, end_index = self.temporal_sample_fn(len(frame_indices))
|
||||
frame_indices = frame_indices[begin_index:end_index]
|
||||
else:
|
||||
frame_indices = frame_indices[: self.num_frames]
|
||||
|
||||
if (
|
||||
server_args.preprocess_config.video_loader_type
|
||||
== VideoLoaderType.TORCHCODEC
|
||||
):
|
||||
video = batch.video_loader[i].get_frames_at(frame_indices).data
|
||||
elif (
|
||||
server_args.preprocess_config.video_loader_type
|
||||
== VideoLoaderType.TORCHVISION
|
||||
):
|
||||
video, _, _ = torchvision.io.read_video(
|
||||
batch.video_loader[i], output_format="TCHW"
|
||||
)
|
||||
video = video[frame_indices]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid video loader type: {server_args.preprocess_config.video_loader_type}"
|
||||
)
|
||||
video = self.video_transform(video)
|
||||
video_pixel_batch.append(video)
|
||||
|
||||
video_pixel_values = torch.stack(video_pixel_batch)
|
||||
video_pixel_values = rearrange(video_pixel_values, "b t c h w -> b c t h w")
|
||||
video_pixel_values = video_pixel_values.to(torch.uint8)
|
||||
|
||||
if server_args.workload_type == WorkloadType.I2V:
|
||||
batch.pil_image = video_pixel_values[:, :, 0, :, :]
|
||||
|
||||
video_pixel_values = video_pixel_values.float() / 255.0
|
||||
batch.latents = video_pixel_values
|
||||
batch.num_frames = [video_pixel_values.shape[2]] * len(batch.video_loader)
|
||||
batch.height = [video_pixel_values.shape[3]] * len(batch.video_loader)
|
||||
batch.width = [video_pixel_values.shape[4]] * len(batch.video_loader)
|
||||
return cast(Req, batch)
|
||||
|
||||
|
||||
class TextTransformStage(PipelineStage):
|
||||
"""
|
||||
Process text data according to the cfg rate.
|
||||
"""
|
||||
|
||||
def __init__(self, cfg_uncondition_drop_rate: float, seed: int) -> None:
|
||||
self.cfg_rate = cfg_uncondition_drop_rate
|
||||
self.rng = random.Random(seed)
|
||||
|
||||
def forward(self, batch: Req, server_args: ServerArgs) -> Req:
|
||||
batch = cast(PreprocessBatch, batch)
|
||||
|
||||
prompts = []
|
||||
for prompt in batch.prompt:
|
||||
if not isinstance(prompt, list):
|
||||
prompt = [prompt]
|
||||
prompt = self.rng.choice(prompt)
|
||||
prompt = prompt if self.rng.random() > self.cfg_rate else ""
|
||||
prompts.append(prompt)
|
||||
|
||||
batch.prompt = prompts
|
||||
return cast(Req, batch)
|
||||
@@ -0,0 +1,147 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from sglang.multimodal_gen import PipelineConfig
|
||||
from sglang.multimodal_gen.configs.models.vaes import WanVAEConfig
|
||||
from sglang.multimodal_gen.runtime.architectures.preprocess.preprocess_pipeline_i2v import (
|
||||
PreprocessPipeline_I2V,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.architectures.preprocess.preprocess_pipeline_ode_trajectory import (
|
||||
PreprocessPipeline_ODE_Trajectory,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.architectures.preprocess.preprocess_pipeline_t2v import (
|
||||
PreprocessPipeline_T2V,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.architectures.preprocess.preprocess_pipeline_text import (
|
||||
PreprocessPipeline_Text,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.distributed import (
|
||||
get_world_size,
|
||||
maybe_init_distributed_environment_and_model_parallel,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def main(args) -> None:
|
||||
args.model_path = maybe_download_model(args.model_path)
|
||||
maybe_init_distributed_environment_and_model_parallel(1, 1)
|
||||
num_gpus = int(os.environ["WORLD_SIZE"])
|
||||
assert num_gpus == 1, "Only support 1 GPU"
|
||||
|
||||
pipeline_config = PipelineConfig.from_pretrained(args.model_path)
|
||||
|
||||
kwargs: dict[str, Any] = {}
|
||||
if args.preprocess_task == "text_only":
|
||||
kwargs = {
|
||||
"text_encoder_cpu_offload": False,
|
||||
}
|
||||
else:
|
||||
# Full config for video/image processing
|
||||
kwargs = {
|
||||
"vae_precision": "fp32",
|
||||
"vae_config": WanVAEConfig(load_encoder=True, load_decoder=True),
|
||||
}
|
||||
pipeline_config.update_config_from_dict(kwargs)
|
||||
|
||||
server_args = ServerArgs(
|
||||
model_path=args.model_path,
|
||||
num_gpus=get_world_size(),
|
||||
dit_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
text_encoder_cpu_offload=False,
|
||||
pipeline_config=pipeline_config,
|
||||
)
|
||||
if args.preprocess_task == "t2v":
|
||||
PreprocessPipeline = PreprocessPipeline_T2V
|
||||
elif args.preprocess_task == "i2v":
|
||||
PreprocessPipeline = PreprocessPipeline_I2V
|
||||
elif args.preprocess_task == "text_only":
|
||||
PreprocessPipeline = PreprocessPipeline_Text
|
||||
elif args.preprocess_task == "ode_trajectory":
|
||||
assert args.flow_shift is not None, "flow_shift is required for ode_trajectory"
|
||||
server_args.pipeline_config.flow_shift = args.flow_shift
|
||||
PreprocessPipeline = PreprocessPipeline_ODE_Trajectory
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid preprocess task: {args.preprocess_task}. "
|
||||
f"Valid options: t2v, i2v, ode_trajectory, text_only"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Preprocess task: %s using %s",
|
||||
args.preprocess_task,
|
||||
PreprocessPipeline.__name__,
|
||||
)
|
||||
|
||||
pipeline = PreprocessPipeline(args.model_path, server_args)
|
||||
pipeline.forward(batch=None, server_args=server_args, args=args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
# dataset & dataloader
|
||||
parser.add_argument("--model_path", type=str, default="data/mochi")
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--data_merge_path", type=str, required=True)
|
||||
parser.add_argument("--num_frames", type=int, default=163)
|
||||
parser.add_argument(
|
||||
"--dataloader_num_workers",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--preprocess_video_batch_size",
|
||||
type=int,
|
||||
default=2,
|
||||
help="Batch size (per device) for the training dataloader.",
|
||||
)
|
||||
parser.add_argument("--samples_per_file", type=int, default=64)
|
||||
parser.add_argument(
|
||||
"--flush_frequency",
|
||||
type=int,
|
||||
default=256,
|
||||
help="how often to save to parquet files",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num_latent_t", type=int, default=28, help="Number of latent timesteps."
|
||||
)
|
||||
parser.add_argument("--max_height", type=int, default=480)
|
||||
parser.add_argument("--max_width", type=int, default=848)
|
||||
parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
|
||||
parser.add_argument("--group_frame", action="store_true") # TODO
|
||||
parser.add_argument("--group_resolution", action="store_true") # TODO
|
||||
parser.add_argument("--flow_shift", type=float, default=None)
|
||||
parser.add_argument(
|
||||
"--preprocess_task",
|
||||
type=str,
|
||||
default="t2v",
|
||||
choices=["t2v", "i2v", "text_only", "ode_trajectory"],
|
||||
help="Type of preprocessing task to run",
|
||||
)
|
||||
parser.add_argument("--train_fps", type=int, default=30)
|
||||
parser.add_argument("--use_image_num", type=int, default=0)
|
||||
parser.add_argument("--text_max_length", type=int, default=256)
|
||||
parser.add_argument("--speed_factor", type=float, default=1.0)
|
||||
parser.add_argument("--drop_short_ratio", type=float, default=1.0)
|
||||
parser.add_argument("--do_temporal_sample", default=False, action="store_true")
|
||||
# text encoder & vae & diffusion model
|
||||
parser.add_argument("--text_encoder_name", type=str, default="google/t5-v1_1-xxl")
|
||||
parser.add_argument("--cache_dir", type=str, default="./cache_dir")
|
||||
parser.add_argument("--training_cfg_rate", type=float, default=0.0)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="The output directory where the model predictions and checkpoints will be written.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
main(args)
|
||||
@@ -0,0 +1,26 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed import (
|
||||
maybe_init_distributed_environment_and_model_parallel,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
from sglang.multimodal_gen.runtime.workflow.workflow_base import WorkflowBase
|
||||
from sglang.multimodal_gen.utils import FlexibleArgumentParser
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def main(server_args: ServerArgs) -> None:
|
||||
maybe_init_distributed_environment_and_model_parallel(1, 1)
|
||||
preprocess_workflow_cls = WorkflowBase.get_workflow_cls(server_args)
|
||||
preprocess_workflow = preprocess_workflow_cls(server_args)
|
||||
preprocess_workflow.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = FlexibleArgumentParser()
|
||||
parser = ServerArgs.add_cli_args(parser)
|
||||
args = parser.parse_args()
|
||||
server_args = ServerArgs.from_cli_args(args)
|
||||
main(server_args)
|
||||
@@ -0,0 +1 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
@@ -0,0 +1,118 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import (
|
||||
ComposedPipelineBase,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines.preprocess.preprocess_stages import (
|
||||
TextTransformStage,
|
||||
VideoTransformStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages import (
|
||||
EncodingStage,
|
||||
ImageEncodingStage,
|
||||
TextEncodingStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines.stages.image_encoding import (
|
||||
ImageVAEEncodingStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
|
||||
|
||||
class PreprocessPipelineI2V(ComposedPipelineBase):
|
||||
_required_config_modules = [
|
||||
"image_encoder",
|
||||
"image_processor",
|
||||
"text_encoder",
|
||||
"tokenizer",
|
||||
"vae",
|
||||
]
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||
assert server_args.preprocess_config is not None
|
||||
self.add_stage(
|
||||
stage_name="text_transform_stage",
|
||||
stage=TextTransformStage(
|
||||
cfg_uncondition_drop_rate=server_args.preprocess_config.training_cfg_rate,
|
||||
seed=server_args.preprocess_config.seed,
|
||||
),
|
||||
)
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[self.get_module("text_encoder")],
|
||||
tokenizers=[self.get_module("tokenizer")],
|
||||
),
|
||||
)
|
||||
self.add_stage(
|
||||
stage_name="video_transform_stage",
|
||||
stage=VideoTransformStage(
|
||||
train_fps=server_args.preprocess_config.train_fps,
|
||||
num_frames=server_args.preprocess_config.num_frames,
|
||||
max_height=server_args.preprocess_config.max_height,
|
||||
max_width=server_args.preprocess_config.max_width,
|
||||
do_temporal_sample=server_args.preprocess_config.do_temporal_sample,
|
||||
),
|
||||
)
|
||||
if (
|
||||
self.get_module("image_encoder") is not None
|
||||
and self.get_module("image_processor") is not None
|
||||
):
|
||||
self.add_stage(
|
||||
stage_name="image_encoding_stage",
|
||||
stage=ImageEncodingStage(
|
||||
image_encoder=self.get_module("image_encoder"),
|
||||
image_processor=self.get_module("image_processor"),
|
||||
),
|
||||
)
|
||||
self.add_stage(
|
||||
stage_name="image_vae_encoding_stage",
|
||||
stage=ImageVAEEncodingStage(
|
||||
vae=self.get_module("vae"),
|
||||
),
|
||||
)
|
||||
self.add_stage(
|
||||
stage_name="video_encoding_stage",
|
||||
stage=EncodingStage(
|
||||
vae=self.get_module("vae"),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class PreprocessPipelineT2V(ComposedPipelineBase):
|
||||
_required_config_modules = ["text_encoder", "tokenizer", "vae"]
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||
assert server_args.preprocess_config is not None
|
||||
self.add_stage(
|
||||
stage_name="text_transform_stage",
|
||||
stage=TextTransformStage(
|
||||
cfg_uncondition_drop_rate=server_args.preprocess_config.training_cfg_rate,
|
||||
seed=server_args.preprocess_config.seed,
|
||||
),
|
||||
)
|
||||
self.add_stage(
|
||||
stage_name="prompt_encoding_stage",
|
||||
stage=TextEncodingStage(
|
||||
text_encoders=[self.get_module("text_encoder")],
|
||||
tokenizers=[self.get_module("tokenizer")],
|
||||
),
|
||||
)
|
||||
self.add_stage(
|
||||
stage_name="video_transform_stage",
|
||||
stage=VideoTransformStage(
|
||||
train_fps=server_args.preprocess_config.train_fps,
|
||||
num_frames=server_args.preprocess_config.num_frames,
|
||||
max_height=server_args.preprocess_config.max_height,
|
||||
max_width=server_args.preprocess_config.max_width,
|
||||
do_temporal_sample=server_args.preprocess_config.do_temporal_sample,
|
||||
),
|
||||
)
|
||||
self.add_stage(
|
||||
stage_name="video_encoding_stage",
|
||||
stage=EncodingStage(
|
||||
vae=self.get_module("vae"),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
EntryClass = [PreprocessPipelineI2V, PreprocessPipelineT2V]
|
||||
55
python/sglang/multimodal_gen/runtime/distributed/__init__.py
Normal file
55
python/sglang/multimodal_gen/runtime/distributed/__init__.py
Normal file
@@ -0,0 +1,55 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed.communication_op import *
|
||||
from sglang.multimodal_gen.runtime.distributed.group_coordinator import (
|
||||
get_local_torch_device,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
|
||||
cleanup_dist_env_and_memory,
|
||||
get_dp_group,
|
||||
get_dp_rank,
|
||||
get_dp_world_size,
|
||||
get_sp_group,
|
||||
get_sp_parallel_rank,
|
||||
get_sp_world_size,
|
||||
get_tp_group,
|
||||
get_tp_rank,
|
||||
get_tp_world_size,
|
||||
get_world_group,
|
||||
get_world_rank,
|
||||
get_world_size,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
maybe_init_distributed_environment_and_model_parallel,
|
||||
model_parallel_is_initialized,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.distributed.utils import *
|
||||
|
||||
__all__ = [
|
||||
# Initialization
|
||||
"init_distributed_environment",
|
||||
"initialize_model_parallel",
|
||||
"cleanup_dist_env_and_memory",
|
||||
"model_parallel_is_initialized",
|
||||
"maybe_init_distributed_environment_and_model_parallel",
|
||||
# World group
|
||||
"get_world_group",
|
||||
"get_world_rank",
|
||||
"get_world_size",
|
||||
# Data parallel group
|
||||
"get_dp_group",
|
||||
"get_dp_rank",
|
||||
"get_dp_world_size",
|
||||
# Sequence parallel group
|
||||
"get_sp_group",
|
||||
"get_sp_parallel_rank",
|
||||
"get_sp_world_size",
|
||||
# Tensor parallel group
|
||||
"get_tp_group",
|
||||
"get_tp_rank",
|
||||
"get_tp_world_size",
|
||||
# Get torch device
|
||||
"get_local_torch_device",
|
||||
]
|
||||
@@ -0,0 +1,55 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/communication_op.py
|
||||
|
||||
import torch
|
||||
import torch.distributed
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
|
||||
get_cfg_group,
|
||||
get_sp_group,
|
||||
get_tp_group,
|
||||
)
|
||||
|
||||
|
||||
def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
|
||||
"""All-reduce the input tensor across model parallel group."""
|
||||
return get_tp_group().all_reduce(input_)
|
||||
|
||||
|
||||
def tensor_model_parallel_all_gather(
|
||||
input_: torch.Tensor, dim: int = -1
|
||||
) -> torch.Tensor:
|
||||
"""All-gather the input tensor across model parallel group."""
|
||||
return get_tp_group().all_gather(input_, dim)
|
||||
|
||||
|
||||
# TODO: remove model, make it sequence_parallel
|
||||
def sequence_model_parallel_all_to_all_4D(
|
||||
input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1
|
||||
) -> torch.Tensor:
|
||||
"""All-to-all communication of 4D tensors (e.g. QKV matrices) across sequence parallel group."""
|
||||
return get_sp_group().all_to_all_4D(input_, scatter_dim, gather_dim)
|
||||
|
||||
|
||||
def sequence_model_parallel_all_gather(
|
||||
input_: torch.Tensor, dim: int = -1
|
||||
) -> torch.Tensor:
|
||||
"""All-gather the input tensor across model parallel group."""
|
||||
return get_sp_group().all_gather(input_, dim)
|
||||
|
||||
|
||||
def cfg_model_parallel_all_gather(
|
||||
input_: torch.Tensor, dim: int = -1, separate_tensors: bool = False
|
||||
) -> torch.Tensor:
|
||||
"""All-gather the input tensor across model parallel group."""
|
||||
return get_cfg_group().all_gather(input_, dim, separate_tensors)
|
||||
|
||||
|
||||
def cfg_model_parallel_all_reduce(
|
||||
input_: torch.Tensor,
|
||||
op: torch._C._distributed_c10d.ReduceOp = torch._C._distributed_c10d.ReduceOp.SUM,
|
||||
) -> torch.Tensor:
|
||||
"""All-reduce the input tensor across CFG parallel group."""
|
||||
return get_cfg_group().all_reduce(input_, op=op)
|
||||
@@ -0,0 +1 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
@@ -0,0 +1,297 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/base_device_communicator.py
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch import Tensor
|
||||
from torch.distributed import ProcessGroup, ReduceOp
|
||||
|
||||
|
||||
class DistributedAutograd:
|
||||
"""Collection of autograd functions for distributed operations.
|
||||
|
||||
This class provides custom autograd functions for distributed operations like all_reduce,
|
||||
all_gather, and all_to_all. Each operation is implemented as a static inner class with
|
||||
proper forward and backward implementations.
|
||||
"""
|
||||
|
||||
class AllReduce(torch.autograd.Function):
|
||||
"""Differentiable all_reduce operation.
|
||||
|
||||
The gradient of all_reduce is another all_reduce operation since the operation
|
||||
combines values from all ranks equally.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx: Any,
|
||||
group: ProcessGroup,
|
||||
input_: Tensor,
|
||||
op: dist.ReduceOp | None = None,
|
||||
) -> Tensor:
|
||||
ctx.group = group
|
||||
ctx.op = op
|
||||
output = input_.clone()
|
||||
dist.all_reduce(output, group=group, op=op)
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx: Any, grad_output: Tensor) -> tuple[None, Tensor, None]:
|
||||
grad_output = grad_output.clone()
|
||||
dist.all_reduce(grad_output, group=ctx.group, op=ctx.op)
|
||||
return None, grad_output, None
|
||||
|
||||
class AllGather(torch.autograd.Function):
|
||||
"""Differentiable all_gather operation.
|
||||
|
||||
The operation gathers tensors from all ranks and concatenates them along a specified dimension.
|
||||
The backward pass uses reduce_scatter to efficiently distribute gradients back to source ranks.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx: Any, group: ProcessGroup, input_: Tensor, world_size: int, dim: int
|
||||
) -> Tensor:
|
||||
ctx.group = group
|
||||
ctx.world_size = world_size
|
||||
ctx.dim = dim
|
||||
ctx.input_shape = input_.shape
|
||||
|
||||
input_size = input_.size()
|
||||
output_size = (input_size[0] * world_size,) + input_size[1:]
|
||||
output_tensor = torch.empty(
|
||||
output_size, dtype=input_.dtype, device=input_.device
|
||||
)
|
||||
|
||||
dist.all_gather_into_tensor(output_tensor, input_, group=group)
|
||||
|
||||
output_tensor = output_tensor.reshape((world_size,) + input_size)
|
||||
output_tensor = output_tensor.movedim(0, dim)
|
||||
output_tensor = output_tensor.reshape(
|
||||
input_size[:dim]
|
||||
+ (world_size * input_size[dim],)
|
||||
+ input_size[dim + 1 :]
|
||||
)
|
||||
return output_tensor
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx: Any, grad_output: Tensor) -> tuple[None, Tensor, None, None]:
|
||||
# Split the gradient tensor along the gathered dimension
|
||||
dim_size = grad_output.size(ctx.dim) // ctx.world_size
|
||||
grad_chunks = grad_output.reshape(
|
||||
grad_output.shape[: ctx.dim]
|
||||
+ (ctx.world_size, dim_size)
|
||||
+ grad_output.shape[ctx.dim + 1 :]
|
||||
)
|
||||
grad_chunks = grad_chunks.movedim(ctx.dim, 0)
|
||||
|
||||
# Each rank only needs its corresponding gradient
|
||||
grad_input = torch.empty(
|
||||
ctx.input_shape, dtype=grad_output.dtype, device=grad_output.device
|
||||
)
|
||||
dist.reduce_scatter_tensor(
|
||||
grad_input, grad_chunks.contiguous(), group=ctx.group
|
||||
)
|
||||
|
||||
return None, grad_input, None, None
|
||||
|
||||
class AllToAll4D(torch.autograd.Function):
|
||||
"""Differentiable all_to_all operation specialized for 4D tensors.
|
||||
|
||||
This operation is particularly useful for attention operations where we need to
|
||||
redistribute data across ranks for efficient parallel processing.
|
||||
|
||||
The operation supports two modes:
|
||||
1. scatter_dim=2, gather_dim=1: Used for redistributing attention heads
|
||||
2. scatter_dim=1, gather_dim=2: Used for redistributing sequence dimensions
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx: Any,
|
||||
group: ProcessGroup,
|
||||
input_: Tensor,
|
||||
world_size: int,
|
||||
scatter_dim: int,
|
||||
gather_dim: int,
|
||||
) -> Tensor:
|
||||
ctx.group = group
|
||||
ctx.world_size = world_size
|
||||
ctx.scatter_dim = scatter_dim
|
||||
ctx.gather_dim = gather_dim
|
||||
|
||||
if world_size == 1:
|
||||
return input_
|
||||
|
||||
assert (
|
||||
input_.dim() == 4
|
||||
), f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}"
|
||||
|
||||
if scatter_dim == 2 and gather_dim == 1:
|
||||
bs, shard_seqlen, hn, hd = input_.shape
|
||||
seqlen = shard_seqlen * world_size
|
||||
shard_hn = hn // world_size
|
||||
|
||||
input_ = input_.transpose(0, 2).contiguous() # hn, shard_seqlen, bs, hd
|
||||
output = torch.empty_like(input_)
|
||||
|
||||
dist.all_to_all_single(
|
||||
output, input_, group=group
|
||||
) # hn, shard_seqlen, bs, hd
|
||||
|
||||
output = torch.cat(
|
||||
output.split(shard_hn), dim=1
|
||||
) # sharded hn, seqlen, bs, hd
|
||||
|
||||
output = output.transpose(
|
||||
0, 2
|
||||
).contiguous() # bs, seqlen, sharded_hn, hd
|
||||
|
||||
return output
|
||||
elif scatter_dim == 1 and gather_dim == 2:
|
||||
bs, seqlen, shard_hn, hd = input_.shape
|
||||
hn = shard_hn * world_size
|
||||
shard_seqlen = seqlen // world_size
|
||||
|
||||
input_ = input_.transpose(0, 2).contiguous() # shard_hn, seqlen, bs, hd
|
||||
|
||||
input_ = (
|
||||
input_.reshape(shard_hn, world_size, shard_seqlen, bs, hd)
|
||||
.transpose(0, 1)
|
||||
.reshape(shard_hn * world_size, shard_seqlen, bs, hd)
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
output = torch.empty_like(input_)
|
||||
|
||||
dist.all_to_all_single(output, input_, group=group)
|
||||
|
||||
output = output.transpose(
|
||||
0, 2
|
||||
).contiguous() # bs, seqlen, sharded_hn, hd
|
||||
|
||||
return output
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Invalid scatter_dim={scatter_dim}, gather_dim={gather_dim}. "
|
||||
f"Only (scatter_dim=2, gather_dim=1) and (scatter_dim=1, gather_dim=2) are supported."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def backward(
|
||||
ctx: Any, grad_output: Tensor
|
||||
) -> tuple[None, Tensor, None, None, None]:
|
||||
if ctx.world_size == 1:
|
||||
return None, grad_output, None, None, None
|
||||
|
||||
# For backward pass, we swap scatter_dim and gather_dim
|
||||
output = DistributedAutograd.AllToAll4D.apply(
|
||||
ctx.group, grad_output, ctx.world_size, ctx.gather_dim, ctx.scatter_dim
|
||||
)
|
||||
return None, output, None, None, None
|
||||
|
||||
|
||||
class DeviceCommunicatorBase:
|
||||
"""
|
||||
Base class for device-specific communicator with autograd support.
|
||||
It can use the `cpu_group` to initialize the communicator.
|
||||
If the device has PyTorch integration (PyTorch can recognize its
|
||||
communication backend), the `device_group` will also be given.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cpu_group: ProcessGroup,
|
||||
device: torch.device | None = None,
|
||||
device_group: ProcessGroup | None = None,
|
||||
unique_name: str = "",
|
||||
):
|
||||
self.device = device or torch.device("cpu")
|
||||
self.cpu_group = cpu_group
|
||||
self.device_group = device_group
|
||||
self.unique_name = unique_name
|
||||
self.rank = dist.get_rank(cpu_group)
|
||||
self.world_size = dist.get_world_size(cpu_group)
|
||||
self.ranks = dist.get_process_group_ranks(cpu_group)
|
||||
self.global_rank = dist.get_rank()
|
||||
self.global_world_size = dist.get_world_size()
|
||||
self.rank_in_group = dist.get_group_rank(self.cpu_group, self.global_rank)
|
||||
|
||||
def all_reduce(
|
||||
self, input_: torch.Tensor, op: dist.ReduceOp | None = ReduceOp.SUM
|
||||
) -> torch.Tensor:
|
||||
"""Performs an all_reduce operation with gradient support."""
|
||||
return DistributedAutograd.AllReduce.apply(self.device_group, input_, op)
|
||||
|
||||
def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
|
||||
"""Performs an all_gather operation with gradient support."""
|
||||
if dim < 0:
|
||||
dim += input_.dim()
|
||||
return DistributedAutograd.AllGather.apply(
|
||||
self.device_group, input_, self.world_size, dim
|
||||
)
|
||||
|
||||
def all_to_all_4D(
|
||||
self, input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1
|
||||
) -> torch.Tensor:
|
||||
"""Performs a 4D all-to-all operation with gradient support."""
|
||||
return DistributedAutograd.AllToAll4D.apply(
|
||||
self.device_group, input_, self.world_size, scatter_dim, gather_dim
|
||||
)
|
||||
|
||||
def gather(
|
||||
self, input_: torch.Tensor, dst: int = 0, dim: int = -1
|
||||
) -> torch.Tensor | None:
|
||||
"""
|
||||
NOTE: We assume that the input tensor is on the same device across
|
||||
all the ranks.
|
||||
NOTE: `dst` is the local rank of the destination rank.
|
||||
"""
|
||||
world_size = self.world_size
|
||||
assert (
|
||||
-input_.dim() <= dim < input_.dim()
|
||||
), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
|
||||
if dim < 0:
|
||||
# Convert negative dim to positive.
|
||||
dim += input_.dim()
|
||||
|
||||
# Allocate output tensor.
|
||||
if self.rank_in_group == dst:
|
||||
gather_list = [torch.empty_like(input_) for _ in range(world_size)]
|
||||
else:
|
||||
gather_list = None
|
||||
# Gather.
|
||||
torch.distributed.gather(
|
||||
input_, gather_list, dst=self.ranks[dst], group=self.device_group
|
||||
)
|
||||
if self.rank_in_group == dst:
|
||||
output_tensor = torch.cat(gather_list, dim=dim)
|
||||
else:
|
||||
output_tensor = None
|
||||
return output_tensor
|
||||
|
||||
def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
|
||||
"""Sends a tensor to the destination rank in a non-blocking way"""
|
||||
"""NOTE: `dst` is the local rank of the destination rank."""
|
||||
if dst is None:
|
||||
dst = (self.rank_in_group + 1) % self.world_size
|
||||
torch.distributed.send(tensor, self.ranks[dst], self.device_group)
|
||||
|
||||
def recv(
|
||||
self, size: torch.Size, dtype: torch.dtype, src: int | None = None
|
||||
) -> torch.Tensor:
|
||||
"""Receives a tensor from the source rank."""
|
||||
"""NOTE: `src` is the local rank of the source rank."""
|
||||
if src is None:
|
||||
src = (self.rank_in_group - 1) % self.world_size
|
||||
|
||||
tensor = torch.empty(size, dtype=dtype, device=self.device)
|
||||
torch.distributed.recv(tensor, self.ranks[src], self.device_group)
|
||||
return tensor
|
||||
|
||||
def destroy(self) -> None:
|
||||
pass
|
||||
@@ -0,0 +1,161 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Adapted from: https://github.com/vllm-project/vllm/blob/main/vllm/distributed/device_communicators/cpu_communicator.py
|
||||
|
||||
import os
|
||||
|
||||
import torch
|
||||
from torch.distributed import ProcessGroup
|
||||
|
||||
from .base_device_communicator import DeviceCommunicatorBase
|
||||
|
||||
|
||||
class CpuCommunicator(DeviceCommunicatorBase):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cpu_group: ProcessGroup,
|
||||
device: torch.device | None = None,
|
||||
device_group: ProcessGroup | None = None,
|
||||
unique_name: str = "",
|
||||
):
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
from sglang.multimodal_gen.runtime.platforms.interface import CpuArchEnum
|
||||
|
||||
super().__init__(cpu_group, device, device_group, unique_name)
|
||||
self.dist_module = torch.distributed
|
||||
|
||||
if (
|
||||
(current_platform.get_cpu_architecture() == CpuArchEnum.X86)
|
||||
and hasattr(torch.ops._C, "init_shm_manager")
|
||||
and unique_name.startswith("tp")
|
||||
):
|
||||
self.dist_module = _CPUSHMDistributed(self)
|
||||
|
||||
def all_reduce(
|
||||
self,
|
||||
input_: torch.Tensor,
|
||||
op: torch.distributed.ReduceOp | None = torch.distributed.ReduceOp.SUM,
|
||||
) -> torch.Tensor:
|
||||
self.dist_module.all_reduce(input_, group=self.device_group, op=op)
|
||||
return input_
|
||||
|
||||
def gather(
|
||||
self, input_: torch.Tensor, dst: int = 0, dim: int = -1
|
||||
) -> torch.Tensor | None:
|
||||
"""
|
||||
NOTE: We assume that the input tensor is on the same device across
|
||||
all the ranks.
|
||||
NOTE: `dst` is the local rank of the destination rank.
|
||||
"""
|
||||
world_size = self.world_size
|
||||
assert (
|
||||
-input_.dim() <= dim < input_.dim()
|
||||
), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
|
||||
if dim < 0:
|
||||
# Convert negative dim to positive.
|
||||
dim += input_.dim()
|
||||
|
||||
# Allocate output tensor.
|
||||
if self.rank_in_group == dst:
|
||||
gather_list = [torch.empty_like(input_) for _ in range(world_size)]
|
||||
else:
|
||||
gather_list = None
|
||||
|
||||
# Gather.
|
||||
self.dist_module.gather(
|
||||
input_, gather_list, dst=self.ranks[dst], group=self.device_group
|
||||
)
|
||||
|
||||
if self.rank_in_group == dst:
|
||||
output_tensor = torch.cat(gather_list, dim=dim)
|
||||
else:
|
||||
output_tensor = None
|
||||
return output_tensor
|
||||
|
||||
def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
|
||||
if dim < 0:
|
||||
# Convert negative dim to positive.
|
||||
dim += input_.dim()
|
||||
input_size = input_.size()
|
||||
# NOTE: we have to use concat-style all-gather here,
|
||||
# stack-style all-gather has compatibility issues with
|
||||
# torch.compile . see https://github.com/pytorch/pytorch/issues/138795
|
||||
output_size = (input_size[0] * self.world_size,) + input_size[1:]
|
||||
# Allocate output tensor.
|
||||
output_tensor = torch.empty(
|
||||
output_size, dtype=input_.dtype, device=input_.device
|
||||
)
|
||||
# All-gather.
|
||||
self.dist_module.all_gather_into_tensor(
|
||||
output_tensor, input_, group=self.device_group
|
||||
)
|
||||
|
||||
# Reshape
|
||||
output_tensor = output_tensor.reshape((self.world_size,) + input_size)
|
||||
output_tensor = output_tensor.movedim(0, dim)
|
||||
output_tensor = output_tensor.reshape(
|
||||
input_size[:dim]
|
||||
+ (self.world_size * input_size[dim],)
|
||||
+ input_size[dim + 1 :]
|
||||
)
|
||||
return output_tensor
|
||||
|
||||
|
||||
class _CPUSHMDistributed:
|
||||
|
||||
def __init__(self, communicator: CpuCommunicator):
|
||||
instance_identifier = os.environ["VLLM_DIST_IDENT"]
|
||||
unique_name = communicator.unique_name
|
||||
instance_identifier = f"{instance_identifier}-{unique_name}"
|
||||
self.communicator = communicator
|
||||
|
||||
group_ranks = [str(rank) for rank in self.communicator.ranks]
|
||||
shm_group_identifier = f"[{'-'.join(group_ranks)}]"
|
||||
self.group_name = f"{instance_identifier}-{shm_group_identifier}-cpushm"
|
||||
|
||||
self.handle = self._init_cpu_shm()
|
||||
|
||||
def _init_cpu_shm(self) -> int:
|
||||
handle = torch.ops._C.init_shm_manager(
|
||||
self.group_name,
|
||||
self.communicator.world_size,
|
||||
self.communicator.rank,
|
||||
)
|
||||
torch.distributed.barrier(self.communicator.device_group)
|
||||
torch.ops._C.join_shm_manager(
|
||||
handle,
|
||||
self.group_name,
|
||||
)
|
||||
torch.distributed.barrier(self.communicator.device_group)
|
||||
|
||||
return int(handle)
|
||||
|
||||
def all_reduce(
|
||||
self, input: torch.Tensor, group: ProcessGroup | None = None
|
||||
) -> None:
|
||||
torch.ops._C.shm_allreduce(self.handle, input)
|
||||
|
||||
def gather(
|
||||
self,
|
||||
input: torch.Tensor,
|
||||
gather_list: list[torch.Tensor] | None,
|
||||
dst: int = -1,
|
||||
group: ProcessGroup | None = None,
|
||||
) -> None:
|
||||
# Note: different from the torch gather, here we use local dst rank.
|
||||
torch.ops._C.shm_gather(
|
||||
self.handle,
|
||||
input,
|
||||
gather_list,
|
||||
torch.distributed.get_group_rank(group, dst),
|
||||
)
|
||||
|
||||
def all_gather_into_tensor(
|
||||
self,
|
||||
output: torch.Tensor,
|
||||
input: torch.Tensor,
|
||||
group: ProcessGroup | None = None,
|
||||
) -> None:
|
||||
torch.ops._C.shm_all_gather(self.handle, input, output)
|
||||
@@ -0,0 +1,79 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/cuda_communicator.py
|
||||
|
||||
import torch
|
||||
from torch.distributed import ProcessGroup
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed.device_communicators.base_device_communicator import (
|
||||
DeviceCommunicatorBase,
|
||||
)
|
||||
|
||||
|
||||
class CudaCommunicator(DeviceCommunicatorBase):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cpu_group: ProcessGroup,
|
||||
device: torch.device | None = None,
|
||||
device_group: ProcessGroup | None = None,
|
||||
unique_name: str = "",
|
||||
):
|
||||
super().__init__(cpu_group, device, device_group, unique_name)
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed.device_communicators.pynccl import (
|
||||
PyNcclCommunicator,
|
||||
)
|
||||
|
||||
self.pynccl_comm: PyNcclCommunicator | None = None
|
||||
if self.world_size > 1:
|
||||
self.pynccl_comm = PyNcclCommunicator(
|
||||
group=self.cpu_group,
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
def all_reduce(self, input_, op: torch.distributed.ReduceOp | None = None):
|
||||
pynccl_comm = self.pynccl_comm
|
||||
assert pynccl_comm is not None
|
||||
out = pynccl_comm.all_reduce(input_, op=op)
|
||||
if out is None:
|
||||
# fall back to the default all-reduce using PyTorch.
|
||||
# this usually happens during testing.
|
||||
# when we run the model, allreduce only happens for the TP
|
||||
# group, where we always have either custom allreduce or pynccl.
|
||||
out = input_.clone()
|
||||
torch.distributed.all_reduce(out, group=self.device_group, op=op)
|
||||
return out
|
||||
|
||||
def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
|
||||
"""Sends a tensor to the destination rank in a non-blocking way"""
|
||||
"""NOTE: `dst` is the local rank of the destination rank."""
|
||||
if dst is None:
|
||||
dst = (self.rank_in_group + 1) % self.world_size
|
||||
|
||||
pynccl_comm = self.pynccl_comm
|
||||
if pynccl_comm is not None and not pynccl_comm.disabled:
|
||||
pynccl_comm.send(tensor, dst)
|
||||
else:
|
||||
torch.distributed.send(tensor, self.ranks[dst], self.device_group)
|
||||
|
||||
def recv(
|
||||
self, size: torch.Size, dtype: torch.dtype, src: int | None = None
|
||||
) -> torch.Tensor:
|
||||
"""Receives a tensor from the source rank."""
|
||||
"""NOTE: `src` is the local rank of the source rank."""
|
||||
if src is None:
|
||||
src = (self.rank_in_group - 1) % self.world_size
|
||||
|
||||
tensor = torch.empty(size, dtype=dtype, device=self.device)
|
||||
pynccl_comm = self.pynccl_comm
|
||||
if pynccl_comm is not None and not pynccl_comm.disabled:
|
||||
pynccl_comm.recv(tensor, src)
|
||||
else:
|
||||
torch.distributed.recv(tensor, self.ranks[src], self.device_group)
|
||||
return tensor
|
||||
|
||||
def destroy(self) -> None:
|
||||
if self.pynccl_comm is not None:
|
||||
self.pynccl_comm = None
|
||||
@@ -0,0 +1,258 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl.py
|
||||
|
||||
# ===================== import region =====================
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch.distributed import ProcessGroup, ReduceOp
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed.device_communicators.pynccl_wrapper import (
|
||||
NCCLLibrary,
|
||||
buffer_type,
|
||||
cudaStream_t,
|
||||
ncclComm_t,
|
||||
ncclDataTypeEnum,
|
||||
ncclRedOpTypeEnum,
|
||||
ncclUniqueId,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.distributed.utils import StatelessProcessGroup
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
from sglang.multimodal_gen.utils import current_stream
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class PyNcclCommunicator:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
group: ProcessGroup | StatelessProcessGroup,
|
||||
device: int | str | torch.device,
|
||||
library_path: str | None = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
group: the process group to work on. If None, it will use the
|
||||
default process group.
|
||||
device: the device to bind the PyNcclCommunicator to. If None,
|
||||
it will be bind to f"cuda:{local_rank}".
|
||||
library_path: the path to the NCCL library. If None, it will
|
||||
use the default library path.
|
||||
It is the caller's responsibility to make sure each communicator
|
||||
is bind to a unique device.
|
||||
"""
|
||||
if not isinstance(group, StatelessProcessGroup):
|
||||
assert dist.is_initialized()
|
||||
assert (
|
||||
dist.get_backend(group) != dist.Backend.NCCL
|
||||
), "PyNcclCommunicator should be attached to a non-NCCL group."
|
||||
# note: this rank is the rank in the group
|
||||
self.rank = dist.get_rank(group)
|
||||
self.world_size = dist.get_world_size(group)
|
||||
else:
|
||||
self.rank = group.rank
|
||||
self.world_size = group.world_size
|
||||
|
||||
self.group = group
|
||||
|
||||
# if world_size == 1, no need to create communicator
|
||||
if self.world_size == 1:
|
||||
self.available = False
|
||||
self.disabled = True
|
||||
return
|
||||
try:
|
||||
self.nccl = NCCLLibrary(library_path)
|
||||
except Exception:
|
||||
# disable because of missing NCCL library
|
||||
# e.g. in a non-GPU environment
|
||||
self.available = False
|
||||
self.disabled = True
|
||||
return
|
||||
|
||||
self.available = True
|
||||
self.disabled = False
|
||||
|
||||
logger.info("sgl-diffusion is using nccl==%s", self.nccl.ncclGetVersion())
|
||||
|
||||
if self.rank == 0:
|
||||
# get the unique id from NCCL
|
||||
self.unique_id = self.nccl.ncclGetUniqueId()
|
||||
else:
|
||||
# construct an empty unique id
|
||||
self.unique_id = ncclUniqueId()
|
||||
|
||||
if not isinstance(group, StatelessProcessGroup):
|
||||
tensor = torch.ByteTensor(list(self.unique_id.internal))
|
||||
ranks = dist.get_process_group_ranks(group)
|
||||
# arg `src` in `broadcast` is the global rank
|
||||
dist.broadcast(tensor, src=ranks[0], group=group)
|
||||
byte_list = tensor.tolist()
|
||||
for i, byte in enumerate(byte_list):
|
||||
self.unique_id.internal[i] = byte
|
||||
else:
|
||||
self.unique_id = group.broadcast_obj(self.unique_id, src=0)
|
||||
if isinstance(device, int):
|
||||
device = torch.device(f"cuda:{device}")
|
||||
elif isinstance(device, str):
|
||||
device = torch.device(device)
|
||||
# now `device` is a `torch.device` object
|
||||
assert isinstance(device, torch.device)
|
||||
self.device = device
|
||||
# nccl communicator and stream will use this device
|
||||
# `torch.cuda.device` is a context manager that changes the
|
||||
# current cuda device to the specified one
|
||||
with torch.cuda.device(device):
|
||||
self.comm: ncclComm_t = self.nccl.ncclCommInitRank(
|
||||
self.world_size, self.unique_id, self.rank
|
||||
)
|
||||
|
||||
stream = current_stream()
|
||||
# A small all_reduce for warmup.
|
||||
data = torch.zeros(1, device=device)
|
||||
self.all_reduce(data)
|
||||
if stream is not None:
|
||||
stream.synchronize()
|
||||
del data
|
||||
|
||||
def all_reduce(
|
||||
self, in_tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM, stream=None
|
||||
) -> torch.Tensor:
|
||||
if self.disabled:
|
||||
return None
|
||||
# nccl communicator created on a specific device
|
||||
# will only work on tensors on the same device
|
||||
# otherwise it will cause "illegal memory access"
|
||||
assert in_tensor.device == self.device, (
|
||||
f"this nccl communicator is created to work on {self.device}, "
|
||||
f"but the input tensor is on {in_tensor.device}"
|
||||
)
|
||||
|
||||
out_tensor = torch.empty_like(in_tensor)
|
||||
|
||||
if stream is None:
|
||||
stream = current_stream()
|
||||
self.nccl.ncclAllReduce(
|
||||
buffer_type(in_tensor.data_ptr()),
|
||||
buffer_type(out_tensor.data_ptr()),
|
||||
in_tensor.numel(),
|
||||
ncclDataTypeEnum.from_torch(in_tensor.dtype),
|
||||
ncclRedOpTypeEnum.from_torch(op),
|
||||
self.comm,
|
||||
cudaStream_t(stream.cuda_stream),
|
||||
)
|
||||
return out_tensor
|
||||
|
||||
def all_gather(
|
||||
self, output_tensor: torch.Tensor, input_tensor: torch.Tensor, stream=None
|
||||
):
|
||||
if self.disabled:
|
||||
return
|
||||
# nccl communicator created on a specific device
|
||||
# will only work on tensors on the same device
|
||||
# otherwise it will cause "illegal memory access"
|
||||
assert input_tensor.device == self.device, (
|
||||
f"this nccl communicator is created to work on {self.device}, "
|
||||
f"but the input tensor is on {input_tensor.device}"
|
||||
)
|
||||
if stream is None:
|
||||
stream = current_stream()
|
||||
self.nccl.ncclAllGather(
|
||||
buffer_type(input_tensor.data_ptr()),
|
||||
buffer_type(output_tensor.data_ptr()),
|
||||
input_tensor.numel(),
|
||||
ncclDataTypeEnum.from_torch(input_tensor.dtype),
|
||||
self.comm,
|
||||
cudaStream_t(stream.cuda_stream),
|
||||
)
|
||||
|
||||
def reduce_scatter(
|
||||
self,
|
||||
output_tensor: torch.Tensor,
|
||||
input_tensor: torch.Tensor,
|
||||
op: ReduceOp = ReduceOp.SUM,
|
||||
stream=None,
|
||||
):
|
||||
if self.disabled:
|
||||
return
|
||||
# nccl communicator created on a specific device
|
||||
# will only work on tensors on the same device
|
||||
# otherwise it will cause "illegal memory access"
|
||||
assert input_tensor.device == self.device, (
|
||||
f"this nccl communicator is created to work on {self.device}, "
|
||||
f"but the input tensor is on {input_tensor.device}"
|
||||
)
|
||||
if stream is None:
|
||||
stream = current_stream()
|
||||
self.nccl.ncclReduceScatter(
|
||||
buffer_type(input_tensor.data_ptr()),
|
||||
buffer_type(output_tensor.data_ptr()),
|
||||
output_tensor.numel(),
|
||||
ncclDataTypeEnum.from_torch(input_tensor.dtype),
|
||||
ncclRedOpTypeEnum.from_torch(op),
|
||||
self.comm,
|
||||
cudaStream_t(stream.cuda_stream),
|
||||
)
|
||||
|
||||
def send(self, tensor: torch.Tensor, dst: int, stream=None):
|
||||
if self.disabled:
|
||||
return
|
||||
assert tensor.device == self.device, (
|
||||
f"this nccl communicator is created to work on {self.device}, "
|
||||
f"but the input tensor is on {tensor.device}"
|
||||
)
|
||||
if stream is None:
|
||||
stream = current_stream()
|
||||
self.nccl.ncclSend(
|
||||
buffer_type(tensor.data_ptr()),
|
||||
tensor.numel(),
|
||||
ncclDataTypeEnum.from_torch(tensor.dtype),
|
||||
dst,
|
||||
self.comm,
|
||||
cudaStream_t(stream.cuda_stream),
|
||||
)
|
||||
|
||||
def recv(self, tensor: torch.Tensor, src: int, stream=None):
|
||||
if self.disabled:
|
||||
return
|
||||
assert tensor.device == self.device, (
|
||||
f"this nccl communicator is created to work on {self.device}, "
|
||||
f"but the input tensor is on {tensor.device}"
|
||||
)
|
||||
if stream is None:
|
||||
stream = current_stream()
|
||||
self.nccl.ncclRecv(
|
||||
buffer_type(tensor.data_ptr()),
|
||||
tensor.numel(),
|
||||
ncclDataTypeEnum.from_torch(tensor.dtype),
|
||||
src,
|
||||
self.comm,
|
||||
cudaStream_t(stream.cuda_stream),
|
||||
)
|
||||
|
||||
def broadcast(self, tensor: torch.Tensor, src: int, stream=None):
|
||||
if self.disabled:
|
||||
return
|
||||
assert tensor.device == self.device, (
|
||||
f"this nccl communicator is created to work on {self.device}, "
|
||||
f"but the input tensor is on {tensor.device}"
|
||||
)
|
||||
if stream is None:
|
||||
stream = current_stream()
|
||||
if src == self.rank:
|
||||
sendbuff = buffer_type(tensor.data_ptr())
|
||||
# NCCL requires the sender also to have a receive buffer
|
||||
recvbuff = buffer_type(tensor.data_ptr())
|
||||
else:
|
||||
sendbuff = buffer_type()
|
||||
recvbuff = buffer_type(tensor.data_ptr())
|
||||
self.nccl.ncclBroadcast(
|
||||
sendbuff,
|
||||
recvbuff,
|
||||
tensor.numel(),
|
||||
ncclDataTypeEnum.from_torch(tensor.dtype),
|
||||
src,
|
||||
self.comm,
|
||||
cudaStream_t(stream.cuda_stream),
|
||||
)
|
||||
@@ -0,0 +1,450 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl_wrapper.py
|
||||
|
||||
# This file is a pure Python wrapper for the NCCL library.
|
||||
# The main purpose is to use NCCL combined with CUDA graph.
|
||||
# Before writing this script, we tried the following approach:
|
||||
# 1. We tried to use `cupy`, it calls NCCL correctly, but `cupy` itself
|
||||
# often gets stuck when initializing the NCCL communicator.
|
||||
# 2. We tried to use `torch.distributed`, but `torch.distributed.all_reduce`
|
||||
# contains many other potential cuda APIs, that are not allowed during
|
||||
# capturing the CUDA graph. For further details, please check
|
||||
# https://discuss.pytorch.org/t/pytorch-cudagraph-with-nccl-operation-failed/ .
|
||||
#
|
||||
# Another rejected idea is to write a C/C++ binding for NCCL. It is usually
|
||||
# doable, but we often encounter issues related with nccl versions, and need
|
||||
# to switch between different versions of NCCL. See
|
||||
# https://github.com/NVIDIA/nccl/issues/1234 for more details.
|
||||
# A C/C++ binding is not flexible enough to handle this. It requires
|
||||
# recompilation of the code every time we want to switch between different
|
||||
# versions. This current implementation, with a **pure** Python wrapper, is
|
||||
# more flexible. We can easily switch between different versions of NCCL by
|
||||
# changing the environment variable `SGL_DIFFUSION_NCCL_SO_PATH`, or the `so_file`
|
||||
# variable in the code.
|
||||
|
||||
# TODO(will): support SGL_DIFFUSION_NCCL_SO_PATH
|
||||
|
||||
import ctypes
|
||||
import platform
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.distributed import ReduceOp
|
||||
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
from sglang.multimodal_gen.utils import find_nccl_library
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# === export types and functions from nccl to Python ===
|
||||
# for the original nccl definition, please check
|
||||
# https://github.com/NVIDIA/nccl/blob/master/src/nccl.h.in
|
||||
|
||||
ncclResult_t = ctypes.c_int
|
||||
ncclComm_t = ctypes.c_void_p
|
||||
|
||||
|
||||
class ncclUniqueId(ctypes.Structure):
|
||||
_fields_ = [("internal", ctypes.c_byte * 128)]
|
||||
|
||||
|
||||
cudaStream_t = ctypes.c_void_p
|
||||
buffer_type = ctypes.c_void_p
|
||||
|
||||
ncclDataType_t = ctypes.c_int
|
||||
|
||||
|
||||
class ncclDataTypeEnum:
|
||||
ncclInt8 = 0
|
||||
ncclChar = 0
|
||||
ncclUint8 = 1
|
||||
ncclInt32 = 2
|
||||
ncclInt = 2
|
||||
ncclUint32 = 3
|
||||
ncclInt64 = 4
|
||||
ncclUint64 = 5
|
||||
ncclFloat16 = 6
|
||||
ncclHalf = 6
|
||||
ncclFloat32 = 7
|
||||
ncclFloat = 7
|
||||
ncclFloat64 = 8
|
||||
ncclDouble = 8
|
||||
ncclBfloat16 = 9
|
||||
ncclNumTypes = 10
|
||||
|
||||
@classmethod
|
||||
def from_torch(cls, dtype: torch.dtype) -> int:
|
||||
if dtype == torch.int8:
|
||||
return cls.ncclInt8
|
||||
if dtype == torch.uint8:
|
||||
return cls.ncclUint8
|
||||
if dtype == torch.int32:
|
||||
return cls.ncclInt32
|
||||
if dtype == torch.int64:
|
||||
return cls.ncclInt64
|
||||
if dtype == torch.float16:
|
||||
return cls.ncclFloat16
|
||||
if dtype == torch.float32:
|
||||
return cls.ncclFloat32
|
||||
if dtype == torch.float64:
|
||||
return cls.ncclFloat64
|
||||
if dtype == torch.bfloat16:
|
||||
return cls.ncclBfloat16
|
||||
raise ValueError(f"Unsupported dtype: {dtype}")
|
||||
|
||||
|
||||
ncclRedOp_t = ctypes.c_int
|
||||
|
||||
|
||||
class ncclRedOpTypeEnum:
|
||||
ncclSum = 0
|
||||
ncclProd = 1
|
||||
ncclMax = 2
|
||||
ncclMin = 3
|
||||
ncclAvg = 4
|
||||
ncclNumOps = 5
|
||||
|
||||
@classmethod
|
||||
def from_torch(cls, op: ReduceOp) -> int:
|
||||
if op == ReduceOp.SUM:
|
||||
return cls.ncclSum
|
||||
if op == ReduceOp.PRODUCT:
|
||||
return cls.ncclProd
|
||||
if op == ReduceOp.MAX:
|
||||
return cls.ncclMax
|
||||
if op == ReduceOp.MIN:
|
||||
return cls.ncclMin
|
||||
if op == ReduceOp.AVG:
|
||||
return cls.ncclAvg
|
||||
raise ValueError(f"Unsupported op: {op}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class Function:
|
||||
name: str
|
||||
restype: Any
|
||||
argtypes: list[Any]
|
||||
|
||||
|
||||
class NCCLLibrary:
|
||||
exported_functions = [
|
||||
# const char* ncclGetErrorString(ncclResult_t result)
|
||||
Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]),
|
||||
# ncclResult_t ncclGetVersion(int *version);
|
||||
Function("ncclGetVersion", ncclResult_t, [ctypes.POINTER(ctypes.c_int)]),
|
||||
# ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId);
|
||||
Function("ncclGetUniqueId", ncclResult_t, [ctypes.POINTER(ncclUniqueId)]),
|
||||
# ncclResult_t ncclCommInitRank(
|
||||
# ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank);
|
||||
# note that ncclComm_t is a pointer type, so the first argument
|
||||
# is a pointer to a pointer
|
||||
Function(
|
||||
"ncclCommInitRank",
|
||||
ncclResult_t,
|
||||
[ctypes.POINTER(ncclComm_t), ctypes.c_int, ncclUniqueId, ctypes.c_int],
|
||||
),
|
||||
# ncclResult_t ncclAllReduce(
|
||||
# const void* sendbuff, void* recvbuff, size_t count,
|
||||
# ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
|
||||
# cudaStream_t stream);
|
||||
# note that cudaStream_t is a pointer type, so the last argument
|
||||
# is a pointer
|
||||
Function(
|
||||
"ncclAllReduce",
|
||||
ncclResult_t,
|
||||
[
|
||||
buffer_type,
|
||||
buffer_type,
|
||||
ctypes.c_size_t,
|
||||
ncclDataType_t,
|
||||
ncclRedOp_t,
|
||||
ncclComm_t,
|
||||
cudaStream_t,
|
||||
],
|
||||
),
|
||||
# ncclResult_t ncclAllGather(
|
||||
# const void* sendbuff, void* recvbuff, size_t count,
|
||||
# ncclDataType_t datatype, ncclComm_t comm,
|
||||
# cudaStream_t stream);
|
||||
# note that cudaStream_t is a pointer type, so the last argument
|
||||
# is a pointer
|
||||
Function(
|
||||
"ncclAllGather",
|
||||
ncclResult_t,
|
||||
[
|
||||
buffer_type,
|
||||
buffer_type,
|
||||
ctypes.c_size_t,
|
||||
ncclDataType_t,
|
||||
ncclComm_t,
|
||||
cudaStream_t,
|
||||
],
|
||||
),
|
||||
# ncclResult_t ncclReduceScatter(
|
||||
# const void* sendbuff, void* recvbuff, size_t count,
|
||||
# ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
|
||||
# cudaStream_t stream);
|
||||
# note that cudaStream_t is a pointer type, so the last argument
|
||||
# is a pointer
|
||||
Function(
|
||||
"ncclReduceScatter",
|
||||
ncclResult_t,
|
||||
[
|
||||
buffer_type,
|
||||
buffer_type,
|
||||
ctypes.c_size_t,
|
||||
ncclDataType_t,
|
||||
ncclRedOp_t,
|
||||
ncclComm_t,
|
||||
cudaStream_t,
|
||||
],
|
||||
),
|
||||
# ncclResult_t ncclSend(
|
||||
# const void* sendbuff, size_t count, ncclDataType_t datatype,
|
||||
# int dest, ncclComm_t comm, cudaStream_t stream);
|
||||
Function(
|
||||
"ncclSend",
|
||||
ncclResult_t,
|
||||
[
|
||||
buffer_type,
|
||||
ctypes.c_size_t,
|
||||
ncclDataType_t,
|
||||
ctypes.c_int,
|
||||
ncclComm_t,
|
||||
cudaStream_t,
|
||||
],
|
||||
),
|
||||
# ncclResult_t ncclRecv(
|
||||
# void* recvbuff, size_t count, ncclDataType_t datatype,
|
||||
# int src, ncclComm_t comm, cudaStream_t stream);
|
||||
Function(
|
||||
"ncclRecv",
|
||||
ncclResult_t,
|
||||
[
|
||||
buffer_type,
|
||||
ctypes.c_size_t,
|
||||
ncclDataType_t,
|
||||
ctypes.c_int,
|
||||
ncclComm_t,
|
||||
cudaStream_t,
|
||||
],
|
||||
),
|
||||
# ncclResult_t ncclBroadcast(
|
||||
# const void* sendbuff, void* recvbuff, size_t count,
|
||||
# ncclDataType_t datatype, int root, ncclComm_t comm,
|
||||
# cudaStream_t stream);
|
||||
Function(
|
||||
"ncclBroadcast",
|
||||
ncclResult_t,
|
||||
[
|
||||
buffer_type,
|
||||
buffer_type,
|
||||
ctypes.c_size_t,
|
||||
ncclDataType_t,
|
||||
ctypes.c_int,
|
||||
ncclComm_t,
|
||||
cudaStream_t,
|
||||
],
|
||||
),
|
||||
# be cautious! this is a collective call, it will block until all
|
||||
# processes in the communicator have called this function.
|
||||
# because Python object destruction can happen in random order,
|
||||
# it is better not to call it at all.
|
||||
# ncclResult_t ncclCommDestroy(ncclComm_t comm);
|
||||
Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]),
|
||||
]
|
||||
|
||||
# class attribute to store the mapping from the path to the library
|
||||
# to avoid loading the same library multiple times
|
||||
path_to_library_cache: dict[str, Any] = {}
|
||||
|
||||
# class attribute to store the mapping from library path
|
||||
# to the corresponding dictionary
|
||||
path_to_dict_mapping: dict[str, dict[str, Any]] = {}
|
||||
|
||||
def __init__(self, so_file: str | None = None):
|
||||
|
||||
so_file = so_file or find_nccl_library()
|
||||
|
||||
try:
|
||||
if so_file not in NCCLLibrary.path_to_dict_mapping:
|
||||
lib = ctypes.CDLL(so_file)
|
||||
NCCLLibrary.path_to_library_cache[so_file] = lib
|
||||
self.lib = NCCLLibrary.path_to_library_cache[so_file]
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to load NCCL library from %s ."
|
||||
"It is expected if you are not running on NVIDIA/AMD GPUs."
|
||||
"Otherwise, the nccl library might not exist, be corrupted "
|
||||
"or it does not support the current platform %s."
|
||||
"If you already have the library, please set the "
|
||||
"environment variable SGL_DIFFUSION_NCCL_SO_PATH"
|
||||
" to point to the correct nccl library path.",
|
||||
so_file,
|
||||
platform.platform(),
|
||||
)
|
||||
raise e
|
||||
|
||||
if so_file not in NCCLLibrary.path_to_dict_mapping:
|
||||
_funcs: dict[str, Any] = {}
|
||||
for func in NCCLLibrary.exported_functions:
|
||||
f = getattr(self.lib, func.name)
|
||||
f.restype = func.restype
|
||||
f.argtypes = func.argtypes
|
||||
_funcs[func.name] = f
|
||||
NCCLLibrary.path_to_dict_mapping[so_file] = _funcs
|
||||
self._funcs = NCCLLibrary.path_to_dict_mapping[so_file]
|
||||
|
||||
def ncclGetErrorString(self, result: ncclResult_t) -> str:
|
||||
return str(self._funcs["ncclGetErrorString"](result).decode("utf-8"))
|
||||
|
||||
def NCCL_CHECK(self, result: ncclResult_t) -> None:
|
||||
if result != 0:
|
||||
error_str = self.ncclGetErrorString(result)
|
||||
raise RuntimeError(f"NCCL error: {error_str}")
|
||||
|
||||
def ncclGetVersion(self) -> str:
|
||||
version = ctypes.c_int()
|
||||
self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version)))
|
||||
version_str = str(version.value)
|
||||
# something like 21903 --> "2.19.3"
|
||||
major = version_str[0].lstrip("0")
|
||||
minor = version_str[1:3].lstrip("0")
|
||||
patch = version_str[3:].lstrip("0")
|
||||
return f"{major}.{minor}.{patch}"
|
||||
|
||||
def ncclGetUniqueId(self) -> ncclUniqueId:
|
||||
unique_id = ncclUniqueId()
|
||||
self.NCCL_CHECK(self._funcs["ncclGetUniqueId"](ctypes.byref(unique_id)))
|
||||
return unique_id
|
||||
|
||||
def ncclCommInitRank(
|
||||
self, world_size: int, unique_id: ncclUniqueId, rank: int
|
||||
) -> ncclComm_t:
|
||||
comm = ncclComm_t()
|
||||
self.NCCL_CHECK(
|
||||
self._funcs["ncclCommInitRank"](
|
||||
ctypes.byref(comm), world_size, unique_id, rank
|
||||
)
|
||||
)
|
||||
return comm
|
||||
|
||||
def ncclAllReduce(
|
||||
self,
|
||||
sendbuff: buffer_type,
|
||||
recvbuff: buffer_type,
|
||||
count: int,
|
||||
datatype: int,
|
||||
op: int,
|
||||
comm: ncclComm_t,
|
||||
stream: cudaStream_t,
|
||||
) -> None:
|
||||
# `datatype` actually should be `ncclDataType_t`
|
||||
# and `op` should be `ncclRedOp_t`
|
||||
# both are aliases of `ctypes.c_int`
|
||||
# when we pass int to a function, it will be converted to `ctypes.c_int`
|
||||
# by ctypes automatically
|
||||
self.NCCL_CHECK(
|
||||
self._funcs["ncclAllReduce"](
|
||||
sendbuff, recvbuff, count, datatype, op, comm, stream
|
||||
)
|
||||
)
|
||||
|
||||
def ncclReduceScatter(
|
||||
self,
|
||||
sendbuff: buffer_type,
|
||||
recvbuff: buffer_type,
|
||||
count: int,
|
||||
datatype: int,
|
||||
op: int,
|
||||
comm: ncclComm_t,
|
||||
stream: cudaStream_t,
|
||||
) -> None:
|
||||
# `datatype` actually should be `ncclDataType_t`
|
||||
# and `op` should be `ncclRedOp_t`
|
||||
# both are aliases of `ctypes.c_int`
|
||||
# when we pass int to a function, it will be converted to `ctypes.c_int`
|
||||
# by ctypes automatically
|
||||
self.NCCL_CHECK(
|
||||
self._funcs["ncclReduceScatter"](
|
||||
sendbuff, recvbuff, count, datatype, op, comm, stream
|
||||
)
|
||||
)
|
||||
|
||||
def ncclAllGather(
|
||||
self,
|
||||
sendbuff: buffer_type,
|
||||
recvbuff: buffer_type,
|
||||
count: int,
|
||||
datatype: int,
|
||||
comm: ncclComm_t,
|
||||
stream: cudaStream_t,
|
||||
) -> None:
|
||||
# `datatype` actually should be `ncclDataType_t`
|
||||
# which is an aliases of `ctypes.c_int`
|
||||
# when we pass int to a function, it will be converted to `ctypes.c_int`
|
||||
# by ctypes automatically
|
||||
self.NCCL_CHECK(
|
||||
self._funcs["ncclAllGather"](
|
||||
sendbuff, recvbuff, count, datatype, comm, stream
|
||||
)
|
||||
)
|
||||
|
||||
def ncclSend(
|
||||
self,
|
||||
sendbuff: buffer_type,
|
||||
count: int,
|
||||
datatype: int,
|
||||
dest: int,
|
||||
comm: ncclComm_t,
|
||||
stream: cudaStream_t,
|
||||
) -> None:
|
||||
self.NCCL_CHECK(
|
||||
self._funcs["ncclSend"](sendbuff, count, datatype, dest, comm, stream)
|
||||
)
|
||||
|
||||
def ncclRecv(
|
||||
self,
|
||||
recvbuff: buffer_type,
|
||||
count: int,
|
||||
datatype: int,
|
||||
src: int,
|
||||
comm: ncclComm_t,
|
||||
stream: cudaStream_t,
|
||||
) -> None:
|
||||
self.NCCL_CHECK(
|
||||
self._funcs["ncclRecv"](recvbuff, count, datatype, src, comm, stream)
|
||||
)
|
||||
|
||||
def ncclBroadcast(
|
||||
self,
|
||||
sendbuff: buffer_type,
|
||||
recvbuff: buffer_type,
|
||||
count: int,
|
||||
datatype: int,
|
||||
root: int,
|
||||
comm: ncclComm_t,
|
||||
stream: cudaStream_t,
|
||||
) -> None:
|
||||
self.NCCL_CHECK(
|
||||
self._funcs["ncclBroadcast"](
|
||||
sendbuff, recvbuff, count, datatype, root, comm, stream
|
||||
)
|
||||
)
|
||||
|
||||
def ncclCommDestroy(self, comm: ncclComm_t) -> None:
|
||||
self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm))
|
||||
|
||||
|
||||
__all__ = [
|
||||
"NCCLLibrary",
|
||||
"ncclDataTypeEnum",
|
||||
"ncclRedOpTypeEnum",
|
||||
"ncclUniqueId",
|
||||
"ncclComm_t",
|
||||
"cudaStream_t",
|
||||
"buffer_type",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
1144
python/sglang/multimodal_gen/runtime/distributed/parallel_state.py
Normal file
1144
python/sglang/multimodal_gen/runtime/distributed/parallel_state.py
Normal file
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user