chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

diffusion: refactor task type of models (#13118)

Browse Source
This commit is contained in:
Mick
2025-11-12 12:39:34 +08:00
committed by GitHub
parent dd909a511c
commit 33d1aeb07f
9 changed files with 96 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
python/sglang/multimodal_gen/configs/pipelines/base.py
View File

@@ -4,7 +4,7 @@
import json
from collections.abc import Callable
from dataclasses import asdict, dataclass, field, fields
from enum import Enum
from enum import Enum, auto
from typing import Any
import torch
@@ -28,6 +28,19 @@ from sglang.multimodal_gen.utils import (
logger = init_logger(__name__)
# NOTE: possible duplication with DataType, WorkloadType
# this may focus on the model's original ability
class ModelTaskType(Enum):
I2V = auto() # Image to Video
T2V = auto() # Text to Video
TI2V = auto() # Text and Image to Video
T2I = auto() # Text to Image
I2I = auto() # Image to Image
def is_image_task(self):
return self == ModelTaskType.T2I or self == ModelTaskType.I2I
class STA_Mode(str, Enum):
"""STA (Sliding Tile Attention) modes."""
@@ -51,11 +64,11 @@ def postprocess_text(output: BaseEncoderOutput, _text_inputs) -> torch.tensor:
class PipelineConfig:
"""Base configuration for all pipeline architectures."""
task_type: ModelTaskType
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
@@ -113,9 +126,6 @@ class PipelineConfig:
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

8
python/sglang/multimodal_gen/configs/pipelines/flux.py
View File

@@ -13,7 +13,11 @@ from sglang.multimodal_gen.configs.models.encoders import (
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.base import (
ModelTaskType,
PipelineConfig,
preprocess_text,
)
from sglang.multimodal_gen.configs.pipelines.hunyuan import (
clip_postprocess_text,
clip_preprocess_text,
@@ -29,7 +33,7 @@ class FluxPipelineConfig(PipelineConfig):
# FIXME: duplicate with SamplingParams.guidance_scale?
embedded_cfg_scale: float = 3.5
is_image_gen: bool = True
task_type: ModelTaskType = ModelTaskType.T2I
vae_tiling: bool = False

6
python/sglang/multimodal_gen/configs/pipelines/qwen_image.py
View File

@@ -10,7 +10,7 @@ from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAECo
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
from sglang.multimodal_gen.configs.pipelines.base import ModelTaskType, PipelineConfig
def _extract_masked_hidden(hidden_states: torch.Tensor, mask: torch.Tensor):
@@ -64,7 +64,7 @@ def _pack_latents(latents, batch_size, num_channels_latents, height, width):
class QwenImagePipelineConfig(PipelineConfig):
should_use_guidance: bool = False
is_image_gen: bool = True
task_type: ModelTaskType = ModelTaskType.T2I
vae_tiling: bool = False
@@ -194,7 +194,7 @@ class QwenImagePipelineConfig(PipelineConfig):
class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
ti2i_task = True
task_type: ModelTaskType = ModelTaskType.I2I
def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
# TODO: lots of duplications here

7
python/sglang/multimodal_gen/configs/pipelines/wan.py
View File

@@ -14,7 +14,7 @@ from sglang.multimodal_gen.configs.models.encoders import (
T5Config,
)
from sglang.multimodal_gen.configs.models.vaes import WanVAEConfig
from sglang.multimodal_gen.configs.pipelines.base import PipelineConfig
from sglang.multimodal_gen.configs.pipelines.base import ModelTaskType, PipelineConfig
def t5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor:
@@ -37,6 +37,7 @@ def t5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tenso
class WanT2V480PConfig(PipelineConfig):
"""Base configuration for Wan T2V 1.3B pipeline architecture."""
task_type: ModelTaskType = ModelTaskType.T2V
# WanConfig-specific parameters with defaults
# DiT
dit_config: DiTConfig = field(default_factory=WanVideoConfig)
@@ -84,7 +85,7 @@ class WanI2V480PConfig(WanT2V480PConfig):
"""Base configuration for Wan I2V 14B 480P pipeline architecture."""
# WanConfig-specific parameters with defaults
i2v_task: bool = True
task_type: ModelTaskType = ModelTaskType.I2V
# Precision for each component
image_encoder_config: EncoderConfig = field(default_factory=CLIPVisionConfig)
image_encoder_precision: str = "fp32"
@@ -129,7 +130,7 @@ class FastWan2_1_T2V_480P_Config(WanT2V480PConfig):
@dataclass
class Wan2_2_TI2V_5B_Config(WanT2V480PConfig):
flow_shift: float | None = 5.0
ti2v_task: bool = True
task_type: ModelTaskType = ModelTaskType.TI2V
expand_timesteps: bool = True
# ti2v, 5B
vae_stride = (4, 16, 16)

2
python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
View File

@@ -259,7 +259,7 @@ class DiffGenerator:
# TODO: simplify
data_type = (
DataType.IMAGE
if self.server_args.pipeline_config.is_image_gen
if self.server_args.pipeline_config.task_type.is_image_task()
or pretrained_sampling_params.num_frames == 1
else DataType.VIDEO
)

93
python/sglang/multimodal_gen/runtime/entrypoints/utils.py
View File

@@ -50,60 +50,59 @@ def prepare_sampling_params(
f"num_frames={sampling_params.num_frames}"
)
temporal_scale_factor = (
pipeline_config.vae_config.arch_config.temporal_compression_ratio
)
# settle num_frames
if server_args.pipeline_config.is_image_gen:
if pipeline_config.task_type.is_image_task():
# settle num_frames
logger.debug(f"Setting num_frames to 1 because this is a image-gen model")
sampling_params.num_frames = 1
num_frames = sampling_params.num_frames
num_gpus = server_args.num_gpus
use_temporal_scaling_frames = pipeline_config.vae_config.use_temporal_scaling_frames
# Adjust number of frames based on number of GPUs
if use_temporal_scaling_frames:
orig_latent_num_frames = (num_frames - 1) // temporal_scale_factor + 1
else: # stepvideo only
orig_latent_num_frames = sampling_params.num_frames // 17 * 3
if orig_latent_num_frames % server_args.num_gpus != 0:
# Adjust latent frames to be divisible by number of GPUs
if sampling_params.num_frames_round_down:
# Ensure we have at least 1 batch per GPU
new_latent_num_frames = (
max(1, (orig_latent_num_frames // num_gpus)) * num_gpus
)
else:
new_latent_num_frames = (
math.ceil(orig_latent_num_frames / num_gpus) * num_gpus
)
sampling_params.data_type = DataType.IMAGE
else:
# Adjust number of frames based on number of GPUs for video task
use_temporal_scaling_frames = (
pipeline_config.vae_config.use_temporal_scaling_frames
)
num_frames = sampling_params.num_frames
num_gpus = server_args.num_gpus
temporal_scale_factor = (
pipeline_config.vae_config.arch_config.temporal_compression_ratio
)
if use_temporal_scaling_frames:
# Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor
new_num_frames = (new_latent_num_frames - 1) * temporal_scale_factor + 1
orig_latent_num_frames = (num_frames - 1) // temporal_scale_factor + 1
else: # stepvideo only
# Find the least common multiple of 3 and num_gpus
divisor = math.lcm(3, num_gpus)
# Round up to the nearest multiple of this LCM
new_latent_num_frames = (
(new_latent_num_frames + divisor - 1) // divisor
) * divisor
# Convert back to actual frames using the StepVideo formula
new_num_frames = new_latent_num_frames // 3 * 17
orig_latent_num_frames = sampling_params.num_frames // 17 * 3
logger.info(
"Adjusting number of frames from %s to %s based on number of GPUs (%s)",
sampling_params.num_frames,
new_num_frames,
server_args.num_gpus,
)
sampling_params.num_frames = new_num_frames
if orig_latent_num_frames % server_args.num_gpus != 0:
# Adjust latent frames to be divisible by number of GPUs
if sampling_params.num_frames_round_down:
# Ensure we have at least 1 batch per GPU
new_latent_num_frames = (
max(1, (orig_latent_num_frames // num_gpus)) * num_gpus
)
else:
new_latent_num_frames = (
math.ceil(orig_latent_num_frames / num_gpus) * num_gpus
)
if pipeline_config.is_image_gen:
sampling_params.data_type = DataType.IMAGE
if use_temporal_scaling_frames:
# Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor
new_num_frames = (new_latent_num_frames - 1) * temporal_scale_factor + 1
else: # stepvideo only
# Find the least common multiple of 3 and num_gpus
divisor = math.lcm(3, num_gpus)
# Round up to the nearest multiple of this LCM
new_latent_num_frames = (
(new_latent_num_frames + divisor - 1) // divisor
) * divisor
# Convert back to actual frames using the StepVideo formula
new_num_frames = new_latent_num_frames // 3 * 17
logger.info(
"Adjusting number of frames from %s to %s based on number of GPUs (%s)",
sampling_params.num_frames,
new_num_frames,
server_args.num_gpus,
)
sampling_params.num_frames = new_num_frames
sampling_params.set_output_file_ext()
sampling_params.log(server_args=server_args)

22
python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py
View File

@@ -19,7 +19,7 @@ import torch.profiler
from einops import rearrange
from tqdm.auto import tqdm
from sglang.multimodal_gen.configs.pipelines.base import STA_Mode
from sglang.multimodal_gen.configs.pipelines.base import ModelTaskType, STA_Mode
from sglang.multimodal_gen.runtime.distributed import (
cfg_model_parallel_all_reduce,
get_local_torch_device,
@@ -269,7 +269,10 @@ class DenoisingStage(PipelineStage):
None,
)
# FIXME: should probably move to latent preparation stage, to handle with offload
if server_args.pipeline_config.ti2v_task and batch.pil_image is not None:
if (
server_args.pipeline_config.task_type == ModelTaskType.TI2V
and batch.pil_image is not None
):
# Wan2.2 TI2V directly replaces the first frame of the latent with
# the image latent instead of appending along the channel dim
assert batch.image_latent is None, "TI2V task should not have image latents"
@@ -334,7 +337,7 @@ class DenoisingStage(PipelineStage):
# Shard z and reserved_frames_mask for TI2V if SP is enabled
if (
server_args.pipeline_config.ti2v_task
server_args.pipeline_config.task_type == ModelTaskType.TI2V
and batch.pil_image is not None
and get_sp_world_size() > 1
):
@@ -681,7 +684,10 @@ class DenoisingStage(PipelineStage):
):
bsz = batch.raw_latent_shape[0]
# expand timestep
if server_args.pipeline_config.ti2v_task and batch.pil_image is not None:
if (
server_args.pipeline_config.task_type == ModelTaskType.TI2V
and batch.pil_image is not None
):
# Explicitly cast t_device to the target float type at the beginning.
# This ensures any precision-based rounding (e.g., float32(999.0) -> bfloat16(1000.0))
# is applied consistently *before* it's used by any rank.
@@ -723,7 +729,10 @@ class DenoisingStage(PipelineStage):
"""
For Wan2.2 ti2v task, global first frame should be replaced with encoded image after each timestep
"""
if server_args.pipeline_config.ti2v_task and batch.pil_image is not None:
if (
server_args.pipeline_config.task_type == ModelTaskType.TI2V
and batch.pil_image is not None
):
# Apply TI2V mask blending with SP-aware z and reserved_frames_mask.
# This ensures the first frame is always the condition image after each step.
# This is only applied on rank 0, where z is not None.
@@ -814,7 +823,8 @@ class DenoisingStage(PipelineStage):
latent_model_input = latents.to(target_dtype)
if batch.image_latent is not None:
assert (
not server_args.pipeline_config.ti2v_task
not server_args.pipeline_config.task_type
== ModelTaskType.TI2V
), "image latents should not be provided for TI2V task"
latent_model_input = torch.cat(
[latent_model_input, batch.image_latent], dim=1

5
python/sglang/multimodal_gen/runtime/pipelines/stages/input_validation.py
View File

@@ -10,6 +10,7 @@ import torchvision.transforms.functional as TF
from PIL import Image
from sglang.multimodal_gen.configs.pipelines import WanI2V480PConfig
from sglang.multimodal_gen.configs.pipelines.base import ModelTaskType
from sglang.multimodal_gen.configs.pipelines.qwen_image import (
QwenImageEditPipelineConfig,
)
@@ -128,8 +129,8 @@ class InputValidationStage(PipelineStage):
batch.width = width
batch.height = height
elif (
server_args.pipeline_config.ti2v_task
or server_args.pipeline_config.ti2i_task
server_args.pipeline_config.task_type == ModelTaskType.TI2V
or server_args.pipeline_config.task_type == ModelTaskType.I2I
) and batch.pil_image is not None:
# further processing for ti2v task
img = batch.pil_image

2
python/sglang/multimodal_gen/runtime/server_args.py
View File

@@ -800,7 +800,7 @@ class ServerArgs:
def check_server_sp_args(self):
if self.pipeline_config.is_image_gen:
if self.pipeline_config.task_type.is_image_task():
if (
(self.sp_degree and self.sp_degree > 1)
or (self.ulysses_degree and self.ulysses_degree > 1)