[diffusion] fix: fix diffusers backend issues in diffusion ci gt workflow (#20173)

This commit is contained in:
Yuhao Yang
2026-03-10 00:51:48 +08:00
committed by GitHub
parent f947bcbd89
commit ecca8c553d
3 changed files with 60 additions and 11 deletions
+2 -2
View File
@@ -29,7 +29,7 @@ jobs:
strategy: strategy:
matrix: matrix:
part: [0, 1] part: [0, 1]
timeout-minutes: 60 timeout-minutes: 150
steps: steps:
- name: Checkout code - name: Checkout code
uses: actions/checkout@v4 uses: actions/checkout@v4
@@ -63,7 +63,7 @@ jobs:
strategy: strategy:
matrix: matrix:
part: [0, 1] part: [0, 1]
timeout-minutes: 60 timeout-minutes: 150
steps: steps:
- name: Checkout code - name: Checkout code
uses: actions/checkout@v4 uses: actions/checkout@v4
+42 -8
View File
@@ -336,11 +336,17 @@ class ModelInfo:
pipeline_config_cls: Type[PipelineConfig] pipeline_config_cls: Type[PipelineConfig]
def _get_diffusers_model_info() -> ModelInfo: def _get_diffusers_model_info(
model_path: Optional[str] = None,
model_id: Optional[str] = None,
) -> ModelInfo:
""" """
Get model info for diffusers backend. Get model info for diffusers backend.
Returns a ModelInfo with DiffusersPipeline and generic configs. Returns a ModelInfo with DiffusersPipeline and generic configs.
When model_path is provided and has a registered native config,
inherits task_type from it so that validation (e.g. accepts_image_input)
works correctly even under the diffusers backend.
""" """
from sglang.multimodal_gen.configs.pipeline_configs.diffusers_generic import ( from sglang.multimodal_gen.configs.pipeline_configs.diffusers_generic import (
DiffusersGenericPipelineConfig, DiffusersGenericPipelineConfig,
@@ -352,10 +358,34 @@ def _get_diffusers_model_info() -> ModelInfo:
DiffusersPipeline, DiffusersPipeline,
) )
pipeline_config_cls = DiffusersGenericPipelineConfig
# If there is a registered native config for this model, inherit its task_type
if model_path is not None:
config_info = _get_config_info(model_path, model_id=model_id)
if config_info is not None:
native_task_type = config_info.pipeline_config_cls.task_type
if native_task_type != DiffusersGenericPipelineConfig.task_type:
pipeline_config_cls = dataclasses.make_dataclass(
"DiffusersGenericPipelineConfig",
[
(
"task_type",
type(native_task_type),
dataclasses.field(default=native_task_type),
)
],
bases=(DiffusersGenericPipelineConfig,),
)
logger.debug(
"Inherited task_type=%s from native config for diffusers backend",
native_task_type.name,
)
return ModelInfo( return ModelInfo(
pipeline_cls=DiffusersPipeline, pipeline_cls=DiffusersPipeline,
sampling_param_cls=DiffusersGenericSamplingParams, sampling_param_cls=DiffusersGenericSamplingParams,
pipeline_config_cls=DiffusersGenericPipelineConfig, pipeline_config_cls=pipeline_config_cls,
) )
@@ -392,7 +422,7 @@ def get_model_info(
logger.info( logger.info(
"Using diffusers backend for model '%s' (explicitly requested)", model_path "Using diffusers backend for model '%s' (explicitly requested)", model_path
) )
return _get_diffusers_model_info() return _get_diffusers_model_info(model_path=model_path, model_id=model_id)
# For AUTO or SGLANG backend, try native implementation first # For AUTO or SGLANG backend, try native implementation first
# 1. Discover all available pipeline classes and cache them # 1. Discover all available pipeline classes and cache them
@@ -407,7 +437,7 @@ def get_model_info(
"Falling back to diffusers backend.", "Falling back to diffusers backend.",
model_path, model_path,
) )
return _get_diffusers_model_info(model_path) return _get_diffusers_model_info(model_path=model_path, model_id=model_id)
# 2. Get pipeline class - check non-diffusers models first # 2. Get pipeline class - check non-diffusers models first
pipeline_class_name = get_non_diffusers_pipeline_name(model_path) pipeline_class_name = get_non_diffusers_pipeline_name(model_path)
@@ -427,7 +457,9 @@ def get_model_info(
logger.error(f"Could not read model config for '{model_path}': {e}") logger.error(f"Could not read model config for '{model_path}': {e}")
if backend == Backend.AUTO: if backend == Backend.AUTO:
logger.info("Falling back to diffusers backend") logger.info("Falling back to diffusers backend")
return _get_diffusers_model_info() return _get_diffusers_model_info(
model_path=model_path, model_id=model_id
)
return None return None
pipeline_class_name = config.get("_class_name") pipeline_class_name = config.get("_class_name")
@@ -437,7 +469,9 @@ def get_model_info(
) )
if backend == Backend.AUTO: if backend == Backend.AUTO:
logger.info("Falling back to diffusers backend") logger.info("Falling back to diffusers backend")
return _get_diffusers_model_info() return _get_diffusers_model_info(
model_path=model_path, model_id=model_id
)
return None return None
pipeline_cls = _PIPELINE_REGISTRY.get(pipeline_class_name) pipeline_cls = _PIPELINE_REGISTRY.get(pipeline_class_name)
@@ -447,7 +481,7 @@ def get_model_info(
f"Pipeline class '{pipeline_class_name}' specified in '{model_path}' has no native sglang support. " f"Pipeline class '{pipeline_class_name}' specified in '{model_path}' has no native sglang support. "
f"Falling back to diffusers backend." f"Falling back to diffusers backend."
) )
return _get_diffusers_model_info() return _get_diffusers_model_info(model_path=model_path, model_id=model_id)
else: else:
logger.error( logger.error(
f"Pipeline class '{pipeline_class_name}' specified in '{model_path}' is not a registered EntryClass in the framework. " f"Pipeline class '{pipeline_class_name}' specified in '{model_path}' is not a registered EntryClass in the framework. "
@@ -464,7 +498,7 @@ def get_model_info(
f"Could not resolve native configuration for model '{model_path}'. " f"Could not resolve native configuration for model '{model_path}'. "
f"Falling back to diffusers backend." f"Falling back to diffusers backend."
) )
return _get_diffusers_model_info() return _get_diffusers_model_info(model_path=model_path, model_id=model_id)
else: else:
logger.error( logger.error(
f"Could not resolve configuration for model '{model_path}'. " f"Could not resolve configuration for model '{model_path}'. "
@@ -456,7 +456,22 @@ class DiffusersPipeline(ComposedPipelineBase):
else: else:
raise raise
pipe = pipe.to(get_local_torch_device()) # Use CPU offload (all-or-nothing in diffusers) if any component offload is requested.
any_offload = (
server_args.dit_cpu_offload
or server_args.text_encoder_cpu_offload
or server_args.image_encoder_cpu_offload
or server_args.vae_cpu_offload
)
if any_offload:
device = get_local_torch_device()
gpu_id = device.index if device.index is not None else 0
pipe.enable_model_cpu_offload(gpu_id=gpu_id)
logger.info(
"Enabled model CPU offload for diffusers pipeline (gpu_id=%d)", gpu_id
)
else:
pipe = pipe.to(get_local_torch_device())
# Apply VAE memory optimizations from pipeline config # Apply VAE memory optimizations from pipeline config
self._apply_vae_optimizations(pipe, server_args) self._apply_vae_optimizations(pipe, server_args)
# Apply attention backend if specified # Apply attention backend if specified