[diffusion] fix: fix diffusers backend issues in diffusion ci gt workflow (#20173)
This commit is contained in:
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user