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

[diffusion] refactor: deprecate workloadtype (#15267)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Mick
2025-12-20 11:24:31 +08:00
committed by GitHub
co-authored by gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
parent 5529ab5895
commit c0f9b51992
17 changed files with 130 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/pipeline_configs/base.py
+34 -4
View File
@@ -19,6 +19,7 @@ from sglang.multimodal_gen.configs.models import (
VAEConfig,
)
from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput
from sglang.multimodal_gen.configs.sample.sampling_params import DataType
from sglang.multimodal_gen.configs.utils import update_config_from_args
from sglang.multimodal_gen.runtime.distributed import (
get_sp_parallel_rank,
@@ -36,17 +37,42 @@ from sglang.multimodal_gen.utils import (
logger = init_logger(__name__)
# NOTE: possible duplication with DataType, WorkloadType
# NOTE: possible duplication with DataType
# this may focus on the model's original ability
class ModelTaskType(Enum):
# TODO: check if I2V/TI2V models can work w/wo text
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
TI2I = auto() # Image to Image or Text-Image to Image
def is_image_gen(self):
return self == ModelTaskType.T2I or self == ModelTaskType.I2I
def is_image_gen(self) -> bool:
return (
self == ModelTaskType.T2I
or self == ModelTaskType.I2I
or self == ModelTaskType.TI2I
)
def requires_image_input(self) -> bool:
return self == ModelTaskType.I2V or self == ModelTaskType.I2I
def accepts_image_input(self) -> bool:
return (
self == ModelTaskType.I2V
or self == ModelTaskType.I2I
or self == ModelTaskType.TI2I
or self == ModelTaskType.TI2V
)
def data_type(self) -> DataType:
if self.is_image_gen():
return DataType.IMAGE
else:
return DataType.VIDEO
class STA_Mode(str, Enum):
@@ -121,6 +147,9 @@ class PipelineConfig:
model_path: str = ""
pipeline_config_path: str | None = None
# precision and autocast
enable_autocast: bool = True
# generation parameters
# controls the timestep embedding generation
should_use_guidance: bool = True
@@ -651,11 +680,12 @@ class ImagePipelineConfig(PipelineConfig):
sp_world_size, rank_in_sp_group = get_sp_world_size(), get_sp_parallel_rank()
seq_len = latents.shape[1]
# TODO: reuse code in PipelineConfig::shard_latents_for_sp
# Pad to next multiple of SP degree if needed
if seq_len % sp_world_size != 0:
pad_len = sp_world_size - (seq_len % sp_world_size)
pad = torch.zeros(
(latents.shape[0], pad_len, latents.shape[2]),
(*latents.shape[:1], pad_len, *latents.shape[2:]),
dtype=latents.dtype,
device=latents.device,
)
python/sglang/multimodal_gen/configs/pipeline_configs/flux.py
+3 -1
View File
@@ -53,6 +53,8 @@ class FluxPipelineConfig(ImagePipelineConfig):
# VAE
vae_config: VAEConfig = field(default_factory=FluxVAEConfig)
enable_autocast: bool = False
# Text encoding stage
text_encoder_configs: tuple[EncoderConfig, ...] = field(
default_factory=lambda: (CLIPTextConfig(), T5Config())
@@ -405,7 +407,7 @@ def flux2_pack_latents(latents):
class Flux2PipelineConfig(FluxPipelineConfig):
embedded_cfg_scale: float = 4.0
task_type: ModelTaskType = ModelTaskType.I2I
task_type: ModelTaskType = ModelTaskType.TI2I
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))
python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py
+2
View File
@@ -80,6 +80,8 @@ class QwenImagePipelineConfig(ImagePipelineConfig):
# VAE
vae_config: VAEConfig = field(default_factory=QwenImageVAEConfig)
enable_autocast: bool = False
# Text encoding stage
text_encoder_configs: tuple[EncoderConfig, ...] = field(
default_factory=lambda: (Qwen2_5VLConfig(),)
python/sglang/multimodal_gen/configs/sample/sampling_params.py
+30 -8
View File
@@ -16,7 +16,6 @@ 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 StoreBoolean, align_to
@@ -195,33 +194,55 @@ class SamplingParams:
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
if self.height is None:
self.height_not_provided = True
self._validate()
# Allow env var to override num_inference_steps (for faster CI testing on AMD)
env_steps = os.environ.get("SGLANG_TEST_NUM_INFERENCE_STEPS")
if env_steps is not None and self.num_inference_steps is not None:
self.num_inference_steps = int(env_steps)
def _validate(self):
"""
check if the sampling params is correct by itself
"""
if self.prompt_path and not self.prompt_path.endswith(".txt"):
raise ValueError("prompt_path must be a txt file")
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 _validate_with_pipeline_config(self, pipeline_config):
"""
check if the sampling params is compatible and valid with server_args
"""
if pipeline_config.task_type.requires_image_input():
# requires image input
if self.image_path is None:
raise ValueError(
f"Served model with task type '{pipeline_config.task_type.name}' requires an 'image_path' input, but none was provided"
)
def _adjust(
self,
server_args: ServerArgs,
server_args,
):
"""
final adjustment, called after merged with user params
"""
# TODO: SamplingParams should not rely on ServerArgs
pipeline_config = server_args.pipeline_config
if not isinstance(self.prompt, str):
raise TypeError(f"`prompt` must be a string, but got {type(self.prompt)}")
self.data_type = server_args.pipeline_config.task_type.data_type()
# Process negative prompt
if self.negative_prompt is not None and not self.negative_prompt.isspace():
# avoid stripping default negative prompt: ' ' for qwen-image
@@ -260,7 +281,6 @@ class SamplingParams:
# settle num_frames
logger.debug(f"num_frames set to 1 for image generation model")
self.num_frames = 1
self.data_type = DataType.IMAGE
elif self.adjust_frames:
# NOTE: We must apply adjust_num_frames BEFORE the SP alignment logic below.
# If we apply it after, adjust_num_frames might modify the frame count
@@ -320,7 +340,7 @@ class SamplingParams:
self.num_frames = new_num_frames
self._set_output_file_name()
self.log(server_args=server_args)
self.log(pipeline_config=server_args.pipeline_config)
@classmethod
def from_pretrained(cls, model_path: str, **kwargs) -> "SamplingParams":
@@ -339,6 +359,8 @@ class SamplingParams:
sampling_params._merge_with_user_params(user_sampling_params)
sampling_params._adjust(server_args)
sampling_params._validate_with_pipeline_config(server_args.pipeline_config)
return sampling_params
def output_size_str(self) -> str:
@@ -658,7 +680,7 @@ class SamplingParams:
def output_file_path(self):
return os.path.join(self.output_path, self.output_file_name)
def log(self, server_args: ServerArgs):
def log(self, pipeline_config):
# 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)
@@ -680,9 +702,9 @@ class SamplingParams:
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}
embedded_guidance_scale: {pipeline_config.embedded_cfg_scale}
n_tokens: {self.n_tokens}
flow_shift: {server_args.pipeline_config.flow_shift}
flow_shift: {pipeline_config.flow_shift}
image_path: {self.image_path}
save_output: {self.save_output}
output_file_path: {self.output_file_path()}