[diffusion] fix: fix PicklingError with --backend diffusers on non-T2I models (#21472)
Co-authored-by: Yufeng He <40085740+universeplayer@users.noreply.github.com>
This commit is contained in:
@@ -78,3 +78,52 @@ class DiffusersGenericPipelineConfig(PipelineConfig):
|
|||||||
Pass through - diffusers handles frame count.
|
Pass through - diffusers handles frame count.
|
||||||
"""
|
"""
|
||||||
return num_frames
|
return num_frames
|
||||||
|
|
||||||
|
|
||||||
|
# Static subclasses for each non-default task type.
|
||||||
|
# These exist so that _get_diffusers_model_info() can swap the task_type without
|
||||||
|
# dynamically creating a class via make_dataclass -- dynamic classes break pickle
|
||||||
|
# when multiprocessing uses the 'spawn' start method (see #21453).
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class DiffusersT2VPipelineConfig(DiffusersGenericPipelineConfig):
|
||||||
|
task_type: ModelTaskType = ModelTaskType.T2V
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class DiffusersI2VPipelineConfig(DiffusersGenericPipelineConfig):
|
||||||
|
task_type: ModelTaskType = ModelTaskType.I2V
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class DiffusersTI2VPipelineConfig(DiffusersGenericPipelineConfig):
|
||||||
|
task_type: ModelTaskType = ModelTaskType.TI2V
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class DiffusersI2IPipelineConfig(DiffusersGenericPipelineConfig):
|
||||||
|
task_type: ModelTaskType = ModelTaskType.I2I
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class DiffusersTI2IPipelineConfig(DiffusersGenericPipelineConfig):
|
||||||
|
task_type: ModelTaskType = ModelTaskType.TI2I
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class DiffusersI2MPipelineConfig(DiffusersGenericPipelineConfig):
|
||||||
|
task_type: ModelTaskType = ModelTaskType.I2M
|
||||||
|
|
||||||
|
|
||||||
|
DIFFUSERS_TASK_TYPE_TO_CONFIG: dict[
|
||||||
|
ModelTaskType, type[DiffusersGenericPipelineConfig]
|
||||||
|
] = {
|
||||||
|
ModelTaskType.T2I: DiffusersGenericPipelineConfig,
|
||||||
|
ModelTaskType.T2V: DiffusersT2VPipelineConfig,
|
||||||
|
ModelTaskType.I2V: DiffusersI2VPipelineConfig,
|
||||||
|
ModelTaskType.TI2V: DiffusersTI2VPipelineConfig,
|
||||||
|
ModelTaskType.I2I: DiffusersI2IPipelineConfig,
|
||||||
|
ModelTaskType.TI2I: DiffusersTI2IPipelineConfig,
|
||||||
|
ModelTaskType.I2M: DiffusersI2MPipelineConfig,
|
||||||
|
}
|
||||||
|
|||||||
@@ -11,7 +11,6 @@ import dataclasses
|
|||||||
import importlib
|
import importlib
|
||||||
import os
|
import os
|
||||||
import pkgutil
|
import pkgutil
|
||||||
import sys
|
|
||||||
from functools import lru_cache
|
from functools import lru_cache
|
||||||
from typing import (
|
from typing import (
|
||||||
TYPE_CHECKING,
|
TYPE_CHECKING,
|
||||||
@@ -453,6 +452,7 @@ def _get_diffusers_model_info(
|
|||||||
works correctly even under the diffusers backend.
|
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 (
|
||||||
|
DIFFUSERS_TASK_TYPE_TO_CONFIG,
|
||||||
DiffusersGenericPipelineConfig,
|
DiffusersGenericPipelineConfig,
|
||||||
)
|
)
|
||||||
from sglang.multimodal_gen.configs.sample.diffusers_generic import (
|
from sglang.multimodal_gen.configs.sample.diffusers_generic import (
|
||||||
@@ -465,36 +465,17 @@ def _get_diffusers_model_info(
|
|||||||
sampling_param_cls = DiffusersGenericSamplingParams
|
sampling_param_cls = DiffusersGenericSamplingParams
|
||||||
pipeline_config_cls = DiffusersGenericPipelineConfig
|
pipeline_config_cls = DiffusersGenericPipelineConfig
|
||||||
|
|
||||||
# If there is a registered native config for this model, inherit its task_type
|
# If there is a registered native config for this model, inherit its task_type.
|
||||||
|
# We use pre-defined static subclasses instead of make_dataclass so the config
|
||||||
|
# class is pickle-safe for multiprocessing spawn (fixes #21453).
|
||||||
if model_path is not None:
|
if model_path is not None:
|
||||||
config_info = _get_config_info(model_path, model_id=model_id)
|
config_info = _get_config_info(model_path, model_id=model_id)
|
||||||
if config_info is not None:
|
if config_info is not None:
|
||||||
sampling_param_cls = config_info.sampling_param_cls
|
sampling_param_cls = config_info.sampling_param_cls
|
||||||
native_task_type = config_info.pipeline_config_cls.task_type
|
native_task_type = config_info.pipeline_config_cls.task_type
|
||||||
if native_task_type != DiffusersGenericPipelineConfig.task_type:
|
if native_task_type != DiffusersGenericPipelineConfig.task_type:
|
||||||
pipeline_config_cls = dataclasses.make_dataclass(
|
pipeline_config_cls = DIFFUSERS_TASK_TYPE_TO_CONFIG.get(
|
||||||
"DiffusersGenericPipelineConfig",
|
native_task_type, DiffusersGenericPipelineConfig
|
||||||
[
|
|
||||||
(
|
|
||||||
"task_type",
|
|
||||||
type(native_task_type),
|
|
||||||
dataclasses.field(default=native_task_type),
|
|
||||||
)
|
|
||||||
],
|
|
||||||
bases=(DiffusersGenericPipelineConfig,),
|
|
||||||
)
|
|
||||||
# make_dataclass sets __module__="types"; fix for pickle.
|
|
||||||
pipeline_config_cls.__module__ = (
|
|
||||||
DiffusersGenericPipelineConfig.__module__
|
|
||||||
)
|
|
||||||
pipeline_config_cls.__qualname__ = (
|
|
||||||
DiffusersGenericPipelineConfig.__qualname__
|
|
||||||
)
|
|
||||||
parent_module = sys.modules[DiffusersGenericPipelineConfig.__module__]
|
|
||||||
setattr(
|
|
||||||
parent_module,
|
|
||||||
DiffusersGenericPipelineConfig.__name__,
|
|
||||||
pipeline_config_cls,
|
|
||||||
)
|
)
|
||||||
logger.debug(
|
logger.debug(
|
||||||
"Inherited task_type=%s from native config for diffusers backend",
|
"Inherited task_type=%s from native config for diffusers backend",
|
||||||
|
|||||||
Reference in New Issue
Block a user