From 7be1a8c70cdf2a4e878eca5f6f24b845ba181ef0 Mon Sep 17 00:00:00 2001 From: Adarsh Shirawalmath <114558126+adarshxs@users.noreply.github.com> Date: Tue, 6 Jan 2026 10:00:57 +0530 Subject: [PATCH] [diffusion] feat: support diffusers backend - run any model supported by diffusers (#14112) --- .../configs/pipeline_configs/__init__.py | 4 + .../pipeline_configs/diffusers_generic.py | 84 +++ .../multimodal_gen/configs/sample/__init__.py | 5 +- .../configs/sample/diffusers_generic.py | 52 ++ .../configs/sample/sampling_params.py | 8 + python/sglang/multimodal_gen/registry.py | 119 +++- .../runtime/entrypoints/cli/generate.py | 27 +- .../entrypoints/diffusion_generator.py | 15 +- .../runtime/entrypoints/openai/image_api.py | 31 +- .../runtime/entrypoints/openai/protocol.py | 16 +- .../runtime/entrypoints/openai/video_api.py | 13 +- .../runtime/entrypoints/utils.py | 28 +- .../runtime/managers/gpu_worker.py | 18 +- .../runtime/pipelines/diffusers_pipeline.py | 640 ++++++++++++++++++ .../runtime/pipelines_core/__init__.py | 4 +- .../runtime/pipelines_core/schedule_batch.py | 3 + .../multimodal_gen/runtime/server_args.py | 110 +++ 17 files changed, 1120 insertions(+), 57 deletions(-) create mode 100644 python/sglang/multimodal_gen/configs/pipeline_configs/diffusers_generic.py create mode 100644 python/sglang/multimodal_gen/configs/sample/diffusers_generic.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py index dbb394ba8..65324c5d0 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py @@ -4,6 +4,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import ( PipelineConfig, SlidingTileAttnConfig, ) +from sglang.multimodal_gen.configs.pipeline_configs.diffusers_generic import ( + DiffusersGenericPipelineConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.flux import FluxPipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.flux_finetuned import ( Flux2FinetunedPipelineConfig, @@ -22,6 +25,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.wan import ( from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipelineConfig __all__ = [ + "DiffusersGenericPipelineConfig", "HunyuanConfig", "FastHunyuanConfig", "FluxPipelineConfig", diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/diffusers_generic.py b/python/sglang/multimodal_gen/configs/pipeline_configs/diffusers_generic.py new file mode 100644 index 000000000..3eac5306f --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/diffusers_generic.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Generic pipeline configuration for diffusers backend. + +This module provides a minimal pipeline configuration that works with the diffusers backend. +Since diffusers handles its own model loading and configuration, this config is intentionally minimal. +""" + +from dataclasses import dataclass, field +from typing import Any + +from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig +from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ModelTaskType, + PipelineConfig, +) + + +@dataclass +class DiffusersGenericPipelineConfig(PipelineConfig): + """ + Generic pipeline configuration for diffusers backend. + + This is a minimal configuration since the diffusers backend handles most + configuration internally. It provides sensible defaults for the required fields. + """ + + # default to T2I since it's the most common + task_type: ModelTaskType = ModelTaskType.T2I + + dit_precision: str = "bf16" + vae_precision: str = "bf16" + + should_use_guidance: bool = True + embedded_cfg_scale: float = 1.0 + flow_shift: float | None = None + disable_autocast: bool = True # let diffusers handle dtype + + # diffusers handles its own loading + dit_config: DiTConfig = field(default_factory=DiTConfig) + vae_config: VAEConfig = field(default_factory=VAEConfig) + image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig) + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: (EncoderConfig(),) + ) + text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp16",)) + + # VAE settings + vae_tiling: bool = False # diffusers handles this + vae_slicing: bool = False # slice VAE decode for lower memory usage + vae_sp: bool = False + + # Attention backend for diffusers models (e.g., "flash", "_flash_3_hub", "sage", "xformers") + # See: https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends + diffusers_attention_backend: str | None = None + + # Quantization config for pipeline-level quantization + # See: https://huggingface.co/docs/diffusers/main/en/quantization/overview + # Use PipelineQuantizationConfig for component-level control: + # from diffusers.quantizers import PipelineQuantizationConfig + # quantization_config = PipelineQuantizationConfig( + # quant_backend="bitsandbytes_4bit", + # quant_kwargs={"load_in_4bit": True, "bnb_4bit_compute_dtype": torch.bfloat16}, + # components_to_quantize=["transformer", "text_encoder_2"], + # ) + quantization_config: Any = None + + def check_pipeline_config(self) -> None: + """ + Override to skip most validation since diffusers handles its own config. + """ + pass + + def adjust_size(self, width, height, image): + """ + Pass through - diffusers handles size adjustments. + """ + return width, height + + def adjust_num_frames(self, num_frames): + """ + Pass through - diffusers handles frame count. + """ + return num_frames diff --git a/python/sglang/multimodal_gen/configs/sample/__init__.py b/python/sglang/multimodal_gen/configs/sample/__init__.py index 1d2574de0..0a81bce02 100644 --- a/python/sglang/multimodal_gen/configs/sample/__init__.py +++ b/python/sglang/multimodal_gen/configs/sample/__init__.py @@ -1,5 +1,8 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo +from sglang.multimodal_gen.configs.sample.diffusers_generic import ( + DiffusersGenericSamplingParams, +) from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams -__all__ = ["SamplingParams"] +__all__ = ["SamplingParams", "DiffusersGenericSamplingParams"] diff --git a/python/sglang/multimodal_gen/configs/sample/diffusers_generic.py b/python/sglang/multimodal_gen/configs/sample/diffusers_generic.py new file mode 100644 index 000000000..d30ea5ddb --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/diffusers_generic.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Generic sampling parameters for diffusers backend. + +This module provides generic sampling parameters that work with any diffusers pipeline. +""" + +from dataclasses import dataclass, field +from typing import Any + +from sglang.multimodal_gen.configs.sample.sampling_params import ( + DataType, + SamplingParams, +) + + +@dataclass +class DiffusersGenericSamplingParams(SamplingParams): + """ + Generic sampling parameters for diffusers backend. + + These parameters cover the most common options across different diffusers pipelines. + The diffusers pipeline will use whichever parameters it supports. + + For pipeline-specific parameters, use `diffusers_kwargs` dict which will be + passed directly to the diffusers pipeline call. + """ + + # Override defaults with more conservative values that work across pipelines + num_frames: int = 1 # default to image generation + height: int = 1024 + width: int = 1024 + num_inference_steps: int = 30 + guidance_scale: float = 7.5 + negative_prompt: str = "" + + # extra kwargs to pass directly to the diffusers pipeline + # example: {"output_type": "latent", "return_dict": False} + diffusers_kwargs: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.num_frames > 1: + self.data_type = DataType.VIDEO + else: + self.data_type = DataType.IMAGE + + if self.width is None: + self.width_not_provided = True + self.width = 1024 + if self.height is None: + self.height_not_provided = True + self.height = 1024 diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index d684bdf45..a3145dffc 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -124,6 +124,7 @@ class SamplingParams: num_inference_steps: int = None guidance_scale: float = None guidance_scale_2: float = None + true_cfg_scale: float = None # for CFG vs guidance distillation (e.g., QwenImage) guidance_rescale: float = 0.0 boundary_ratio: float | None = None @@ -580,6 +581,13 @@ class SamplingParams: default=SamplingParams.return_trajectory_decoded, help="Whether to return the decoded trajectory", ) + parser.add_argument( + "--diffusers-kwargs", + type=str, + default=None, + help="JSON string of extra kwargs to pass to diffusers pipeline. " + 'Example: \'{"output_type": "latent", "clip_skip": 2}\'', + ) parser.add_argument( "--no-override-protected-fields", action="store_true", diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 0979304c3..1f779dd5b 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -12,7 +12,20 @@ import importlib import os import pkgutil from functools import lru_cache -from typing import Any, Callable, Dict, List, Optional, Tuple, Type +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Optional, + Tuple, + Type, + Union, +) + +if TYPE_CHECKING: + from sglang.multimodal_gen.runtime.server_args import Backend from sglang.multimodal_gen.configs.pipeline_configs import ( FastHunyuanConfig, @@ -236,8 +249,34 @@ class ModelInfo: pipeline_config_cls: Type[PipelineConfig] +def _get_diffusers_model_info(model_path: str) -> ModelInfo: + """ + Get model info for diffusers backend. + + Returns a ModelInfo with DiffusersPipeline and generic configs. + """ + from sglang.multimodal_gen.configs.pipeline_configs.diffusers_generic import ( + DiffusersGenericPipelineConfig, + ) + from sglang.multimodal_gen.configs.sample.diffusers_generic import ( + DiffusersGenericSamplingParams, + ) + from sglang.multimodal_gen.runtime.pipelines.diffusers_pipeline import ( + DiffusersPipeline, + ) + + return ModelInfo( + pipeline_cls=DiffusersPipeline, + sampling_param_cls=DiffusersGenericSamplingParams, + pipeline_config_cls=DiffusersGenericPipelineConfig, + ) + + @lru_cache(maxsize=1) -def get_model_info(model_path: str) -> Optional[ModelInfo]: +def get_model_info( + model_path: str, + backend: Optional[Union[str, "Backend"]] = None, +) -> Optional[ModelInfo]: """ Resolves all necessary classes (pipeline, sampling, config) for a given model path. @@ -246,7 +285,31 @@ def get_model_info(model_path: str) -> Optional[ModelInfo]: '_class_name' against an auto-discovered registry of pipeline implementations. 2. Resolves the associated configuration classes (for sampling and pipeline) using a manually registered mapping based on the model path. + + Args: + model_path: Path to the model or HuggingFace model ID + backend: Backend to use ('auto', 'sglang', 'diffusers'). If None, uses 'auto'. + + Returns: + ModelInfo with the resolved pipeline class and config classes, or None if not found. """ + # import Backend enum here to avoid circular imports + from sglang.multimodal_gen.runtime.server_args import Backend + + # Normalize backend + if backend is None: + backend = Backend.AUTO + elif isinstance(backend, str): + backend = Backend.from_string(backend) + + # Handle explicit diffusers backend + if backend == Backend.DIFFUSERS: + logger.info( + "Using diffusers backend for model '%s' (explicitly requested)", model_path + ) + return _get_diffusers_model_info(model_path) + + # For AUTO or SGLANG backend, try native implementation first # 1. Discover all available pipeline classes and cache them _discover_and_register_pipelines() @@ -258,32 +321,55 @@ def get_model_info(model_path: str) -> Optional[ModelInfo]: config = maybe_download_model_index(model_path) except Exception as e: logger.error(f"Could not read model config for '{model_path}': {e}") + if backend == Backend.AUTO: + logger.info("Falling back to diffusers backend") + return _get_diffusers_model_info(model_path) return None pipeline_class_name = config.get("_class_name") if not pipeline_class_name: logger.error(f"'_class_name' not found in model_index.json for '{model_path}'") + if backend == Backend.AUTO: + logger.info("Falling back to diffusers backend") + return _get_diffusers_model_info(model_path) return None pipeline_cls = _PIPELINE_REGISTRY.get(pipeline_class_name) if not pipeline_cls: - logger.error( - f"Pipeline class '{pipeline_class_name}' specified in '{model_path}' is not a registered EntryClass in the framework. " - f"Available pipelines: {list(_PIPELINE_REGISTRY.keys())}" - ) - return None + if backend == Backend.AUTO: + logger.warning( + f"Pipeline class '{pipeline_class_name}' specified in '{model_path}' has no native sglang support. " + f"Falling back to diffusers backend." + ) + return _get_diffusers_model_info(model_path) + else: + logger.error( + f"Pipeline class '{pipeline_class_name}' specified in '{model_path}' is not a registered EntryClass in the framework. " + f"Available pipelines: {list(_PIPELINE_REGISTRY.keys())}. " + f"Consider using --backend diffusers to use vanilla diffusers pipeline." + ) + return None # 3. Get configuration classes (sampling, pipeline config) config_info = _get_config_info(model_path) if not config_info: - logger.error( - f"Could not resolve configuration for model '{model_path}'. " - "It is not a registered model path or detected by any registered model family detectors. " - f"Known model paths: {list(_MODEL_HF_PATH_TO_NAME.keys())}" - ) - return None + if backend == Backend.AUTO: + logger.warning( + f"Could not resolve native configuration for model '{model_path}'. " + f"Falling back to diffusers backend." + ) + return _get_diffusers_model_info(model_path) + else: + logger.error( + f"Could not resolve configuration for model '{model_path}'. " + "It is not a registered model path or detected by any registered model family detectors. " + f"Known model paths: {list(_MODEL_HF_PATH_TO_NAME.keys())}. " + f"Consider using --backend diffusers to use vanilla diffusers pipeline." + ) + return None - # 4. Combine the complete model info + # 4. Combine and return the complete model info + logger.info("Using native sglang backend for model '%s'", model_path) model_info = ModelInfo( pipeline_cls=pipeline_cls, sampling_param_cls=config_info.sampling_param_cls, @@ -312,7 +398,6 @@ def _register_configs(): "FastVideo/FastHunyuan-diffusers", ], ) - # Wan register_configs( sampling_param_cls=WanT2V_1_3B_SamplingParams, @@ -372,7 +457,6 @@ def _register_configs(): "Wan-AI/Wan2.2-TI2V-5B-Diffusers", ], ) - register_configs( sampling_param_cls=Wan2_2_TI2V_5B_SamplingParam, pipeline_config_cls=FastWan2_2_TI2V_5B_Config, @@ -381,7 +465,6 @@ def _register_configs(): "FastVideo/FastWan2.2-TI2V-5B-Diffusers", ], ) - register_configs( sampling_param_cls=Wan2_2_T2V_A14B_SamplingParam, pipeline_config_cls=Wan2_2_T2V_A14B_Config, @@ -399,7 +482,6 @@ def _register_configs(): "FastVideo/FastWan2.1-T2V-1.3B-Diffusers", ], ) - # FLUX register_configs( sampling_param_cls=FluxSamplingParams, @@ -425,7 +507,6 @@ def _register_configs(): ], model_detectors=[lambda hf_id: "z-image" in hf_id.lower()], ) - # Qwen-Image register_configs( sampling_param_cls=QwenImageSamplingParams, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py b/python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py index 3d8628b66..fb6554f4b 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py @@ -5,11 +5,15 @@ import argparse import dataclasses +import json import os from typing import cast from sglang.multimodal_gen import DiffGenerator -from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams +from sglang.multimodal_gen.configs.sample.sampling_params import ( + SamplingParams, + generate_request_id, +) from sglang.multimodal_gen.runtime.entrypoints.cli.cli_types import CLISubcommand from sglang.multimodal_gen.runtime.entrypoints.cli.utils import ( RaiseNotImplementedAction, @@ -89,14 +93,33 @@ def generate_cmd(args: argparse.Namespace): args.request_id = "mocked_fake_id_for_offline_generate" server_args = ServerArgs.from_cli_args(args) + sampling_params_kwargs = SamplingParams.get_cli_args(args) + sampling_params_kwargs["request_id"] = generate_request_id() + + # Handle diffusers-specific kwargs passed via CLI + if hasattr(args, "diffusers_kwargs") and args.diffusers_kwargs: + try: + sampling_params_kwargs["diffusers_kwargs"] = json.loads( + args.diffusers_kwargs + ) + logger.info( + "Parsed diffusers_kwargs: %s", + sampling_params_kwargs["diffusers_kwargs"], + ) + except json.JSONDecodeError as e: + logger.error("Failed to parse --diffusers-kwargs as JSON: %s", e) + raise ValueError( + f"--diffusers-kwargs must be valid JSON. Got: {args.diffusers_kwargs}" + ) from e + generator = DiffGenerator.from_pretrained( model_path=server_args.model_path, server_args=server_args, local_mode=True ) results = generator.generate(sampling_params_kwargs=sampling_params_kwargs) - prompt = sampling_params_kwargs.get("prompt", None) + prompt = sampling_params_kwargs.get("prompt") maybe_dump_performance(args, server_args, prompt, results) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 8551e1ba5..ce3f836f0 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -200,15 +200,20 @@ class DiffGenerator: **sampling_params_kwargs, ) + # Extract diffusers_kwargs if passed + diffusers_kwargs = sampling_params_kwargs.pop("diffusers_kwargs", None) + requests: list[Req] = [] for output_idx, p in enumerate(prompts): sampling_params.prompt = p - requests.append( - prepare_request( - server_args=self.server_args, - sampling_params=sampling_params, - ) + req = prepare_request( + server_args=self.server_args, + sampling_params=sampling_params, ) + # Add diffusers_kwargs to request's extra dict + if diffusers_kwargs: + req.extra["diffusers_kwargs"] = diffusers_kwargs + requests.append(req) results = [] total_start_time = time.perf_counter() diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index 31921ebcd..e5dea3358 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py @@ -59,9 +59,10 @@ def _build_sampling_params_from_request( image_path: Optional[list[str]] = None, seed: Optional[int] = None, generator_device: Optional[str] = None, - negative_prompt: Optional[str] = None, - guidance_scale: Optional[float] = None, num_inference_steps: Optional[int] = None, + guidance_scale: Optional[float] = None, + true_cfg_scale: Optional[float] = None, + negative_prompt: Optional[str] = None, enable_teacache: Optional[bool] = None, num_frames: int = 1, ) -> SamplingParams: @@ -70,14 +71,14 @@ def _build_sampling_params_from_request( else: width, height = _parse_size(size) ext = _choose_ext(output_format, background) + server_args = get_global_server_args() - # Build user params sampling_params = SamplingParams.from_user_sampling_params_args( model_path=server_args.model_path, request_id=request_id, prompt=prompt, image_path=image_path, - num_frames=num_frames, # image + num_frames=num_frames, width=width, height=height, num_outputs_per_prompt=max(1, min(int(n or 1), 10)), @@ -90,7 +91,16 @@ def _build_sampling_params_from_request( num_inference_steps=num_inference_steps, enable_teacache=enable_teacache, **({"negative_prompt": negative_prompt} if negative_prompt is not None else {}), + **({"true_cfg_scale": true_cfg_scale} if true_cfg_scale is not None else {}), ) + + if num_inference_steps is not None: + sampling_params.num_inference_steps = num_inference_steps + if guidance_scale is not None: + sampling_params.guidance_scale = guidance_scale + if seed is not None: + sampling_params.seed = seed + return sampling_params @@ -103,6 +113,7 @@ def _build_req_from_sampling(s: SamplingParams) -> Req: async def generations( request: ImageGenerationsRequest, ): + request_id = generate_request_id() sampling = _build_sampling_params_from_request( request_id=request_id, @@ -113,15 +124,19 @@ async def generations( background=request.background, seed=request.seed, generator_device=request.generator_device, - negative_prompt=request.negative_prompt, - guidance_scale=request.guidance_scale, num_inference_steps=request.num_inference_steps, + guidance_scale=request.guidance_scale, + true_cfg_scale=request.true_cfg_scale, + negative_prompt=request.negative_prompt, enable_teacache=request.enable_teacache, ) batch = prepare_request( server_args=get_global_server_args(), sampling_params=sampling, ) + # Add diffusers_kwargs if provided + if request.diffusers_kwargs: + batch.extra["diffusers_kwargs"] = request.diffusers_kwargs # Run synchronously for images and save to disk save_file_path_list, result = await process_generation_batch( @@ -180,6 +195,7 @@ async def edits( user: Optional[str] = Form(None), negative_prompt: Optional[str] = Form(None), guidance_scale: Optional[float] = Form(None), + true_cfg_scale: Optional[float] = Form(None), num_inference_steps: Optional[int] = Form(None), enable_teacache: Optional[bool] = Form(False), num_frames: int = Form(1), @@ -224,9 +240,10 @@ async def edits( generator_device=generator_device, negative_prompt=negative_prompt, guidance_scale=guidance_scale, + true_cfg_scale=true_cfg_scale, num_inference_steps=num_inference_steps, enable_teacache=enable_teacache, - num_frames=num_frames, # image + num_frames=num_frames, ) batch = _build_req_from_sampling(sampling) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py index d606946ae..165c41606 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py @@ -28,13 +28,18 @@ class ImageGenerationsRequest(BaseModel): style: Optional[str] = "vivid" background: Optional[str] = "auto" # transparent | opaque | auto output_format: Optional[str] = None # png | jpeg | webp + user: Optional[str] = None + # SGLang extensions + num_inference_steps: Optional[int] = None + guidance_scale: Optional[float] = None + true_cfg_scale: Optional[float] = ( + None # for CFG vs guidance distillation (e.g., QwenImage) + ) seed: Optional[int] = 1024 generator_device: Optional[str] = "cuda" - user: Optional[str] = None negative_prompt: Optional[str] = None - guidance_scale: Optional[float] = None - num_inference_steps: Optional[int] = None enable_teacache: Optional[bool] = False + diffusers_kwargs: Optional[Dict[str, Any]] = None # kwargs for diffusers backend # Video API protocol models @@ -65,11 +70,16 @@ class VideoGenerationsRequest(BaseModel): num_frames: Optional[int] = None seed: Optional[int] = 1024 generator_device: Optional[str] = "cuda" + # SGLang extensions num_inference_steps: Optional[int] = None guidance_scale: Optional[float] = None guidance_scale_2: Optional[float] = None + true_cfg_scale: Optional[float] = ( + None # for CFG vs guidance distillation (e.g., QwenImage) + ) negative_prompt: Optional[str] = None enable_teacache: Optional[bool] = False + diffusers_kwargs: Optional[Dict[str, Any]] = None # kwargs for diffusers backend class VideoListResponse(BaseModel): diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py index 08982b805..bd571b09d 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py @@ -54,14 +54,13 @@ def _build_sampling_params_from_request( else: width, height = _parse_size(request.size) seconds = request.seconds if request.seconds is not None else 4 - # Prefer user-provided fps/num_frames from request; fallback to defaults fps_default = 24 fps = request.fps if request.fps is not None else fps_default - # If user provides num_frames, use it directly; otherwise derive from seconds * fps derived_num_frames = fps * seconds num_frames = ( request.num_frames if request.num_frames is not None else derived_num_frames ) + server_args = get_global_server_args() sampling_kwargs = { "request_id": request_id, @@ -92,6 +91,13 @@ def _build_sampling_params_from_request( **sampling_kwargs, ) + if request.num_inference_steps is not None: + sampling_params.num_inference_steps = request.num_inference_steps + if request.guidance_scale is not None: + sampling_params.guidance_scale = request.guidance_scale + if request.seed is not None: + sampling_params.seed = request.seed + return sampling_params @@ -261,6 +267,9 @@ async def create_video( server_args=get_global_server_args(), sampling_params=sampling_params, ) + # Add diffusers_kwargs if provided + if req.diffusers_kwargs: + batch.extra["diffusers_kwargs"] = req.diffusers_kwargs # Enqueue the job asynchronously and return immediately asyncio.create_task(_dispatch_job_async(request_id, batch)) return VideoResponse(**job) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index b2166ebef..02ab1d465 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -30,24 +30,30 @@ def prepare_request( server_args: ServerArgs, sampling_params: SamplingParams, ) -> Req: - """ - Settle SamplingParams according to ServerArgs - - """ - # Create a copy of inference args to avoid modifying the original. - # Filter out fields not defined in Req to avoid unexpected-kw TypeError. params_dict = shallow_asdict(sampling_params) - req_field_names = {f.name for f in dataclasses.fields(Req)} - filtered_params = {k: v for k, v in params_dict.items() if k in req_field_names} - req = Req(**filtered_params, VSA_sparsity=server_args.VSA_sparsity) + + diffusers_kwargs = params_dict.pop("diffusers_kwargs", None) + extra = params_dict.get("extra") or {} + if diffusers_kwargs: + extra["diffusers_kwargs"] = diffusers_kwargs + params_dict["extra"] = extra + + # Filter to only fields that exist in Req + req_fields = {f.name for f in dataclasses.fields(Req)} + filtered_params = {k: v for k, v in params_dict.items() if k in req_fields} + + req = Req( + **filtered_params, + VSA_sparsity=server_args.VSA_sparsity, + ) + req.adjust_size(server_args) if (req.width is not None and req.width <= 0) or ( req.height is not None and req.height <= 0 ): raise ValueError( - f"Height, width must be positive integers, got " - f"height={req.height}, width={req.width}" + f"Height and width must be positive, got height={req.height}, width={req.width}" ) return req diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index b31af0dee..a2210e6ad 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -116,7 +116,6 @@ class GPUWorker: Execute a forward pass. """ assert self.pipeline is not None - # TODO: dealing with first req for now req = batch[0] output_batch = None try: @@ -125,7 +124,18 @@ class GPUWorker: start_time = time.monotonic() - output_batch = self.pipeline.forward(req, self.server_args) + result = self.pipeline.forward(req, self.server_args) + + if isinstance(result, Req): + output_batch = OutputBatch( + output=result.output, + timings=result.timings, + trajectory_timesteps=getattr(result, "trajectory_timesteps", None), + trajectory_latents=getattr(result, "trajectory_latents", None), + trajectory_decoded=getattr(result, "trajectory_decoded", None), + ) + else: + output_batch = result if self.rank == 0: peak_memory_bytes = torch.cuda.max_memory_allocated() @@ -155,10 +165,6 @@ class GPUWorker: f"Error executing request {req.request_id}: {e}", exc_info=True ) if output_batch is None: - from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import ( - OutputBatch, - ) - output_batch = OutputBatch() output_batch.error = f"Error executing request {req.request_id}: {e}" finally: diff --git a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py new file mode 100644 index 000000000..a962063f2 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py @@ -0,0 +1,640 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Diffusers backend pipeline wrapper. + +This module provides a wrapper that allows running any diffusers-supported model +through sglang's infrastructure using vanilla diffusers pipelines. +""" + +import argparse +import inspect +import re +import warnings +from io import BytesIO +from typing import Any + +import numpy as np +import requests +import torch +import torchvision.transforms as T +from diffusers import DiffusionPipeline +from PIL import Image + +from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig +from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import ( + PipelineExecutor, +) +from sglang.multimodal_gen.runtime.pipelines_core.executors.sync_executor import ( + SyncExecutor, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages import PipelineStage +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__) + + +class DiffusersExecutionStage(PipelineStage): + """Pipeline stage that wraps diffusers pipeline execution.""" + + def __init__(self, diffusers_pipe: DiffusionPipeline): + super().__init__() + self.diffusers_pipe = diffusers_pipe + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + """Execute the diffusers pipeline.""" + + kwargs = self._build_pipeline_kwargs(batch, server_args) + + # Filter kwargs to only those supported by the pipeline, warn about ignored args + kwargs, _ = self._filter_pipeline_kwargs(kwargs) + + # Request tensor output for cleaner handling + if "output_type" not in kwargs: + kwargs["output_type"] = "pt" + + with torch.no_grad(), warnings.catch_warnings(record=True): + warnings.simplefilter("always") + try: + output = self.diffusers_pipe(**kwargs) + except TypeError as e: + # Some pipelines don't support output_type="pt" + if "output_type" in str(e): + kwargs.pop("output_type", None) + output = self.diffusers_pipe(**kwargs) + else: + raise + + batch.output = self._extract_output(output) + if batch.output is not None: + batch.output = self._postprocess_output(batch.output) + + return batch + + def _filter_pipeline_kwargs( + self, kwargs: dict, *, strict: bool = False + ) -> tuple[dict, list[str]]: + """Filter kwargs to those accepted by the pipeline's __call__. + + Args: + kwargs: Arguments to filter + strict: If True, raise ValueError on unsupported args; otherwise warn + + Returns: + Tuple of (filtered_kwargs, ignored_keys) + """ + try: + sig = inspect.signature(self.diffusers_pipe.__call__) + except (ValueError, TypeError): + return kwargs, [] + + params = sig.parameters + accepts_var_kwargs = any( + p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values() + ) + if accepts_var_kwargs: + return kwargs, [] + + valid = set(params.keys()) - {"self"} + + filtered = {} + ignored = [] + for k, v in kwargs.items(): + if k in valid: + filtered[k] = v + else: + ignored.append(k) + + if ignored: + pipe_name = type(self.diffusers_pipe).__name__ + msg = ( + f"Pipeline '{pipe_name}' does not support: {', '.join(sorted(ignored))}. " + "These arguments will be ignored." + ) + if strict: + raise ValueError(msg) + logger.warning(msg) + + return filtered, ignored + + def _extract_output(self, output: Any) -> torch.Tensor | None: + """Extract tensor output from pipeline result.""" + for attr in ["images", "frames", "video", "sample", "pred_original_sample"]: + if not hasattr(output, attr): + continue + + data = getattr(output, attr) + if data is None: + continue + + result = self._convert_to_tensor(data) + if result is not None: + logger.info( + "Extracted output from '%s': shape=%s, dtype=%s", + attr, + result.shape, + result.dtype, + ) + return result + + logger.warning("Could not extract output from pipeline result") + return None + + def _convert_to_tensor(self, data: Any) -> torch.Tensor | None: + """Convert various data formats to a tensor.""" + if isinstance(data, torch.Tensor): + return data + + if isinstance(data, np.ndarray): + tensor = torch.from_numpy(data).float() + if tensor.max() > 1.0: + tensor = tensor / 255.0 + # (B, H, W, C) -> (B, C, H, W) or (B, T, H, W, C) -> (B, C, T, H, W) + if tensor.ndim == 4: + tensor = tensor.permute(0, 3, 1, 2) + elif tensor.ndim == 5: + tensor = tensor.permute(0, 4, 1, 2, 3) + return tensor + + if hasattr(data, "mode"): # PIL Image + return T.ToTensor()(data) + + if isinstance(data, list) and len(data) > 0: + return self._convert_list_to_tensor(data) + + return None + + def _convert_list_to_tensor(self, data: list) -> torch.Tensor | None: + """Convert a list of items to a tensor.""" + first = data[0] + + # Nested list (e.g., [[frame1, frame2, ...]] for video batches) + if isinstance(first, list) and len(first) > 0: + data = first + first = data[0] + + if hasattr(first, "mode"): # PIL images + tensors = [T.ToTensor()(img) for img in data] + stacked = torch.stack(tensors) + if len(tensors) > 1: + return stacked.permute(1, 0, 2, 3) # (T, C, H, W) -> (C, T, H, W) + return stacked[0] + + if isinstance(first, torch.Tensor): + stacked = torch.stack(data) + if len(data) > 1: + return stacked.permute(1, 0, 2, 3) + return stacked[0] + + if isinstance(first, np.ndarray): + tensors = [torch.from_numpy(arr).float() for arr in data] + if tensors[0].max() > 1.0: + tensors = [t / 255.0 for t in tensors] + if tensors[0].ndim == 3: + tensors = [t.permute(2, 0, 1) for t in tensors] + stacked = torch.stack(tensors) + if len(data) > 1: + return stacked.permute(1, 0, 2, 3) + return stacked[0] + + return None + + def _postprocess_output(self, output: torch.Tensor) -> torch.Tensor: + """Post-process output tensor to ensure valid values and correct shape.""" + output = output.cpu().float() + + # Handle NaN or Inf values + if torch.isnan(output).any() or torch.isinf(output).any(): + logger.warning("Output contains invalid values, fixing...") + output = torch.nan_to_num(output, nan=0.5, posinf=1.0, neginf=0.0) + + # Normalize to [0, 1] range if needed + min_val, max_val = output.min().item(), output.max().item() + if min_val < -0.5 or max_val > 1.5: + output = (output + 1) / 2 + + output = output.clamp(0, 1) + + # Ensure correct shape for downstream processing + output = self._fix_output_shape(output) + + logger.info("Final output tensor shape: %s", output.shape) + return output + + def _fix_output_shape(self, output: torch.Tensor) -> torch.Tensor: + """Fix tensor shape for downstream processing. + + Expected: (B, C, H, W) for images or (B, C, T, H, W) for videos. + """ + if output.dim() == 5: + # Video: (B, T, C, H, W) -> (B, C, T, H, W) + return output.permute(0, 2, 1, 3, 4) + + if output.dim() == 4: + if output.shape[0] == 1 or output.shape[1] in [1, 3, 4]: + return output # Already (B, C, H, W) + # (T, C, H, W) -> (1, C, T, H, W) + return output.unsqueeze(0).permute(0, 2, 1, 3, 4) + + if output.dim() == 3: + c, h, w = output.shape + if c > 4 and w <= 4: + output = output.permute(2, 0, 1) + if output.shape[0] == 1: + output = output.repeat(3, 1, 1) + return output.unsqueeze(0) + + if output.dim() == 2: + return output.unsqueeze(0).repeat(3, 1, 1).unsqueeze(0) + + return output + + def _build_pipeline_kwargs(self, batch: Req, server_args: ServerArgs) -> dict: + """Build kwargs dict for diffusers pipeline call.""" + kwargs = {} + + if batch.prompt is not None: + kwargs["prompt"] = batch.prompt + + if batch.negative_prompt: + kwargs["negative_prompt"] = batch.negative_prompt + + if batch.num_inference_steps is not None: + kwargs["num_inference_steps"] = batch.num_inference_steps + + if batch.guidance_scale is not None: + kwargs["guidance_scale"] = batch.guidance_scale + + if batch.true_cfg_scale is not None: + kwargs["true_cfg_scale"] = batch.true_cfg_scale + + if batch.height is not None: + kwargs["height"] = batch.height + + if batch.width is not None: + kwargs["width"] = batch.width + + if batch.num_frames is not None and batch.num_frames > 1: + kwargs["num_frames"] = batch.num_frames + + # Generator for reproducibility + if batch.generator is not None: + kwargs["generator"] = batch.generator + elif batch.seed is not None: + device = self._get_pipeline_device() + kwargs["generator"] = torch.Generator(device=device).manual_seed(batch.seed) + + # Image input for img2img or inpainting + image = self._load_input_image(batch) + if image is not None: + kwargs["image"] = image + + if batch.num_outputs_per_prompt > 1: + kwargs["num_images_per_prompt"] = batch.num_outputs_per_prompt + + # Extra diffusers-specific kwargs + if batch.extra: + diffusers_kwargs = batch.extra.get("diffusers_kwargs", {}) + if diffusers_kwargs: + kwargs.update(diffusers_kwargs) + + return kwargs + + def _get_pipeline_device(self) -> str: + """Get the device the pipeline is running on.""" + for attr in ["unet", "transformer", "vae"]: + component = getattr(self.diffusers_pipe, attr, None) + if component is not None: + try: + return next(component.parameters()).device + except StopIteration: + pass + return "cuda" if torch.cuda.is_available() else "cpu" + + def _load_input_image(self, batch: Req) -> Image.Image | None: + """Load input image from batch.""" + # Check for PIL image in condition_image or pixel_values + if batch.condition_image is not None and isinstance( + batch.condition_image, Image.Image + ): + return batch.condition_image + if batch.pixel_values is not None and isinstance( + batch.pixel_values, Image.Image + ): + return batch.pixel_values + + if not batch.image_path: + return None + + try: + if batch.image_path.startswith(("http://", "https://")): + response = requests.get(batch.image_path, timeout=30) + response.raise_for_status() + return Image.open(BytesIO(response.content)).convert("RGB") + return Image.open(batch.image_path).convert("RGB") + except Exception as e: + logger.error("Failed to load image from %s: %s", batch.image_path, e) + return None + + +class DiffusersPipeline(ComposedPipelineBase): + """ + Pipeline wrapper that uses vanilla diffusers pipelines. + + This allows running any diffusers-supported model through sglang's infrastructure + without requiring native sglang implementation. + """ + + pipeline_name = "DiffusersPipeline" + is_video_pipeline = False + _required_config_modules: list[str] = [] + + def __init__( + self, + model_path: str, + server_args: ServerArgs, + required_config_modules: list[str] | None = None, + loaded_modules: dict[str, torch.nn.Module] | None = None, + executor: PipelineExecutor | None = None, + ): + self.server_args = server_args + self.model_path = model_path + self._stages: list[PipelineStage] = [] + self._stage_name_mapping: dict[str, PipelineStage] = {} + self.modules: dict[str, Any] = {} + self.memory_usages: dict[str, float] = {} + self.post_init_called = False + self.executor = executor or SyncExecutor(server_args=server_args) + + logger.info("Loading diffusers pipeline from %s", model_path) + self.diffusers_pipe = self._load_diffusers_pipeline(model_path, server_args) + self._detect_pipeline_type() + + def _load_diffusers_pipeline(self, model_path: str, server_args: ServerArgs) -> Any: + """Load the diffusers pipeline. + + Optimizations applied: + - device_map: Loads models directly to GPU, warming up CUDA caching allocator + to avoid small tensor allocations during inference. + - Parallel shard loading: When using device_map with accelerate, model shards + are loaded in parallel for faster initialization. + """ + + original_model_path = model_path # Keep original for custom_pipeline + model_path = maybe_download_model(model_path) + self.model_path = model_path + + dtype = self._get_dtype(server_args) + device_map = self._get_device_map(server_args) + logger.info( + "Loading diffusers pipeline with dtype=%s, device_map=%s", + dtype, + device_map, + ) + + # Build common kwargs for from_pretrained + load_kwargs = { + "torch_dtype": dtype, + "trust_remote_code": server_args.trust_remote_code, + "revision": server_args.revision, + } + + # Add device_map for direct GPU loading and parallel shard loading + # This warms up CUDA caching allocator and enables parallel loading via accelerate + if device_map is not None: + load_kwargs["device_map"] = device_map + + # Add quantization config if provided (e.g., BitsAndBytesConfig for 4/8-bit) + config = getattr(server_args, "pipeline_config", None) + if config is not None: + quant_config = getattr(config, "quantization_config", None) + if quant_config is not None: + load_kwargs["quantization_config"] = quant_config + logger.info( + "Using quantization config: %s", type(quant_config).__name__ + ) + + try: + pipe = DiffusionPipeline.from_pretrained(model_path, **load_kwargs) + except AttributeError as e: + if "has no attribute" in str(e): + # Custom pipeline class not in diffusers - try loading with custom_pipeline + logger.info( + "Pipeline class not found in diffusers, trying custom_pipeline from repo..." + ) + try: + custom_kwargs = { + **load_kwargs, + "custom_pipeline": original_model_path, + } + custom_kwargs["trust_remote_code"] = True + pipe = DiffusionPipeline.from_pretrained( + model_path, **custom_kwargs + ) + except Exception as e2: + match = re.search(r"has no attribute (\w+)", str(e)) + class_name = match.group(1) if match else "unknown" + raise RuntimeError( + f"Pipeline class '{class_name}' not found in diffusers and no custom pipeline.py in repo. " + f"Try: pip install --upgrade diffusers (some pipelines require latest version). " + f"Original error: {e}" + ) from e2 + else: + raise + except Exception as e: + # Only retry with float32 for dtype-related errors + if "dtype" in str(e).lower() or "float" in str(e).lower(): + logger.warning( + "Failed with dtype=%s, falling back to float32: %s", dtype, e + ) + load_kwargs["torch_dtype"] = torch.float32 + pipe = DiffusionPipeline.from_pretrained(model_path, **load_kwargs) + else: + raise + + # Only move to device if device_map wasn't used (already on device) + if device_map is None: + if torch.cuda.is_available(): + pipe = pipe.to("cuda") + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + pipe = pipe.to("mps") + + # Apply VAE memory optimizations from pipeline config + self._apply_vae_optimizations(pipe, server_args) + + # Apply attention backend if specified + self._apply_attention_backend(pipe, server_args) + + logger.info("Loaded diffusers pipeline: %s", pipe.__class__.__name__) + return pipe + + def _apply_vae_optimizations(self, pipe: Any, server_args: ServerArgs) -> None: + """Apply VAE memory optimizations (tiling, slicing) from pipeline config.""" + config = getattr(server_args, "pipeline_config", None) + if config is None: + return + + # VAE slicing: decode latents slice-by-slice for lower peak memory + # https://huggingface.co/docs/diffusers/optimization/memory#vae-slicing + if getattr(config, "vae_slicing", False): + if hasattr(pipe, "enable_vae_slicing"): + pipe.enable_vae_slicing() + logger.info("Enabled VAE slicing for lower memory usage") + + # VAE tiling: decode latents tile-by-tile for large images + # https://huggingface.co/docs/diffusers/optimization/memory#vae-tiling + if getattr(config, "vae_tiling", False): + if hasattr(pipe, "enable_vae_tiling"): + pipe.enable_vae_tiling() + logger.info("Enabled VAE tiling for large image support") + + def _apply_attention_backend(self, pipe: Any, server_args: ServerArgs) -> None: + """Apply attention backend setting from pipeline config or server_args. + + See: https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends + Available backends: flash, _flash_3_hub, sage, xformers, native, etc. + """ + backend = getattr(server_args, "diffusers_attention_backend", None) + + if backend is None: + config = getattr(server_args, "pipeline_config", None) + if config is not None: + backend = getattr(config, "diffusers_attention_backend", None) + + if backend is None: + return + + for component_name in ["transformer", "unet"]: + component = getattr(pipe, component_name, None) + if component is not None and hasattr(component, "set_attention_backend"): + try: + component.set_attention_backend(backend) + logger.info( + "Set attention backend '%s' on %s", backend, component_name + ) + except Exception as e: + logger.warning( + "Failed to set attention backend '%s' on %s: %s", + backend, + component_name, + e, + ) + + def _get_device_map(self, server_args: ServerArgs) -> str | None: + """ + Determine device_map for pipeline loading. + """ + if not torch.cuda.is_available(): + return None + return "cuda" + + def _get_dtype(self, server_args: ServerArgs) -> torch.dtype: + """ + Determine the dtype to use for model loading. + """ + dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + + if hasattr(server_args, "pipeline_config") and server_args.pipeline_config: + dit_precision = getattr(server_args.pipeline_config, "dit_precision", None) + if dit_precision == "fp16": + dtype = torch.float16 + elif dit_precision == "bf16": + dtype = torch.bfloat16 + elif dit_precision == "fp32": + dtype = torch.float32 + + return dtype + + def _detect_pipeline_type(self): + """Detect if this is an image or video pipeline.""" + pipe_class_name = self.diffusers_pipe.__class__.__name__.lower() + video_indicators = ["video", "animat", "cogvideo", "wan", "hunyuan"] + self.is_video_pipeline = any(ind in pipe_class_name for ind in video_indicators) + logger.info( + "Detected pipeline type: %s", + "video" if self.is_video_pipeline else "image", + ) + + def load_modules( + self, + server_args: ServerArgs, + loaded_modules: dict[str, torch.nn.Module] | None = None, + ) -> dict[str, Any]: + """Skip sglang's module loading - diffusers handles it.""" + return {"diffusers_pipeline": self.diffusers_pipe} + + def create_pipeline_stages(self, server_args: ServerArgs): + """Create the execution stage wrapping the diffusers pipeline.""" + self.add_stage( + stage_name="diffusers_execution", + stage=DiffusersExecutionStage(self.diffusers_pipe), + ) + + def initialize_pipeline(self, server_args: ServerArgs): + """Initialize the pipeline.""" + pass + + def post_init(self) -> None: + """Post initialization hook.""" + if self.post_init_called: + return + self.post_init_called = True + self.initialize_pipeline(self.server_args) + self.create_pipeline_stages(self.server_args) + + def add_stage(self, stage_name: str, stage: PipelineStage): + """Add a stage to the pipeline.""" + self._stages.append(stage) + self._stage_name_mapping[stage_name] = stage + setattr(self, stage_name, stage) + + @property + def stages(self) -> list[PipelineStage]: + """List of stages in the pipeline.""" + return self._stages + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + """Execute the pipeline on the given batch.""" + if not self.post_init_called: + self.post_init() + return self.executor.execute(self.stages, batch, server_args) + + @classmethod + def from_pretrained( + cls, + model_path: str, + device: str | None = None, + torch_dtype: torch.dtype | None = None, + pipeline_config: str | PipelineConfig | None = None, + args: argparse.Namespace | None = None, + required_config_modules: list[str] | None = None, + loaded_modules: dict[str, torch.nn.Module] | None = None, + **kwargs, + ) -> "DiffusersPipeline": + """Load a pipeline from a pretrained model using diffusers backend.""" + kwargs["model_path"] = model_path + server_args = ServerArgs.from_kwargs(**kwargs) + + pipe = cls( + model_path, + server_args, + required_config_modules=required_config_modules, + loaded_modules=loaded_modules, + ) + pipe.post_init() + return pipe + + def get_module(self, module_name: str, default_value: Any = None) -> Any: + """Get a module by name.""" + if module_name == "diffusers_pipeline": + return self.diffusers_pipe + return self.modules.get(module_name, default_value) + + +EntryClass = DiffusersPipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py index 27206b531..f5df51d80 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py @@ -42,7 +42,9 @@ def build_pipeline( 3. based on the config, determine the pipeline class """ model_path = server_args.model_path - model_info = get_model_info(model_path) + model_info = get_model_info(model_path, backend=server_args.backend) + if model_info is None: + raise ValueError(f"Unsupported model: {model_path}") pipeline_cls = model_info.pipeline_cls diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index eca3a893f..83ff71c5a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -136,6 +136,9 @@ class Req: ) guidance_scale: float = 1.0 guidance_scale_2: float | None = None + true_cfg_scale: float | None = ( + None # for CFG vs guidance distillation (e.g., QwenImage) + ) guidance_rescale: float = 0.0 eta: float = 0.0 sigmas: list[float] | None = None diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index f76ed7c47..2c5c602dc 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -143,13 +143,98 @@ def _sanitize_for_logging(obj: Any, key_hint: str | None = None) -> Any: return "" +class ExecutionMode(str, Enum): + """ + Enumeration for different pipeline modes. + + Inherits from str to allow string comparison for backward compatibility. + """ + + INFERENCE = "inference" + + @classmethod + def from_string(cls, value: str) -> "ExecutionMode": + """Convert string to ExecutionMode enum.""" + try: + return cls(value.lower()) + except ValueError: + raise ValueError( + f"Invalid mode: {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 [mode.value for mode in cls] + + +class WorkloadType(str, Enum): + """ + Enumeration for different workload types. + + Inherits from str to allow string comparison for backward compatibility. + """ + + I2V = "i2v" # Image to Video + T2V = "t2v" # Text to Video + T2I = "t2i" # Text to Image + I2I = "i2i" # Image to Image + + @classmethod + def from_string(cls, value: str) -> "WorkloadType": + """Convert string to WorkloadType enum.""" + try: + return cls(value.lower()) + except ValueError: + raise ValueError( + f"Invalid workload 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 [workload.value for workload in cls] + + +class Backend(str, Enum): + """ + Enumeration for different model backends. + - AUTO: Automatically select backend (prefer sglang native, fallback to diffusers) + - SGLANG: Use sglang's native optimized implementation + - DIFFUSERS: Use vanilla diffusers pipeline (supports all diffusers models) + """ + + AUTO = "auto" + SGLANG = "sglang" + DIFFUSERS = "diffusers" + + @classmethod + def from_string(cls, value: str) -> "Backend": + """Convert string to Backend enum.""" + try: + return cls(value.lower()) + except ValueError: + raise ValueError( + f"Invalid backend: {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 [backend.value for backend in cls] + + @dataclasses.dataclass class ServerArgs: # Model and path configuration (for convenience) model_path: str + # Model backend (sglang native or diffusers) + backend: Backend = Backend.AUTO + # Attention attention_backend: str = None + diffusers_attention_backend: str = None # for diffusers backend only # Distributed executor backend nccl_port: Optional[int] = None @@ -355,6 +440,13 @@ class ServerArgs: choices=[e.name.lower() for e in AttentionBackendEnum] + ["fa3", "fa4"], help="The attention backend to use. If not specified, the backend is automatically selected based on hardware and installed packages.", ) + parser.add_argument( + "--diffusers-attention-backend", + type=str, + default=None, + help="Attention backend for diffusers pipelines (e.g., flash, _flash_3_hub, sage, xformers). " + "See: https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends", + ) # HuggingFace specific parameters parser.add_argument( @@ -601,6 +693,14 @@ class ServerArgs: default=ServerArgs.log_level, help="The logging level of all loggers.", ) + parser.add_argument( + "--backend", + type=str, + choices=Backend.choices(), + default=ServerArgs.backend.value, + help="The model backend to use. 'auto' prefers sglang native and falls back to diffusers. " + "'sglang' uses native optimized implementation. 'diffusers' uses vanilla diffusers pipeline.", + ) return parser def url(self): @@ -718,6 +818,16 @@ class ServerArgs: @classmethod def from_kwargs(cls, **kwargs: Any) -> "ServerArgs": + # Convert mode string to enum if necessary + if "mode" in kwargs and isinstance(kwargs["mode"], str): + kwargs["mode"] = ExecutionMode.from_string(kwargs["mode"]) + # Convert workload_type string to enum if necessary + if "workload_type" in kwargs and isinstance(kwargs["workload_type"], str): + kwargs["workload_type"] = WorkloadType.from_string(kwargs["workload_type"]) + # Convert backend string to enum if necessary + if "backend" in kwargs and isinstance(kwargs["backend"], str): + kwargs["backend"] = Backend.from_string(kwargs["backend"]) + kwargs["pipeline_config"] = PipelineConfig.from_kwargs(kwargs) return cls(**kwargs)