diffusion: reduce effort of supporting new model (#12982)

This commit is contained in:
Mick
2025-11-10 21:20:33 +08:00
committed by GitHub
parent afee2843d5
commit 5639145fac
66 changed files with 667 additions and 2385 deletions
@@ -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",
]
@@ -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):
@@ -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):
@@ -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
@@ -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.",
@@ -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