From 5639145facec2c5ea40ec9ca44286eb014a069be Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 10 Nov 2025 21:20:33 +0800 Subject: [PATCH] diffusion: reduce effort of supporting new model (#12982) --- python/sglang/multimodal_gen/README.md | 4 +- .../configs/pipelines/__init__.py | 4 - .../multimodal_gen/configs/pipelines/base.py | 27 +- .../configs/pipelines/qwen_image.py | 10 +- .../configs/pipelines/registry.py | 168 ------- .../multimodal_gen/configs/sample/base.py | 12 +- .../multimodal_gen/configs/sample/registry.py | 122 ----- python/sglang/multimodal_gen/registry.py | 424 +++++++++++++++++ .../runtime/architectures/basic/flux/flux.py | 12 +- .../basic/hunyuan/hunyuan_pipeline.py | 4 +- .../basic/qwen_image/qwen_image.py | 15 +- .../basic/wan/wan_causal_dmd_pipeline.py | 7 +- .../basic/wan/wan_dmd_pipeline.py | 5 +- .../basic/wan/wan_i2v_dmd_pipeline.py | 2 +- .../architectures/basic/wan/wan_pipeline.py | 7 +- .../architectures/preprocess/__init__.py | 1 - .../preprocess/preprocess_pipeline_base.py | 433 ------------------ .../preprocess/preprocess_pipeline_i2v.py | 247 ---------- .../preprocess_pipeline_ode_trajectory.py | 355 -------------- .../preprocess/preprocess_pipeline_t2v.py | 26 -- .../preprocess/preprocess_pipeline_text.py | 200 -------- .../preprocess/preprocess_stages.py | 134 ------ .../architectures/preprocess/v1_preprocess.py | 147 ------ .../preprocess/v1_preprocessing_new.py | 26 -- .../architectures/preprocess/wan/__init__.py | 1 - .../wan/wan_preprocess_pipelines.py | 118 ----- .../entrypoints/diffusion_generator.py | 4 +- .../runtime/entrypoints/openai/image_api.py | 2 +- .../runtime/entrypoints/openai/video_api.py | 2 +- .../runtime/entrypoints/utils.py | 2 +- .../runtime/layers/attention/layer.py | 1 - .../runtime/layers/visual_embedding.py | 6 +- .../runtime/managers/gpu_worker.py | 4 +- .../runtime/managers/scheduler.py | 8 +- .../runtime/managers/schedulerbase.py | 3 +- .../runtime/models/dits/wanvideo.py | 9 +- .../runtime/pipelines/__init__.py | 41 +- .../pipelines/composed_pipeline_base.py | 2 +- .../pipelines/executors/pipeline_executor.py | 2 +- .../pipelines/executors/sync_executor.py | 2 +- .../runtime/pipelines/pipeline_registry.py | 239 ---------- ...peline_batch_info.py => schedule_batch.py} | 0 .../runtime/pipelines/stages/base.py | 2 +- .../pipelines/stages/causal_denoising.py | 2 +- .../runtime/pipelines/stages/conditioning.py | 2 +- .../runtime/pipelines/stages/decoding.py | 2 +- .../runtime/pipelines/stages/denoising.py | 2 +- .../runtime/pipelines/stages/denoising_dmd.py | 2 +- .../runtime/pipelines/stages/encoding.py | 2 +- .../pipelines/stages/image_encoding.py | 2 +- .../pipelines/stages/input_validation.py | 2 +- .../pipelines/stages/latent_preparation.py | 2 +- .../pipelines/stages/stepvideo_encoding.py | 2 +- .../runtime/pipelines/stages/text_encoding.py | 2 +- .../pipelines/stages/timestep_preparation.py | 2 +- .../runtime/scheduler_client.py | 2 +- .../multimodal_gen/runtime/server_args.py | 44 +- .../runtime/sync_scheduler_client.py | 2 +- .../multimodal_gen/runtime/utils/common.py | 82 +++- .../runtime/utils/logging_utils.py | 2 +- .../preprocess/preprocess_workflow.py | 5 +- .../test/cli/test_generate_common.py | 6 +- .../test/cli/test_generate_t2i_perf.py | 5 +- .../test/cli/test_generate_t2v_perf.py | 19 +- .../test/cli/test_generate_ti2v_perf.py | 7 +- .../sglang/multimodal_gen/test/test_utils.py | 14 +- 66 files changed, 667 insertions(+), 2385 deletions(-) delete mode 100644 python/sglang/multimodal_gen/configs/pipelines/registry.py delete mode 100644 python/sglang/multimodal_gen/configs/sample/registry.py create mode 100644 python/sglang/multimodal_gen/registry.py delete mode 100644 python/sglang/multimodal_gen/runtime/architectures/preprocess/__init__.py delete mode 100644 python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_base.py delete mode 100644 python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_i2v.py delete mode 100644 python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_ode_trajectory.py delete mode 100644 python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_t2v.py delete mode 100644 python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_text.py delete mode 100644 python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_stages.py delete mode 100644 python/sglang/multimodal_gen/runtime/architectures/preprocess/v1_preprocess.py delete mode 100644 python/sglang/multimodal_gen/runtime/architectures/preprocess/v1_preprocessing_new.py delete mode 100644 python/sglang/multimodal_gen/runtime/architectures/preprocess/wan/__init__.py delete mode 100644 python/sglang/multimodal_gen/runtime/architectures/preprocess/wan/wan_preprocess_pipelines.py delete mode 100644 python/sglang/multimodal_gen/runtime/pipelines/pipeline_registry.py rename python/sglang/multimodal_gen/runtime/pipelines/{pipeline_batch_info.py => schedule_batch.py} (100%) diff --git a/python/sglang/multimodal_gen/README.md b/python/sglang/multimodal_gen/README.md index 8b2655628..4f7dc7051 100644 --- a/python/sglang/multimodal_gen/README.md +++ b/python/sglang/multimodal_gen/README.md @@ -20,7 +20,7 @@ SGLang Diffusion has the following features: uv pip install 'sglang[diffusion]' --prerelease=allow ``` -For more installation methods (e.g. pypi, uv, docker), check the [docs](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen/docs/install.md). +For more installation methods (e.g. pypi, uv, docker), check [install.md](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen/docs/install.md). ## Inference @@ -61,7 +61,7 @@ sglang generate --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ --save-output ``` -For more usage examples (e.g. OpenAI compatible API, server mode), check the [docs](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen/docs/cli.md). +For more usage examples (e.g. OpenAI compatible API, server mode), check [cli.md](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen/docs/cli.md). ## Contributing diff --git a/python/sglang/multimodal_gen/configs/pipelines/__init__.py b/python/sglang/multimodal_gen/configs/pipelines/__init__.py index 5db869f31..6611a9382 100644 --- a/python/sglang/multimodal_gen/configs/pipelines/__init__.py +++ b/python/sglang/multimodal_gen/configs/pipelines/__init__.py @@ -9,9 +9,6 @@ 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, @@ -33,5 +30,4 @@ __all__ = [ "WanI2V720PConfig", "StepVideoT2VConfig", "SelfForcingWanT2V480PConfig", - "get_pipeline_config_cls_from_name", ] diff --git a/python/sglang/multimodal_gen/configs/pipelines/base.py b/python/sglang/multimodal_gen/configs/pipelines/base.py index 9451f59d3..09a01f526 100644 --- a/python/sglang/multimodal_gen/configs/pipelines/base.py +++ b/python/sglang/multimodal_gen/configs/pipelines/base.py @@ -5,7 +5,7 @@ import json from collections.abc import Callable from dataclasses import asdict, dataclass, field, fields from enum import Enum -from typing import Any, cast +from typing import Any import torch from diffusers.image_processor import VaeImageProcessor @@ -312,19 +312,6 @@ class PipelineConfig: 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 = "" @@ -334,9 +321,7 @@ class PipelineConfig: 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, - ) + from sglang.multimodal_gen.registry import get_model_info prefix_with_dot = ( f"{config_cli_prefix}." if (config_cli_prefix.strip() != "") else "" @@ -352,17 +337,17 @@ class PipelineConfig: 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) + model_info = get_model_info(model_path) # 2. Instantiate PipelineConfig - if pipeline_config_cls is None: + if model_info is None: logger.warning( - "Couldn't find pipeline config for %s. Using the default pipeline config.", + "Couldn't find model info for %s. Using the default pipeline config.", model_path, ) pipeline_config = cls() else: - pipeline_config = pipeline_config_cls() + pipeline_config = model_info.pipeline_config_cls() # 3. Load PipelineConfig from a json file or a PipelineConfig object if provided if isinstance(pipeline_config_or_path, str): diff --git a/python/sglang/multimodal_gen/configs/pipelines/qwen_image.py b/python/sglang/multimodal_gen/configs/pipelines/qwen_image.py index e4b702c84..1b16c6fd4 100644 --- a/python/sglang/multimodal_gen/configs/pipelines/qwen_image.py +++ b/python/sglang/multimodal_gen/configs/pipelines/qwen_image.py @@ -117,15 +117,7 @@ class QwenImagePipelineConfig(PipelineConfig): 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 + return _pack_latents(latents, batch_size, num_channels_latents, height, width) @staticmethod def get_freqs_cis(img_shapes, txt_seq_lens, rotary_emb, device, dtype): diff --git a/python/sglang/multimodal_gen/configs/pipelines/registry.py b/python/sglang/multimodal_gen/configs/pipelines/registry.py deleted file mode 100644 index b9c223399..000000000 --- a/python/sglang/multimodal_gen/configs/pipelines/registry.py +++ /dev/null @@ -1,168 +0,0 @@ -# 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 diff --git a/python/sglang/multimodal_gen/configs/sample/base.py b/python/sglang/multimodal_gen/configs/sample/base.py index 27da2f9f5..ca1ddd1aa 100644 --- a/python/sglang/multimodal_gen/configs/sample/base.py +++ b/python/sglang/multimodal_gen/configs/sample/base.py @@ -205,14 +205,12 @@ class SamplingParams: @classmethod def from_pretrained(cls, model_path: str, **kwargs) -> "SamplingParams": - from sglang.multimodal_gen.configs.sample.registry import ( - get_sampling_param_cls_for_name, - ) + from sglang.multimodal_gen.registry import get_model_info - 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) + model_info = get_model_info(model_path) + logger.debug(f"Found model info: {model_info}") + if model_info is not None: + sampling_params: SamplingParams = model_info.sampling_param_cls(**kwargs) else: logger.warning( "Couldn't find an optimal sampling param for %s. Using the default sampling param.", diff --git a/python/sglang/multimodal_gen/configs/sample/registry.py b/python/sglang/multimodal_gen/configs/sample/registry.py deleted file mode 100644 index 297901fc2..000000000 --- a/python/sglang/multimodal_gen/configs/sample/registry.py +++ /dev/null @@ -1,122 +0,0 @@ -# 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 diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py new file mode 100644 index 000000000..6897f8906 --- /dev/null +++ b/python/sglang/multimodal_gen/registry.py @@ -0,0 +1,424 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Central registry for multimodal models. + +This module provides a centralized registry for multimodal models, including pipelines +and sampling parameters. It allows for easy registration and retrieval of model +information based on model paths or other identifiers. +""" + +import dataclasses +import importlib +import os +import pkgutil +from typing import Any, Callable, Dict, List, Optional, Tuple, Type + +from sglang.multimodal_gen.configs.pipelines import ( + FastHunyuanConfig, + FluxPipelineConfig, + HunyuanConfig, + StepVideoT2VConfig, + WanI2V480PConfig, + WanI2V720PConfig, + WanT2V480PConfig, + WanT2V720PConfig, +) +from sglang.multimodal_gen.configs.pipelines.base import PipelineConfig +from sglang.multimodal_gen.configs.pipelines.qwen_image import ( + QwenImageEditPipelineConfig, + QwenImagePipelineConfig, +) +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, +) +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 +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, +) +from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( + maybe_download_model_index, + verify_model_config_and_directory, +) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +# --- Part 1: Pipeline Discovery --- + +_PIPELINE_REGISTRY: Dict[str, Type[ComposedPipelineBase]] = {} + + +def _discover_and_register_pipelines(): + """ + Automatically discover and register all ComposedPipelineBase subclasses. + This function scans the 'sglang.multimodal_gen.runtime.architectures' package, + finds modules with an 'EntryClass' attribute, and maps the class's 'pipeline_name' + to the class itself in a global registry. + """ + if _PIPELINE_REGISTRY: # E-run only once + return + + package_name = "sglang.multimodal_gen.runtime.architectures" + package = importlib.import_module(package_name) + + for _, pipeline_type_str, ispkg in pkgutil.iter_modules(package.__path__): + if not ispkg: + continue + pipeline_type_package_name = f"{package_name}.{pipeline_type_str}" + pipeline_type_package = importlib.import_module(pipeline_type_package_name) + for _, arch, ispkg_arch in pkgutil.iter_modules(pipeline_type_package.__path__): + if not ispkg_arch: + continue + arch_package_name = f"{pipeline_type_package_name}.{arch}" + arch_package = importlib.import_module(arch_package_name) + for _, module_name, ispkg_module in pkgutil.walk_packages( + arch_package.__path__, arch_package.__name__ + "." + ): + if not ispkg_module: + pipeline_module = importlib.import_module(module_name) + if hasattr(pipeline_module, "EntryClass"): + entry_cls = pipeline_module.EntryClass + if not isinstance(entry_cls, list): + entry_cls_list = [entry_cls] + else: + entry_cls_list = entry_cls + + for cls in entry_cls_list: + if hasattr(cls, "pipeline_name"): + if cls.pipeline_name in _PIPELINE_REGISTRY: + logger.warning( + f"Duplicate pipeline name '{cls.pipeline_name}' found. Overwriting." + ) + _PIPELINE_REGISTRY[cls.pipeline_name] = cls + # else: + # logger.warning( + # f"Pipeline class {cls.__name__} does not have a 'pipeline_name' attribute." + # ) + + +# --- Part 2: Config Registration --- +@dataclasses.dataclass +class ConfigInfo: + """Encapsulates all configuration information required to register a + diffusers model within this framework.""" + + sampling_param_cls: Any + pipeline_config_cls: Type[PipelineConfig] + + +# The central registry mapping a model name to its configuration information +_CONFIG_REGISTRY: Dict[str, ConfigInfo] = {} + +# Mappings from Hugging Face model paths to our internal model names +_MODEL_PATH_TO_NAME: Dict[str, str] = {} + +# Detectors to identify model families from paths or class names +_MODEL_NAME_DETECTORS: List[Tuple[str, Callable[[str], bool]]] = [] + + +def register_configs( + model_name: str, + sampling_param_cls: Any, + pipeline_config_cls: Type[PipelineConfig], + model_path_to_name_mappings: Optional[Dict[str, str]] = None, + model_name_detectors: Optional[List[Tuple[str, Callable[[str], bool]]]] = None, +): + """ + Registers configuration classes for a new model family. + """ + if model_name in _CONFIG_REGISTRY: + logger.warning( + f"Config for model '{model_name}' is already registered and will be overwritten." + ) + + _CONFIG_REGISTRY[model_name] = ConfigInfo( + sampling_param_cls=sampling_param_cls, + pipeline_config_cls=pipeline_config_cls, + ) + if model_path_to_name_mappings: + for path, name in model_path_to_name_mappings.items(): + if path in _MODEL_PATH_TO_NAME: + logger.warning( + f"Model path '{path}' is already mapped to '{_MODEL_PATH_TO_NAME[path]}' and will be overwritten by '{name}'." + ) + _MODEL_PATH_TO_NAME[path] = name + + if model_name_detectors: + _MODEL_NAME_DETECTORS.extend(model_name_detectors) + + +def _get_config_info(model_path: str) -> Optional[ConfigInfo]: + """ + Gets the ConfigInfo for a given model path using mappings and detectors. + """ + # 1. Exact match + if model_path in _MODEL_PATH_TO_NAME: + model_name = _MODEL_PATH_TO_NAME[model_path] + return _CONFIG_REGISTRY.get(model_name) + + # 2. Partial match + for registered_id, model_name in _MODEL_PATH_TO_NAME.items(): + if registered_id in model_path: + return _CONFIG_REGISTRY.get(model_name) + + # 3. Use detectors + if os.path.exists(model_path): + config = verify_model_config_and_directory(model_path) + else: + config = maybe_download_model_index(model_path) + + pipeline_name = config.get("_class_name", "").lower() + + for model_name, detector in _MODEL_NAME_DETECTORS: + if detector(model_path.lower()) or detector(pipeline_name): + return _CONFIG_REGISTRY.get(model_name) + + return None + + +# --- Part 3: Main Resolver --- + + +@dataclasses.dataclass +class ModelInfo: + """ + Encapsulates all configuration information required to register a + diffusers model within this framework. + """ + + pipeline_cls: Type[ComposedPipelineBase] + sampling_param_cls: Any + pipeline_config_cls: Type[PipelineConfig] + + +def get_model_info(model_path: str) -> Optional[ModelInfo]: + """ + Resolves all necessary classes (pipeline, sampling, config) for a given model path. + + This function serves as the main entry point for model resolution. It performs two main tasks: + 1. Dynamically resolves the pipeline class by reading 'model_index.json' and matching + '_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. + """ + # 1. Discover all available pipeline classes and cache them + _discover_and_register_pipelines() + + # 2. Get pipeline class from model's model_index.json + try: + if os.path.exists(model_path): + config = verify_model_config_and_directory(model_path) + else: + config = maybe_download_model_index(model_path) + except Exception as e: + logger.error(f"Could not read model config for '{model_path}': {e}") + 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}'") + 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 + + # 3. Get configuration classes (sampling, pipeline config) + config_info = _get_config_info(model_path) + if not config_info: + logger.warning( + f"No specific configuration registered for '{model_path}'. " + f"Falling back to default SamplingParams and PipelineConfig." + ) + # Fallback to defaults if no specific config is found + from sglang.multimodal_gen.configs.sample.base import SamplingParams + + config_info = ConfigInfo( + sampling_param_cls=SamplingParams, pipeline_config_cls=PipelineConfig + ) + + # 4. Combine and return the complete model info + return ModelInfo( + pipeline_cls=pipeline_cls, + sampling_param_cls=config_info.sampling_param_cls, + pipeline_config_cls=config_info.pipeline_config_cls, + ) + + +# Registration of model configs +def _register_configs(): + # Hunyuan + register_configs( + model_name="hunyuan", + sampling_param_cls=HunyuanSamplingParams, + pipeline_config_cls=HunyuanConfig, + model_path_to_name_mappings={ + "hunyuanvideo-community/HunyuanVideo": "hunyuan", + }, + model_name_detectors=[("hunyuan", lambda id: "hunyuan" in id.lower())], + ) + register_configs( + model_name="fasthunyuan", + sampling_param_cls=FastHunyuanSamplingParam, + pipeline_config_cls=FastHunyuanConfig, + model_path_to_name_mappings={ + "FastVideo/FastHunyuan-diffusers": "fasthunyuan", + }, + ) + + # StepVideo + register_configs( + model_name="stepvideo", + sampling_param_cls=StepVideoT2VSamplingParams, + pipeline_config_cls=StepVideoT2VConfig, + model_path_to_name_mappings={ + "FastVideo/stepvideo-t2v-diffusers": "stepvideo", + }, + model_name_detectors=[("stepvideo", lambda id: "stepvideo" in id.lower())], + ) + + # Wan + register_configs( + model_name="wan-t2v-1.3b", + sampling_param_cls=WanT2V_1_3B_SamplingParams, + pipeline_config_cls=WanT2V480PConfig, + model_path_to_name_mappings={ + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": "wan-t2v-1.3b", + }, + model_name_detectors=[("wanpipeline", lambda id: "wanpipeline" in id.lower())], + ) + register_configs( + model_name="wan-t2v-14b", + sampling_param_cls=WanT2V_14B_SamplingParams, + pipeline_config_cls=WanT2V720PConfig, + model_path_to_name_mappings={ + "Wan-AI/Wan2.1-T2V-14B-Diffusers": "wan-t2v-14b", + }, + ) + register_configs( + model_name="wan-i2v-14b-480p", + sampling_param_cls=WanI2V_14B_480P_SamplingParam, + pipeline_config_cls=WanI2V480PConfig, + model_path_to_name_mappings={ + "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": "wan-i2v-14b-480p", + }, + model_name_detectors=[ + ("wanimagetovideo", lambda id: "wanimagetovideo" in id.lower()) + ], + ) + register_configs( + model_name="wan-i2v-14b-720p", + sampling_param_cls=WanI2V_14B_720P_SamplingParam, + pipeline_config_cls=WanI2V720PConfig, + model_path_to_name_mappings={ + "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": "wan-i2v-14b-720p", + }, + ) + register_configs( + model_name="wan-fun-1.3b-inp", + sampling_param_cls=Wan2_1_Fun_1_3B_InP_SamplingParams, + pipeline_config_cls=WanI2V480PConfig, + model_path_to_name_mappings={ + "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": "wan-fun-1.3b-inp", + }, + ) + register_configs( + model_name="wan-ti2v-5b", + sampling_param_cls=Wan2_2_TI2V_5B_SamplingParam, + pipeline_config_cls=Wan2_2_TI2V_5B_Config, + model_path_to_name_mappings={ + "Wan-AI/Wan2.2-TI2V-5B-Diffusers": "wan-ti2v-5b", + }, + ) + + register_configs( + model_name="fastwan-ti2v-5b", + sampling_param_cls=Wan2_2_TI2V_5B_SamplingParam, + pipeline_config_cls=FastWan2_2_TI2V_5B_Config, + model_path_to_name_mappings={ + "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers": "fastwan-ti2v-5b", + "FastVideo/FastWan2.2-TI2V-5B-Diffusers": "fastwan-ti2v-5b", + }, + ) + + register_configs( + model_name="wan-t2v-a14b", + sampling_param_cls=Wan2_2_T2V_A14B_SamplingParam, + pipeline_config_cls=Wan2_2_T2V_A14B_Config, + model_path_to_name_mappings={ + "Wan-AI/Wan2.2-T2V-A14B-Diffusers": "wan-t2v-a14b", + }, + ) + register_configs( + model_name="wan-i2v-a14b", + sampling_param_cls=Wan2_2_I2V_A14B_SamplingParam, + pipeline_config_cls=Wan2_2_I2V_A14B_Config, + model_path_to_name_mappings={ + "Wan-AI/Wan2.2-I2V-A14B-Diffusers": "wan-i2v-a14b", + }, + ) + register_configs( + model_name="fast-wan-t2v-1.3b", + sampling_param_cls=FastWanT2V480PConfig, + pipeline_config_cls=FastWan2_1_T2V_480P_Config, + model_path_to_name_mappings={ + "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": "fast-wan-t2v-1.3b", + }, + ) + + # FLUX + register_configs( + model_name="flux", + sampling_param_cls=FluxSamplingParams, + pipeline_config_cls=FluxPipelineConfig, + model_path_to_name_mappings={ + "black-forest-labs/FLUX.1-dev": "flux", + }, + model_name_detectors=[("flux", lambda id: "flux" in id.lower())], + ) + + # Qwen-Image + register_configs( + model_name="qwen-image", + sampling_param_cls=QwenImageSamplingParams, + pipeline_config_cls=QwenImagePipelineConfig, + model_path_to_name_mappings={ + "Qwen/Qwen-Image": "qwen-image", + }, + ) + register_configs( + model_name="qwen-image-edit", + sampling_param_cls=QwenImageSamplingParams, + pipeline_config_cls=QwenImageEditPipelineConfig, + model_path_to_name_mappings={ + "Qwen/Qwen-Image-Edit": "qwen-image-edit", + }, + ) + + +_register_configs() diff --git a/python/sglang/multimodal_gen/runtime/architectures/basic/flux/flux.py b/python/sglang/multimodal_gen/runtime/architectures/basic/flux/flux.py index d88e554db..b56947b42 100644 --- a/python/sglang/multimodal_gen/runtime/architectures/basic/flux/flux.py +++ b/python/sglang/multimodal_gen/runtime/architectures/basic/flux/flux.py @@ -1,14 +1,10 @@ # 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.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages import ( ConditioningStage, DecodingStage, diff --git a/python/sglang/multimodal_gen/runtime/architectures/basic/hunyuan/hunyuan_pipeline.py b/python/sglang/multimodal_gen/runtime/architectures/basic/hunyuan/hunyuan_pipeline.py index ffc2c6eec..8edf842c0 100644 --- a/python/sglang/multimodal_gen/runtime/architectures/basic/hunyuan/hunyuan_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/architectures/basic/hunyuan/hunyuan_pipeline.py @@ -9,7 +9,9 @@ using the modular pipeline architecture. """ -from sglang.multimodal_gen.runtime.pipelines import ComposedPipelineBase +from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import ( + ComposedPipelineBase, +) from sglang.multimodal_gen.runtime.pipelines.stages import ( ConditioningStage, DecodingStage, diff --git a/python/sglang/multimodal_gen/runtime/architectures/basic/qwen_image/qwen_image.py b/python/sglang/multimodal_gen/runtime/architectures/basic/qwen_image/qwen_image.py index 649a7f74d..0682c467f 100644 --- a/python/sglang/multimodal_gen/runtime/architectures/basic/qwen_image/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/architectures/basic/qwen_image/qwen_image.py @@ -1,17 +1,13 @@ # 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.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages import ( - ConditioningStage, DecodingStage, DenoisingStage, ImageEncodingStage, @@ -21,6 +17,9 @@ from sglang.multimodal_gen.runtime.pipelines.stages import ( TextEncodingStage, TimestepPreparationStage, ) +from sglang.multimodal_gen.runtime.pipelines.stages.conditioning import ( + ConditioningStage, +) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger diff --git a/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_causal_dmd_pipeline.py b/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_causal_dmd_pipeline.py index 6e1f59be2..31b5bf04f 100644 --- a/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_causal_dmd_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_causal_dmd_pipeline.py @@ -7,7 +7,10 @@ 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 +from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.pipelines.lora_pipeline import LoRAPipeline # isort: off from sglang.multimodal_gen.runtime.pipelines.stages import ( @@ -27,7 +30,7 @@ logger = init_logger(__name__) class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase): - pipeline_name = "WanPipeline" + pipeline_name = "WanCausalDMDPipeline" _required_config_modules = [ "text_encoder", diff --git a/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_dmd_pipeline.py b/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_dmd_pipeline.py index 2b13408e3..984812f1f 100644 --- a/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_dmd_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_dmd_pipeline.py @@ -11,7 +11,10 @@ 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.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 diff --git a/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_i2v_dmd_pipeline.py b/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_i2v_dmd_pipeline.py index b0e256457..ed3fdd4ea 100644 --- a/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_i2v_dmd_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_i2v_dmd_pipeline.py @@ -37,7 +37,7 @@ logger = init_logger(__name__) class WanImageToVideoDmdPipeline(LoRAPipeline, ComposedPipelineBase): - pipeline_name = "WanCausalDMDPipeline" + pipeline_name = "WanImageToVideoDmdPipeline" _required_config_modules = [ "text_encoder", diff --git a/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_pipeline.py b/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_pipeline.py index be49674d6..9d7c592bd 100644 --- a/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/architectures/basic/wan/wan_pipeline.py @@ -11,7 +11,10 @@ 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.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.pipelines.lora_pipeline import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines.stages import ( ConditioningStage, DecodingStage, @@ -32,7 +35,7 @@ class WanPipeline(LoRAPipeline, ComposedPipelineBase): Wan video diffusion pipeline with LoRA support. """ - pipeline_name = "WanImageToVideoPipeline" + pipeline_name = "WanPipeline" _required_config_modules = [ "text_encoder", diff --git a/python/sglang/multimodal_gen/runtime/architectures/preprocess/__init__.py b/python/sglang/multimodal_gen/runtime/architectures/preprocess/__init__.py deleted file mode 100644 index af2eb7d10..000000000 --- a/python/sglang/multimodal_gen/runtime/architectures/preprocess/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo diff --git a/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_base.py b/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_base.py deleted file mode 100644 index e82dd5716..000000000 --- a/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_base.py +++ /dev/null @@ -1,433 +0,0 @@ -# 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 diff --git a/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_i2v.py b/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_i2v.py deleted file mode 100644 index 2c6e8dbbc..000000000 --- a/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_i2v.py +++ /dev/null @@ -1,247 +0,0 @@ -# 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 diff --git a/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_ode_trajectory.py b/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_ode_trajectory.py deleted file mode 100644 index 950b38c36..000000000 --- a/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_ode_trajectory.py +++ /dev/null @@ -1,355 +0,0 @@ -# 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 diff --git a/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_t2v.py b/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_t2v.py deleted file mode 100644 index d47ab9aec..000000000 --- a/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_t2v.py +++ /dev/null @@ -1,26 +0,0 @@ -# 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 diff --git a/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_text.py b/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_text.py deleted file mode 100644 index 3906f09a5..000000000 --- a/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_pipeline_text.py +++ /dev/null @@ -1,200 +0,0 @@ -# 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 diff --git a/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_stages.py b/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_stages.py deleted file mode 100644 index 126ab05d6..000000000 --- a/python/sglang/multimodal_gen/runtime/architectures/preprocess/preprocess_stages.py +++ /dev/null @@ -1,134 +0,0 @@ -# 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) diff --git a/python/sglang/multimodal_gen/runtime/architectures/preprocess/v1_preprocess.py b/python/sglang/multimodal_gen/runtime/architectures/preprocess/v1_preprocess.py deleted file mode 100644 index 8a160069a..000000000 --- a/python/sglang/multimodal_gen/runtime/architectures/preprocess/v1_preprocess.py +++ /dev/null @@ -1,147 +0,0 @@ -# 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) diff --git a/python/sglang/multimodal_gen/runtime/architectures/preprocess/v1_preprocessing_new.py b/python/sglang/multimodal_gen/runtime/architectures/preprocess/v1_preprocessing_new.py deleted file mode 100644 index 59f03618b..000000000 --- a/python/sglang/multimodal_gen/runtime/architectures/preprocess/v1_preprocessing_new.py +++ /dev/null @@ -1,26 +0,0 @@ -# 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) diff --git a/python/sglang/multimodal_gen/runtime/architectures/preprocess/wan/__init__.py b/python/sglang/multimodal_gen/runtime/architectures/preprocess/wan/__init__.py deleted file mode 100644 index af2eb7d10..000000000 --- a/python/sglang/multimodal_gen/runtime/architectures/preprocess/wan/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo diff --git a/python/sglang/multimodal_gen/runtime/architectures/preprocess/wan/wan_preprocess_pipelines.py b/python/sglang/multimodal_gen/runtime/architectures/preprocess/wan/wan_preprocess_pipelines.py deleted file mode 100644 index 47ec436ff..000000000 --- a/python/sglang/multimodal_gen/runtime/architectures/preprocess/wan/wan_preprocess_pipelines.py +++ /dev/null @@ -1,118 +0,0 @@ -# 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] diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index a30a99a5e..c59335465 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -21,6 +21,9 @@ import torch import torchvision from einops import rearrange +from sglang.multimodal_gen.runtime.pipelines import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import OutputBatch + # Suppress verbose logging from imageio, which is triggered when saving images. logging.getLogger("imageio").setLevel(logging.WARNING) logging.getLogger("imageio_ffmpeg").setLevel(logging.WARNING) @@ -32,7 +35,6 @@ from sglang.multimodal_gen.configs.sample.base import DataType, SamplingParams from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request from sglang.multimodal_gen.runtime.launch_server import launch_server from sglang.multimodal_gen.runtime.managers.schedulerbase import SchedulerBase -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import OutputBatch, Req from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs from sglang.multimodal_gen.runtime.sync_scheduler_client import sync_scheduler_client from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger 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 be77cd555..07f67a5a1 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py @@ -24,7 +24,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( post_process_sample, ) from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.scheduler_client import scheduler_client from sglang.multimodal_gen.runtime.server_args import get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger 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 0d5b8353a..ad6b6f000 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py @@ -34,7 +34,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( post_process_sample, ) from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.server_args import get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index 123e3efec..835f860f4 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -16,7 +16,7 @@ logging.getLogger("imageio").setLevel(logging.WARNING) logging.getLogger("imageio_ffmpeg").setLevel(logging.WARNING) from sglang.multimodal_gen.configs.sample.base import DataType, SamplingParams -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req 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 shallow_asdict diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 482ea4efc..b7faea789 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -364,7 +364,6 @@ class USPAttention(nn.Module): ), "USPAttention does not support replicated_qkv." forward_context: ForwardContext = get_forward_context() ctx_attn_metadata = forward_context.attn_metadata - if get_sequence_parallel_world_size() == 1: # No sequence parallelism, just run local attention. out = self.attn_impl.forward(q, k, v, ctx_attn_metadata) diff --git a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py b/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py index 715a1b874..d556ab584 100644 --- a/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py +++ b/python/sglang/multimodal_gen/runtime/layers/visual_embedding.py @@ -101,7 +101,11 @@ class TimestepEmbedder(nn.Module): t, self.frequency_embedding_size, self.max_period, dtype=self.freq_dtype ).to(self.mlp.fc_in.weight.dtype) if timestep_seq_len is not None: - t_freq = t_freq.unflatten(0, (1, timestep_seq_len)) + assert ( + t_freq.shape[0] % timestep_seq_len == 0 + ), "timestep length is not divisible by timestep_seq_len" + batch_size = t_freq.shape[0] // timestep_seq_len + t_freq = t_freq.unflatten(0, (batch_size, timestep_seq_len)) # t_freq = t_freq.to(self.mlp.fc_in.weight.dtype) t_emb = self.mlp(t_freq) return t_emb diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index c6606fa8a..3338482ff 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -16,8 +16,8 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_cfg_group, get_tp_group, ) -from sglang.multimodal_gen.runtime.pipelines import build_pipeline -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import OutputBatch, Req +from sglang.multimodal_gen.runtime.pipelines import Req, build_pipeline +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs from sglang.multimodal_gen.runtime.utils.common import set_cuda_arch from sglang.multimodal_gen.runtime.utils.logging_utils import ( diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index d2e07e9b1..dffd56239 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -6,7 +6,7 @@ from typing import Any import zmq from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import OutputBatch +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.server_args import ( PortArgs, ServerArgs, @@ -42,9 +42,11 @@ class Scheduler: # Inter-process Communication self.context = zmq.Context(io_threads=2) endpoint = server_args.scheduler_endpoint() - logger.info(f"Scheduler listening at endpoint: {endpoint}") if gpu_id == 0: - self.receiver = get_zmq_socket(self.context, zmq.REP, endpoint, True) + self.receiver, actual_endpoint = get_zmq_socket( + self.context, zmq.REP, endpoint, True + ) + logger.info(f"Scheduler bind at endpoint: {actual_endpoint}") else: self.receiver = None diff --git a/python/sglang/multimodal_gen/runtime/managers/schedulerbase.py b/python/sglang/multimodal_gen/runtime/managers/schedulerbase.py index 4bf392250..208ff5fbb 100644 --- a/python/sglang/multimodal_gen/runtime/managers/schedulerbase.py +++ b/python/sglang/multimodal_gen/runtime/managers/schedulerbase.py @@ -6,7 +6,8 @@ from typing import TypeVar import zmq -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import OutputBatch, Req +from sglang.multimodal_gen.runtime.pipelines import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.utils import init_logger diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 2f3caa12f..54f996499 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -349,7 +349,6 @@ class WanTransformerBlock(nn.Module): hidden_states = hidden_states.squeeze(1) bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype - if temb.dim() == 4: # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( @@ -368,12 +367,12 @@ class WanTransformerBlock(nn.Module): shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( e.chunk(6, dim=1) ) + assert shift_msa.dtype == torch.float32 # 1. Self-attention - norm_hidden_states = ( - self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa - ).to(orig_dtype) + norm1 = self.norm1(hidden_states.float()) + norm_hidden_states = (norm1 * (1 + scale_msa) + shift_msa).to(orig_dtype) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) @@ -720,6 +719,7 @@ class WanTransformer3DModel(CachableDiT): encoder_hidden_states_image = None batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size post_patch_num_frames = num_frames // p_t post_patch_height = height // p_h @@ -744,7 +744,6 @@ class WanTransformer3DModel(CachableDiT): hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) - # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v) if timestep.dim() == 2: # ti2v diff --git a/python/sglang/multimodal_gen/runtime/pipelines/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines/__init__.py index 8139975b8..fdb1acbfa 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/__init__.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/__init__.py @@ -9,15 +9,12 @@ This package contains diffusion pipelines for generating videos and images. from typing import cast +from sglang.multimodal_gen.registry import get_model_info 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.pipeline_batch_info import Req -from sglang.multimodal_gen.runtime.pipelines.pipeline_registry import ( - PipelineType, - get_pipeline_registry, -) +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( maybe_download_model, @@ -36,7 +33,6 @@ class PipelineWithLoRA(LoRAPipeline, ComposedPipelineBase): def build_pipeline( server_args: ServerArgs, - pipeline_type: PipelineType | str = PipelineType.BASIC, ) -> PipelineWithLoRA: """ Only works with valid hf diffusers configs. (model_index.json) @@ -45,37 +41,12 @@ def build_pipeline( 2. verify the model config and directory 3. based on the config, determine the pipeline class """ - # Get pipeline type model_path = server_args.model_path - model_path = maybe_download_model(model_path) - # server_args.downloaded_model_path = model_path - logger.info("Model path: %s", model_path) + model_info = get_model_info(model_path) + if model_info is None: + raise ValueError(f"Unsupported model: {model_path}") - config = verify_model_config_and_directory(model_path) - pipeline_name = config.get("_class_name") - if pipeline_name is None: - raise ValueError( - "Model config does not contain a _class_name attribute. " - "Only diffusers format is supported." - ) - - # Get the appropriate pipeline registry based on pipeline_type - logger.info( - "Building pipeline of type: %s", - ( - pipeline_type.value - if isinstance(pipeline_type, PipelineType) - else pipeline_type - ), - ) - pipeline_registry = get_pipeline_registry(pipeline_type) - - if isinstance(pipeline_type, str): - pipeline_type = PipelineType.from_string(pipeline_type) - - pipeline_cls = pipeline_registry.resolve_pipeline_cls( - pipeline_name, pipeline_type, server_args.workload_type - ) + pipeline_cls = model_info.pipeline_cls # instantiate the pipelines pipeline = pipeline_cls(model_path, server_args) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines/composed_pipeline_base.py index d5fcf357e..d3f2c8fae 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/composed_pipeline_base.py @@ -22,7 +22,7 @@ from sglang.multimodal_gen.runtime.loader.component_loader import ( from sglang.multimodal_gen.runtime.pipelines.executors.pipeline_executor import ( PipelineExecutor, ) -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages import PipelineStage from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( diff --git a/python/sglang/multimodal_gen/runtime/pipelines/executors/pipeline_executor.py b/python/sglang/multimodal_gen/runtime/pipelines/executors/pipeline_executor.py index 08dc0ceb5..d826f7cd9 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/executors/pipeline_executor.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/executors/pipeline_executor.py @@ -8,7 +8,7 @@ import time from abc import ABC, abstractmethod from typing import List -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages import PipelineStage from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger diff --git a/python/sglang/multimodal_gen/runtime/pipelines/executors/sync_executor.py b/python/sglang/multimodal_gen/runtime/pipelines/executors/sync_executor.py index 88528c51f..6368593d2 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/executors/sync_executor.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/executors/sync_executor.py @@ -11,7 +11,7 @@ from sglang.multimodal_gen.runtime.pipelines.executors.pipeline_executor import Timer, logger, ) -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages import PipelineStage from sglang.multimodal_gen.runtime.server_args import ServerArgs diff --git a/python/sglang/multimodal_gen/runtime/pipelines/pipeline_registry.py b/python/sglang/multimodal_gen/runtime/pipelines/pipeline_registry.py deleted file mode 100644 index a1605f5ca..000000000 --- a/python/sglang/multimodal_gen/runtime/pipelines/pipeline_registry.py +++ /dev/null @@ -1,239 +0,0 @@ -# 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.6.4.post1/vllm/model_executor/models/registry.py -# and https://github.com/sgl-project/sglang/blob/v0.4.3/python/sglang/srt/models/registry.py -import dataclasses -import importlib -import pkgutil -from collections.abc import Set -from dataclasses import dataclass -from enum import Enum -from functools import lru_cache - -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 WorkloadType -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger - -logger = init_logger(__name__) - -_PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = { - WorkloadType.I2V: "PreprocessPipelineI2V", - WorkloadType.T2V: "PreprocessPipelineT2V", -} - - -class PipelineType(str, Enum): - """ - Enumeration for different pipeline types. - - Inherits from str to allow string comparison for backward compatibility. - """ - - BASIC = "basic" - PREPROCESS = "preprocess" - - @classmethod - def from_string(cls, value: str) -> "PipelineType": - """Convert string to PipelineType enum.""" - try: - return cls(value.lower()) - except ValueError: - raise ValueError( - f"Invalid pipeline type: {value}. Must be one of: {', '.join([t.value for t in cls])}" - ) from None - - @classmethod - def choices(cls) -> list[str]: - """Get all available choices as strings.""" - return [pipeline_type.value for pipeline_type in cls] - - -@dataclass -class _PipelineRegistry: - # Keyed by pipeline_type -> architecture -> pipeline_name - # pipelines[pipeline_type][architecture][pipeline_name] = pipeline_cls - pipelines: dict[str, dict[str, type[ComposedPipelineBase] | None]] = ( - dataclasses.field(default_factory=dict) - ) - - def get_supported_archs( - self, pipeline_name_in_config: str, pipeline_type: PipelineType - ) -> Set[str]: - """Get supported architectures, optionally filtered by pipeline type and workload type.""" - return set(self.pipelines[pipeline_type.value].keys()) - - def _load_preprocess_pipeline_cls( - self, workload_type: WorkloadType - ) -> type[ComposedPipelineBase] | None: - pipeline_name = _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME[workload_type] - - return self.pipelines[PipelineType.PREPROCESS.value][pipeline_name] - - def _try_load_pipeline_cls( - self, - pipeline_name_in_config: str, - pipeline_type: PipelineType, - workload_type: WorkloadType, - ) -> type[ComposedPipelineBase] | type[LoRAPipeline] | None: - """Try to load a pipeline class for the given architecture, pipeline type, and workload type.""" - - if pipeline_type.value not in self.pipelines: - return None - - try: - if pipeline_type == PipelineType.PREPROCESS: - return self._load_preprocess_pipeline_cls(workload_type) - elif pipeline_type == PipelineType.BASIC: - return self.pipelines[pipeline_type.value][pipeline_name_in_config] - else: - raise ValueError(f"Invalid pipeline type: {pipeline_type.value}") - except KeyError as e: - logger.error( - f"Please check if the ComposedPipeline class has been defined associated with {pipeline_type.value}.{pipeline_name_in_config}" - ) - raise e - return None - - def resolve_pipeline_cls( - self, - pipeline_name_in_config: str, - pipeline_type: PipelineType, - workload_type: WorkloadType, - ) -> type[ComposedPipelineBase] | type[LoRAPipeline]: - """Resolve pipeline class based on pipeline name in the config, pipeline type, and workload type.""" - if not pipeline_name_in_config: - logger.warning("No pipeline architecture is specified") - - pipeline_cls = self._try_load_pipeline_cls( - pipeline_name_in_config, pipeline_type, workload_type - ) - if pipeline_cls is not None: - return pipeline_cls - supported_archs = self.get_supported_archs( - pipeline_name_in_config, pipeline_type - ) - raise ValueError( - f"Pipeline architecture '{pipeline_name_in_config}' is not supported for pipeline type '{pipeline_type.value}' " - f"and workload type '{workload_type.value}'. " - f"Supported architectures: {supported_archs}" - ) - - -@lru_cache -def import_pipeline_classes( - pipeline_types: list[PipelineType] | PipelineType | None = None, -) -> dict[str, dict[str, type[ComposedPipelineBase] | None]]: - """ - Import pipeline classes based on the pipeline type and workload type. - - Args: - pipeline_types: The pipeline types to load (basic, preprocess). - If None, loads all types. - - Returns: - A three-level nested dictionary: - {pipeline_type: {architecture_name: {pipeline_name: pipeline_cls}}} - e.g., {"basic": {"wan": {"WanPipeline": WanPipeline}}} - """ - type_to_pipeline_dict: dict[str, dict[str, type[ComposedPipelineBase] | None]] = {} - package_name: str = "sglang.multimodal_gen.runtime.architectures" - - # Determine which pipeline types to scan - if isinstance(pipeline_types, list): - pipeline_types_to_scan = [ - pipeline_type.value for pipeline_type in pipeline_types - ] - elif isinstance(pipeline_types, PipelineType): - pipeline_types_to_scan = [pipeline_types.value] - else: - pipeline_types_to_scan = [pt.value for pt in PipelineType] - - logger.info("Loading pipelines for types: %s", pipeline_types_to_scan) - - for pipeline_type_str in pipeline_types_to_scan: - # Try to load from pipeline-type-specific directory first - pipeline_type_package_name = f"{package_name}.{pipeline_type_str}" - pipeline_dict: dict[str, type[ComposedPipelineBase] | None] = {} - - try: - pipeline_type_package = importlib.import_module(pipeline_type_package_name) - logger.debug("Successfully imported %s", pipeline_type_package_name) - - for _, arch, ispkg in pkgutil.iter_modules(pipeline_type_package.__path__): - - arch_package_name = f"{pipeline_type_package_name}.{arch}" - if ispkg: - arch_package = importlib.import_module(arch_package_name) - for _, module_name, ispkg in pkgutil.walk_packages( - arch_package.__path__, arch_package_name + "." - ): - if not ispkg: - pipeline_module = importlib.import_module(module_name) - if hasattr(pipeline_module, "EntryClass"): - entry_cls_list = pipeline_module.EntryClass - if not isinstance(entry_cls_list, list): - entry_cls_list = [entry_cls_list] - - if isinstance(pipeline_module.EntryClass, list): - pipeline_names = [ - pipeline.__name__ - for pipeline in pipeline_module.EntryClass - ] - else: - pipeline_names = [ - pipeline_module.EntryClass.__name__ - ] - - for entry_cls, pipeline_name in zip( - entry_cls_list, pipeline_names - ): - assert ( - pipeline_name not in pipeline_dict - ), f"Duplicated pipeline implementation for {pipeline_name} in {pipeline_type_str}.{arch_package_name}" - - assert hasattr( - entry_cls, "pipeline_name" - ), f"{entry_cls}" - pipeline_dict[pipeline_name] = entry_cls - - type_to_pipeline_dict[pipeline_type_str] = pipeline_dict - - except ImportError as e: - raise ImportError( - f"Could not import {pipeline_type_package_name} when importing pipeline classes: {e}" - ) from None - - # Log summary - total_pipelines = sum( - len(pipeline_dict) for pipeline_dict in type_to_pipeline_dict.values() - ) - logger.info( - "Loaded %d pipeline classes across %d types", - total_pipelines, - len(pipeline_types_to_scan), - ) - - return type_to_pipeline_dict - - -def get_pipeline_registry( - pipeline_type: PipelineType | str | None = None, -) -> _PipelineRegistry: - """ - Get a pipeline registry for the specified mode, pipeline type, and workload type. - - Args: - pipeline_type: Pipeline type to load. If None and mode is provided, will be derived from mode. - - Returns: - A pipeline registry instance. - """ - if isinstance(pipeline_type, str): - pipeline_type = PipelineType.from_string(pipeline_type) - - pipeline_classes = import_pipeline_classes(pipeline_type) - return _PipelineRegistry(pipeline_classes) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/pipeline_batch_info.py b/python/sglang/multimodal_gen/runtime/pipelines/schedule_batch.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/pipelines/pipeline_batch_info.py rename to python/sglang/multimodal_gen/runtime/pipelines/schedule_batch.py diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/base.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/base.py index eb89dbe7c..6eeed5f78 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/base.py @@ -16,7 +16,7 @@ from enum import Enum, auto import torch import sglang.multimodal_gen.envs as envs -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages.validators import VerificationResult from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/causal_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/causal_denoising.py index 689be4541..27b5ac2a6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/causal_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/causal_denoising.py @@ -5,7 +5,7 @@ import torch # type: ignore 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.models.utils import pred_noise_to_pred_video -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages.denoising import DenoisingStage from sglang.multimodal_gen.runtime.pipelines.stages.validators import ( StageValidators as V, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/conditioning.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/conditioning.py index fb47b2948..38a6fd389 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/conditioning.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/conditioning.py @@ -7,7 +7,7 @@ Conditioning stage for diffusion pipelines. import torch -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages.base import PipelineStage from sglang.multimodal_gen.runtime.pipelines.stages.validators import ( StageValidators as V, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/decoding.py index 0728586f5..f4198bf5d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/decoding.py @@ -17,7 +17,7 @@ from sglang.multimodal_gen.configs.pipelines.qwen_image import ( from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.component_loader import VAELoader from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import OutputBatch, Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import OutputBatch, Req from sglang.multimodal_gen.runtime.pipelines.stages.base import ( PipelineStage, StageParallelismType, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py index 17aae037a..2e3922b11 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py @@ -44,7 +44,7 @@ from sglang.multimodal_gen.runtime.layers.attention.STA_configuration import ( ) from sglang.multimodal_gen.runtime.loader.component_loader import TransformerLoader from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages.base import ( PipelineStage, StageParallelismType, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising_dmd.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising_dmd.py index 1d39aaf8e..890428dcb 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising_dmd.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising_dmd.py @@ -23,7 +23,7 @@ from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler FlowMatchEulerDiscreteScheduler, ) from sglang.multimodal_gen.runtime.models.utils import pred_noise_to_pred_video -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages import DenoisingStage from sglang.multimodal_gen.runtime.pipelines.stages.denoising import ( st_attn_available, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/encoding.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/encoding.py index dbea07442..e06937f7a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/encoding.py @@ -9,7 +9,7 @@ import torch from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages.base import PipelineStage from sglang.multimodal_gen.runtime.pipelines.stages.validators import ( V, # Import validators diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/image_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/image_encoding.py index 0f91451da..421eff15c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/image_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/image_encoding.py @@ -25,7 +25,7 @@ from sglang.multimodal_gen.runtime.models.vision_utils import ( pil_to_numpy, resize, ) -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages.base import PipelineStage from sglang.multimodal_gen.runtime.pipelines.stages.validators import ( StageValidators as V, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/input_validation.py index 9c3fd2fc2..1fd09b3c5 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/input_validation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/input_validation.py @@ -14,7 +14,7 @@ from sglang.multimodal_gen.configs.pipelines.qwen_image import ( QwenImageEditPipelineConfig, ) from sglang.multimodal_gen.runtime.models.vision_utils import load_image, load_video -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages.base import PipelineStage from sglang.multimodal_gen.runtime.pipelines.stages.validators import ( StageValidators, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/latent_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/latent_preparation.py index 55f4fc86e..1bfccf0f4 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/latent_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/latent_preparation.py @@ -7,7 +7,7 @@ Latent preparation stage for diffusion pipelines. from diffusers.utils.torch_utils import randn_tensor from sglang.multimodal_gen.runtime.distributed import get_local_torch_device -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages.base import PipelineStage from sglang.multimodal_gen.runtime.pipelines.stages.validators import ( StageValidators as V, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/stepvideo_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/stepvideo_encoding.py index 54aa6b45c..81e39d7b6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/stepvideo_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/stepvideo_encoding.py @@ -5,7 +5,7 @@ import torch from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages.base import PipelineStage from sglang.multimodal_gen.runtime.pipelines.stages.validators import ( StageValidators as V, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/text_encoding.py index eff5ee1c9..cf6c342f2 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/text_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/text_encoding.py @@ -13,7 +13,7 @@ from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput from sglang.multimodal_gen.configs.pipelines import FluxPipelineConfig 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.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages.base import PipelineStage from sglang.multimodal_gen.runtime.pipelines.stages.validators import ( StageValidators as V, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/timestep_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/timestep_preparation.py index 09c5d22ee..28275f11c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/timestep_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/timestep_preparation.py @@ -18,7 +18,7 @@ from sglang.multimodal_gen.configs.pipelines.qwen_image import ( QwenImagePipelineConfig, ) from sglang.multimodal_gen.runtime.distributed import get_local_torch_device -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines.stages.base import ( PipelineStage, StageParallelismType, diff --git a/python/sglang/multimodal_gen/runtime/scheduler_client.py b/python/sglang/multimodal_gen/runtime/scheduler_client.py index 97cc1165e..5569ae209 100644 --- a/python/sglang/multimodal_gen/runtime/scheduler_client.py +++ b/python/sglang/multimodal_gen/runtime/scheduler_client.py @@ -5,7 +5,7 @@ import asyncio import zmq import zmq.asyncio -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 95a7f12d0..02b1dea0b 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -335,9 +335,12 @@ class ServerArgs: return self.host is None or self.port is None def __post_init__(self): - self.scheduler_port = self.settle_port(self.scheduler_port) + # Add randomization to avoid race condition when multiple servers start simultaneously + initial_scheduler_port = self.scheduler_port + random.randint(0, 100) + self.scheduler_port = self.settle_port(initial_scheduler_port) # TODO: remove hard code - self.master_port = self.settle_port(self.master_port or 30005, 37) + initial_master_port = (self.master_port or 30005) + random.randint(0, 100) + self.master_port = self.settle_port(initial_master_port, 37) if self.moba_config_path: try: with open(self.moba_config_path) as f: @@ -646,14 +649,45 @@ class ServerArgs: scheduler_host = self.host or "localhost" return f"tcp://{scheduler_host}:{self.scheduler_port}" - def settle_port(self, port: int, port_inc: int = 42) -> int: - while True: + def settle_port( + self, port: int, port_inc: int = 42, max_attempts: int = 100 + ) -> int: + """ + Find an available port with retry logic. + + Args: + port: Initial port to check + port_inc: Port increment for each attempt + max_attempts: Maximum number of attempts to find an available port + + Returns: + An available port number + + Raises: + RuntimeError: If no available port is found after max_attempts + """ + attempts = 0 + original_port = port + + while attempts < max_attempts: if is_port_available(port): + if attempts > 0: + logger.info( + f"Port {original_port} was unavailable, using port {port} instead" + ) return port + + attempts += 1 if port < 60000: port += port_inc else: - port -= port_inc + 1 + # Wrap around with randomization to avoid collision + port = 5000 + random.randint(0, 1000) + + raise RuntimeError( + f"Failed to find available port after {max_attempts} attempts " + f"(started from port {original_port})" + ) def post_init_serve(self): """ diff --git a/python/sglang/multimodal_gen/runtime/sync_scheduler_client.py b/python/sglang/multimodal_gen/runtime/sync_scheduler_client.py index 93359f34d..2e246fb82 100644 --- a/python/sglang/multimodal_gen/runtime/sync_scheduler_client.py +++ b/python/sglang/multimodal_gen/runtime/sync_scheduler_client.py @@ -2,7 +2,7 @@ import zmq -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import Req +from sglang.multimodal_gen.runtime.pipelines.schedule_batch import Req from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger diff --git a/python/sglang/multimodal_gen/runtime/utils/common.py b/python/sglang/multimodal_gen/runtime/utils/common.py index c39769ae8..dbf3661f3 100644 --- a/python/sglang/multimodal_gen/runtime/utils/common.py +++ b/python/sglang/multimodal_gen/runtime/utils/common.py @@ -126,8 +126,26 @@ def is_port_available(port): def get_zmq_socket( - context: zmq.Context, socket_type: zmq.SocketType, endpoint: str, bind: bool -) -> zmq.Socket: + context: zmq.Context, + socket_type: zmq.SocketType, + endpoint: str, + bind: bool, + max_bind_retries: int = 10, +) -> tuple[zmq.Socket, str]: + """ + Create and configure a ZMQ socket. + + Args: + context: ZMQ context + socket_type: Type of ZMQ socket + endpoint: Endpoint string (e.g., "tcp://localhost:5555") + bind: Whether to bind (True) or connect (False) + max_bind_retries: Maximum number of retries if bind fails due to address already in use + + Returns: + A tuple of (socket, actual_endpoint). The actual_endpoint may differ from the + requested endpoint if bind retry was needed. + """ mem = psutil.virtual_memory() total_mem = mem.total / 1024**3 available_mem = mem.available / 1024**3 @@ -165,11 +183,67 @@ def get_zmq_socket( raise ValueError(f"Unsupported socket type: {socket_type}") if bind: - socket.bind(endpoint) + # Parse port from endpoint for retry logic + import re + + port_match = re.search(r":(\d+)$", endpoint) + + if port_match and max_bind_retries > 1: + original_port = int(port_match.group(1)) + last_exception = None + + for attempt in range(max_bind_retries): + try: + current_endpoint = endpoint + if attempt > 0: + # Try next port (increment by 42 to match settle_port logic) + current_port = original_port + attempt * 42 + current_endpoint = re.sub( + r":(\d+)$", f":{current_port}", endpoint + ) + logger.info( + f"ZMQ bind failed for port {original_port + (attempt - 1) * 42}, " + f"retrying with port {current_port} (attempt {attempt + 1}/{max_bind_retries})" + ) + + socket.bind(current_endpoint) + + if attempt > 0: + logger.warning( + f"Successfully bound ZMQ socket to {current_endpoint} after {attempt + 1} attempts. " + f"Original port {original_port} was unavailable." + ) + + return socket, current_endpoint + + except zmq.ZMQError as e: + last_exception = e + if e.errno == zmq.EADDRINUSE and attempt < max_bind_retries - 1: + # Address already in use, try next port + continue + elif attempt == max_bind_retries - 1: + # Last attempt failed + logger.error( + f"Failed to bind ZMQ socket after {max_bind_retries} attempts. " + f"Original endpoint: {endpoint}, Last tried port: {original_port + attempt * 42}" + ) + raise + else: + # Different error, raise immediately + raise + + # Should not reach here, but just in case + if last_exception: + raise last_exception + else: + # No retry logic needed (either no port in endpoint or max_bind_retries == 1) + socket.bind(endpoint) + return socket, endpoint else: socket.connect(endpoint) + return socket, endpoint - return socket + return socket, endpoint # https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip diff --git a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py index 64ece7951..be99b3ac2 100644 --- a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py @@ -34,7 +34,7 @@ _warned_main_process = False _FORMAT = ( f"{SGL_DIFFUSION_LOGGING_PREFIX}%(levelname)s %(asctime)s " - "[%(filename)s:%(lineno)d] %(message)s" + "[%(filename)s: %(lineno)d] %(message)s" ) # _FORMAT = "[%(asctime)s] %(message)s" diff --git a/python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow.py b/python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow.py index 3d3a831ae..38fba4214 100644 --- a/python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow.py +++ b/python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow.py @@ -15,7 +15,6 @@ from sglang.multimodal_gen.dataset.dataloader.schema import ( pyarrow_schema_t2v, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import get_world_rank -from sglang.multimodal_gen.runtime.pipelines.pipeline_registry import PipelineType from sglang.multimodal_gen.runtime.server_args import ServerArgs, WorkloadType from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.workflow.preprocess.components import ( @@ -32,9 +31,7 @@ logger = init_logger(__name__) class PreprocessWorkflow(WorkflowBase): def register_pipelines(self) -> None: - self.add_pipeline_config( - "preprocess_pipeline", (PipelineType.PREPROCESS, self.server_args) - ) + self.add_pipeline_config("preprocess_pipeline", self.server_args) def register_components(self) -> None: assert self.server_args.preprocess_config is not None diff --git a/python/sglang/multimodal_gen/test/cli/test_generate_common.py b/python/sglang/multimodal_gen/test/cli/test_generate_common.py index aff1c98aa..49dfa653f 100644 --- a/python/sglang/multimodal_gen/test/cli/test_generate_common.py +++ b/python/sglang/multimodal_gen/test/cli/test_generate_common.py @@ -47,7 +47,8 @@ class TestGenerate(TestCLIBase): command = [ "sglang", "generate", - "--prompt='A curious raccoon'", + "--prompt", + "A curious raccoon", "--output-path=outputs", f"--model-path={self.model_path}", "--save-output", @@ -84,7 +85,8 @@ class TestWanGenerate(TestGenerate): command = [ "sglang", "generate", - "--prompt='A curious raccoon'", + "--prompt", + "A curious raccoon", "--output-path=outputs", f"--model-path={self.model_path}", "--save-output", diff --git a/python/sglang/multimodal_gen/test/cli/test_generate_t2i_perf.py b/python/sglang/multimodal_gen/test/cli/test_generate_t2i_perf.py index bbfde89a4..800828ccf 100644 --- a/python/sglang/multimodal_gen/test/cli/test_generate_t2i_perf.py +++ b/python/sglang/multimodal_gen/test/cli/test_generate_t2i_perf.py @@ -48,7 +48,8 @@ class TestQwenImageEdit(TestGenerateBase): "generate", "--text-encoder-cpu-offload", "--pin-cpu-memory", - f"--prompt='{self.prompt}'", + f"--prompt", + f"{self.prompt}", "--save-output", "--log-level=debug", f"--width={self.width}", @@ -58,7 +59,7 @@ class TestQwenImageEdit(TestGenerateBase): def test_single_gpu(self): self._run_test( - name=f"{self.model_name()}, single gpu", + name=f"{self.model_name()}_single_gpu", args=None, model_path=self.model_path, test_key="test_single_gpu", diff --git a/python/sglang/multimodal_gen/test/cli/test_generate_t2v_perf.py b/python/sglang/multimodal_gen/test/cli/test_generate_t2v_perf.py index b0e3e2a52..3190c15af 100644 --- a/python/sglang/multimodal_gen/test/cli/test_generate_t2v_perf.py +++ b/python/sglang/multimodal_gen/test/cli/test_generate_t2v_perf.py @@ -17,7 +17,7 @@ class TestFastWan2_1_T2V(TestGenerateBase): "test_single_gpu": 13.0, "test_cfg_parallel": 15.0, "test_usp": 15.0, - "test_mixed": 15.0, + "test_mixed": 15.0 * 1.05, } # disabled for vsa @@ -44,22 +44,31 @@ class TestWan2_1_T2V(TestGenerateBase): thresholds = { "test_single_gpu": 76.0 * 1.05, "test_cfg_parallel": 46.5 * 1.05, - "test_usp": 22.5, - "test_mixed": 26.5, + "test_usp": 39.8 * 1.05, + "test_mixed": 37.3 * 1.05, } + def test_mixed(self): + pass + + def test_cfg_parallel(self): + pass + class TestWan2_2_T2V(TestGenerateBase): model_path = "Wan-AI/Wan2.2-T2V-A14B-Diffusers" extra_args = [] data_type: DataType = DataType.VIDEO thresholds = { - "test_single_gpu": 865, + "test_single_gpu": 904.3 * 1.05, "test_cfg_parallel": 446, - "test_usp": 124, + "test_usp": 316 * 1.05, "test_mixed": 159, } + def test_single_gpu(self): + pass + def test_mixed(self): pass diff --git a/python/sglang/multimodal_gen/test/cli/test_generate_ti2v_perf.py b/python/sglang/multimodal_gen/test/cli/test_generate_ti2v_perf.py index 092725bdf..51948f063 100644 --- a/python/sglang/multimodal_gen/test/cli/test_generate_ti2v_perf.py +++ b/python/sglang/multimodal_gen/test/cli/test_generate_ti2v_perf.py @@ -15,7 +15,8 @@ class TestGenerateTI2VBase(TestGenerateBase): cls.base_command = [ "sglang", "generate", - f'--prompt="Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline\'s intricate details and the refreshing atmosphere of the seaside."', + "--prompt", + "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside.", "--image-path", "https://github.com/Wan-Video/Wan2.2/blob/990af50de458c19590c245151197326e208d7191/examples/i2v_input.JPG?raw=true", "--save-output", @@ -36,14 +37,14 @@ class TestGenerateTI2VBase(TestGenerateBase): class TestWan2_1_I2V_14B_480P(TestGenerateTI2VBase): model_path = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" thresholds = { - "test_usp": 530.5 * 1.05, + "test_usp": 557.9 * 1.05, } class TestWan2_1_I2V_14B_720P(TestGenerateTI2VBase): model_path = "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers" thresholds = { - "test_usp": 530.5 * 1.05, + "test_usp": 558.4 * 1.05, } diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index b4d60e268..86847797a 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -19,7 +19,7 @@ logger = init_logger(__name__) def run_command(command) -> Optional[float]: """Runs a command and returns the execution time and status.""" - print(f"Running command: {' '.join(command)}") + print(f"Running command: {shlex.join(command)}") duration = None with subprocess.Popen( @@ -105,7 +105,8 @@ class TestCLIBase(unittest.TestCase): "generate", "--text-encoder-cpu-offload", "--pin-cpu-memory", - "--prompt='A curious raccoon'", + "--prompt", + "A curious raccoon", "--save-output", "--log-level=debug", f"--width={width}", @@ -124,7 +125,7 @@ class TestCLIBase(unittest.TestCase): self.base_command + [f"--model-path={model_path}"] + shlex.split(args or "") - + [f"--output-file-name={name}"] + + ["--output-file-name", f"{name}"] + self.extra_args ) duration = run_command(command) @@ -155,7 +156,8 @@ class TestGenerateBase(TestCLIBase): "generate", # "--text-encoder-cpu-offload", # "--pin-cpu-memory", - f"--prompt='{prompt}'", + f"--prompt", + f"{prompt}", "--save-output", "--log-level=debug", f"--width={width}", @@ -237,7 +239,7 @@ class TestGenerateBase(TestCLIBase): def test_single_gpu(self): """single gpu""" self._run_test( - name=f"{self.model_name()}_single gpu", + name=f"{self.model_name()}_single_gpu", args=None, model_path=self.model_path, test_key="test_single_gpu", @@ -248,7 +250,7 @@ class TestGenerateBase(TestCLIBase): if self.data_type == DataType.IMAGE: return self._run_test( - name=f"{self.model_name()}_cfg parallel", + name=f"{self.model_name()}_cfg_parallel", args="--num-gpus 2 --enable-cfg-parallel", model_path=self.model_path, test_key="test_cfg_parallel",