diff --git a/.github/workflows/diffusion-ci-gt-gen.yml b/.github/workflows/diffusion-ci-gt-gen.yml index b1af552f9..60ef0f698 100644 --- a/.github/workflows/diffusion-ci-gt-gen.yml +++ b/.github/workflows/diffusion-ci-gt-gen.yml @@ -29,7 +29,7 @@ jobs: strategy: matrix: part: [0, 1] - timeout-minutes: 60 + timeout-minutes: 150 steps: - name: Checkout code uses: actions/checkout@v4 @@ -63,7 +63,7 @@ jobs: strategy: matrix: part: [0, 1] - timeout-minutes: 60 + timeout-minutes: 150 steps: - name: Checkout code uses: actions/checkout@v4 diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index c207cd3a4..d61fcb2ee 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -336,11 +336,17 @@ class ModelInfo: 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. 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 ( DiffusersGenericPipelineConfig, @@ -352,10 +358,34 @@ def _get_diffusers_model_info() -> ModelInfo: 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( pipeline_cls=DiffusersPipeline, sampling_param_cls=DiffusersGenericSamplingParams, - pipeline_config_cls=DiffusersGenericPipelineConfig, + pipeline_config_cls=pipeline_config_cls, ) @@ -392,7 +422,7 @@ def get_model_info( logger.info( "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 # 1. Discover all available pipeline classes and cache them @@ -407,7 +437,7 @@ def get_model_info( "Falling back to diffusers backend.", 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 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}") if backend == Backend.AUTO: 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 pipeline_class_name = config.get("_class_name") @@ -437,7 +469,9 @@ def get_model_info( ) if backend == Backend.AUTO: 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 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"Falling back to diffusers backend." ) - return _get_diffusers_model_info() + return _get_diffusers_model_info(model_path=model_path, model_id=model_id) else: logger.error( 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"Falling back to diffusers backend." ) - return _get_diffusers_model_info() + return _get_diffusers_model_info(model_path=model_path, model_id=model_id) else: logger.error( f"Could not resolve configuration for model '{model_path}'. " diff --git a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py index f4f6f87cb..e2037097b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py @@ -456,7 +456,22 @@ class DiffusersPipeline(ComposedPipelineBase): else: 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 self._apply_vae_optimizations(pipe, server_args) # Apply attention backend if specified