[feat]: General diffusers adapter backend to run diffusion models (#2724)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>

vllm_omni/diffusion/data.py

@@ -452,7 +452,14 @@ class OmniDiffusionConfig:
     custom_pipeline_args: dict[str, Any] | None = None
 
     # Diffusion model loading format
-    diffusion_load_format: str = "default"  # "default", "custom_pipeline", "dummy"
+    # "default", "custom_pipeline", "dummy", "diffusers" (HF diffusers adapter)
+    diffusion_load_format: str = "default"
+
+    # Diffusers adapter kwargs
+    # kwargs forwarded to DiffusionPipeline.from_pretrained()
+    diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict)
+    # kwargs forwarded to pipeline.__call__()
+    diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict)
 
     # http server endpoint config, would be ignored in local mode
     host: str | None = None
@@ -648,6 +655,12 @@ class OmniDiffusionConfig:
         elif self.max_cpu_loras < 1:
             raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA")
 
+        if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):
+            raise ValueError(
+                "diffusers_load_kwargs and diffusers_call_kwargs are only "
+                "valid together with diffusion_load_format=diffusers"
+            )
+
     def set_tf_model_config(self, tf_config: "TransformerConfig") -> None:
         """Assign `tf_model_config` and propagate quantization if detected.
 
@@ -686,6 +699,10 @@ class OmniDiffusionConfig:
         """
         from vllm.transformers_utils.config import get_hf_file_to_dict
 
+        # Default model_class_name for diffusers adapter
+        if self.model_class_name is None and self.diffusion_load_format == "diffusers":
+            self.model_class_name = "DiffusersAdapterPipeline"
+
         try:
             config_dict = get_hf_file_to_dict("model_index.json", self.model)
             if config_dict is not None:
@@ -693,32 +710,42 @@ class OmniDiffusionConfig:
                     self.model_class_name = config_dict.get("_class_name", None)
                 self.update_multimodal_support()
 
-                tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model)
-                self.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
+                # Skip transformer config loading for diffusers adapter
+                # (non-DiT models don't have a separate transformer folder/config)
+                if self.diffusion_load_format == "diffusers":
+                    self.tf_model_config = TransformerConfig()
+                else:
+                    tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model)
+                    self.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
             else:
                 raise FileNotFoundError("model_index.json not found")
         except (AttributeError, OSError, ValueError, FileNotFoundError):
-            cfg = get_hf_file_to_dict("config.json", self.model)
-            if cfg is None:
-                raise ValueError(f"Could not find config.json or model_index.json for model {self.model}")
-
-            self.tf_model_config = TransformerConfig.from_dict(cfg)
-            model_type = cfg.get("model_type")
-            architectures = cfg.get("architectures") or []
-
-            if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
-                self.model_class_name = "BagelPipeline"
+            # Skip transformer config loading for diffusers adapter
+            # (non-DiT models don't have a separate transformer folder/config)
+            if self.diffusion_load_format == "diffusers":
                 self.tf_model_config = TransformerConfig()
-                self.update_multimodal_support()
-            elif model_type == "nextstep":
-                if self.model_class_name is None:
-                    self.model_class_name = "NextStep11Pipeline"
-                self.tf_model_config = TransformerConfig()
-                self.update_multimodal_support()
-            elif architectures and len(architectures) == 1:
-                self.model_class_name = architectures[0]
             else:
-                raise
+                cfg = get_hf_file_to_dict("config.json", self.model)
+                if cfg is None:
+                    raise ValueError(f"Could not find config.json or model_index.json for model {self.model}")
+
+                self.tf_model_config = TransformerConfig.from_dict(cfg)
+                model_type = cfg.get("model_type")
+                architectures = cfg.get("architectures") or []
+
+                if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
+                    self.model_class_name = "BagelPipeline"
+                    self.tf_model_config = TransformerConfig()
+                    self.update_multimodal_support()
+                elif model_type == "nextstep":
+                    if self.model_class_name is None:
+                        self.model_class_name = "NextStep11Pipeline"
+                    self.tf_model_config = TransformerConfig()
+                    self.update_multimodal_support()
+                elif architectures and len(architectures) == 1:
+                    self.model_class_name = architectures[0]
+                else:
+                    raise
 
     @classmethod
     def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
@@ -743,6 +770,12 @@ class OmniDiffusionConfig:
             cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER")
             kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none"
 
+        # Falsy-value check for not-None fields (convert potential None values in YAML config to empty containers)
+        if "diffusers_load_kwargs" in kwargs and kwargs["diffusers_load_kwargs"] is None:
+            kwargs["diffusers_load_kwargs"] = {}
+        if "diffusers_call_kwargs" in kwargs and kwargs["diffusers_call_kwargs"] is None:
+            kwargs["diffusers_call_kwargs"] = {}
+
         # Filter kwargs to only include valid fields
         valid_fields = {f.name for f in fields(cls)}
         filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -32,6 +32,7 @@ from vllm.utils.torch_utils import set_default_torch_dtype
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
 from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
+from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
 from vllm_omni.diffusion.registry import initialize_model
 
 if TYPE_CHECKING:
@@ -257,11 +258,14 @@ class DiffusersPipelineLoader:
         self,
         od_config: OmniDiffusionConfig,
         load_device: str,
-        load_format: str = "default",
+        load_format: str | None = "default",
         custom_pipeline_name: str | None = None,
         device: torch.device | None = None,
     ) -> nn.Module:
         """Load a model with the given configurations."""
+        if load_format is None:
+            load_format = "default"
+
         # CPU offload + FP8: load weights on device for FP8 quantization
         if load_device == "cpu" and od_config.quantization_config is not None:
             load_device = device.type
@@ -277,11 +281,21 @@ class DiffusersPipelineLoader:
                 with target_device:
                     if load_format == "default":
                         model = initialize_model(od_config)
+                    elif load_format == "diffusers":
+                        model = DiffusersAdapterPipeline(od_config=od_config, device=target_device)
                     elif load_format == "custom_pipeline":
                         model_cls = resolve_obj_by_qualname(custom_pipeline_name)
                         model = model_cls(od_config=od_config)
+                    else:
+                        # 'dummy' format should not call this function at all
+                        raise ValueError(f"Unknown load_format: {load_format}")
                 logger.debug("Loading weights on %s ...", load_device)
-                if self._is_gguf_quantization(od_config):
+                if load_format == "diffusers":
+                    # DiffusersAdapterPipeline.load_weights() calls
+                    # DiffusionPipeline.from_pretrained() internally — it does
+                    # NOT use our native (customized) pipeline classes.
+                    cast(DiffusersAdapterPipeline, model).load_weights()
+                elif self._is_gguf_quantization(od_config):
                     self._load_weights_with_gguf(model, od_config)
                 else:
                     # Quantization does not happen in `load_weights` but after it

vllm_omni/diffusion/models/diffusers_adapter/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Diffusers backend adapter for vLLM-Omni."""
+
+from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import (
+    DiffusersAdapterPipeline,
+)
+
+__all__ = [
+    "DiffusersAdapterPipeline",
+]

vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py

@@ -0,0 +1,358 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Diffusers backend adapter for vLLM-Omni.
+
+Provides a black-box wrapper around any 🤗 Diffusers pipeline, enabling
+vLLM-Omni to directly serve Diffusers models with near-zero per-model code.
+
+The adapter delegates full pipeline execution to diffusers' ``__call__()``.
+It does NOT support:
+- CFG parallel (diffusers handles CFG via guidance_scale internally)
+- Sequence parallel (requires model-specific attention surgery)
+- TeaCache / Cache-DiT (requires hooking into transformer blocks)
+- Step-wise execution (continuous batching)
+"""
+
+import logging
+import os
+from typing import Any
+
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from torch import nn
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniPromptType
+from vllm_omni.platforms import current_omni_platform
+
+logger = logging.getLogger(__name__)
+
+
+class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
+    """Black-box adapter that delegates full pipeline execution to a diffusers pipeline.
+
+    Usage::
+
+        adapter = DiffusersAdapterPipeline(od_config=od_config)
+        adapter.load_weights()  # calls DiffusionPipeline.from_pretrained()
+        output = adapter.forward(req)
+
+    Step-wise execution is explicitly rejected — diffusers encapsulates the
+    full denoising loop internally. Use native pipelines for continuous
+    batching mode.
+    """
+
+    supports_step_execution: bool = False
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, device: torch.device | None = None):
+        super().__init__()
+        self._pipeline: DiffusionPipeline
+        self.od_config = od_config
+        self.device = device
+        self._capabilities: dict[str, Any] = {}
+        self._raise_unsupported_features()
+
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler,
+            profiler_targets=["forward"],
+        )
+        if od_config.enable_diffusion_pipeline_profiler:
+            logger.info("Profiling enabled for DiffusersAdapterPipeline. Only 'forward' is supported.")
+
+    # ------------------------------------------------------------------
+    # Weight loading
+    # ------------------------------------------------------------------
+
+    def load_weights(self) -> None:
+        """Load the diffusers pipeline via ``DiffusionPipeline.from_pretrained()``."""
+
+        model_id = self.od_config.model
+        dtype = self.od_config.dtype
+
+        load_kwargs = {
+            "torch_dtype": dtype,
+            **self.od_config.diffusers_load_kwargs,
+        }
+        logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}")
+
+        self._pipeline = DiffusionPipeline.from_pretrained(
+            model_id,
+            **load_kwargs,
+        ).to(self.device)
+
+        # CPU offloading
+        if self.od_config.enable_layerwise_offload:
+            self._pipeline.enable_sequential_cpu_offload()
+        elif self.od_config.enable_cpu_offload:
+            self._pipeline.enable_model_cpu_offload()
+
+        # VAE slicing and tiling: try-catch because not all models have VAE
+        if self.od_config.vae_use_slicing:
+            try:
+                self._pipeline.enable_vae_slicing()
+            except Exception as e:
+                logger.warning(
+                    f"Failed to enable VAE slicing for diffusers pipeline {self._pipeline.__class__.__name__}: {e}"
+                )
+        if self.od_config.vae_use_tiling:
+            try:
+                self._pipeline.enable_vae_tiling()
+            except Exception as e:
+                logger.warning(
+                    f"Failed to enable VAE tiling for diffusers pipeline {self._pipeline.__class__.__name__}: {e}"
+                )
+
+        # Attention backend
+        self._set_attention_backend()
+
+    # ------------------------------------------------------------------
+    # Step-wise execution — explicitly rejected
+    # ------------------------------------------------------------------
+
+    def prepare_encode(self, **_: Any) -> Any:
+        raise NotImplementedError(
+            "Step-wise execution is not yet supported with the diffusers backend. "
+            "Use a native pipeline for continuous batching mode."
+        )
+
+    def denoise_step(self, **_: Any) -> torch.Tensor | None:
+        raise NotImplementedError(
+            "Step-wise execution is not yet supported with the diffusers backend. "
+            "Use a native pipeline for continuous batching mode."
+        )
+
+    def step_scheduler(self, **_: Any) -> None:
+        raise NotImplementedError(
+            "Step-wise execution is not yet supported with the diffusers backend. "
+            "Use a native pipeline for continuous batching mode."
+        )
+
+    def post_decode(self, **_: Any) -> Any:
+        raise NotImplementedError(
+            "Step-wise execution is not yet supported with the diffusers backend. "
+            "Use a native pipeline for continuous batching mode."
+        )
+
+    # ------------------------------------------------------------------
+    # Forward pass
+    # ------------------------------------------------------------------
+
+    def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+        """Full delegation to diffusers ``pipeline.__call__()``."""
+
+        kwargs = self._build_call_kwargs(req)
+        logger.debug(f"Calling diffusers pipeline with kwargs: {kwargs}")
+
+        with torch.inference_mode():
+            output = self._pipeline(**kwargs)  # pyright: ignore[reportCallIssue]
+
+        return self._wrap_output(output)
+
+    # ------------------------------------------------------------------
+    # Validation guards
+    # ------------------------------------------------------------------
+
+    def _raise_unsupported_features(self) -> None:
+        """Raise an error for incompatible feature switches."""
+        pc = self.od_config.parallel_config
+        if pc.cfg_parallel_size > 1:
+            raise NotImplementedError(
+                "CFG parallel is not supported with the diffusers backend. "
+                "Diffusers handles CFG internally via guidance_scale."
+            )
+        if pc.sequence_parallel_size is not None and pc.sequence_parallel_size > 1:
+            raise NotImplementedError(
+                "Sequence parallel is not supported with the diffusers backend. "
+                "It requires model-specific attention surgery."
+            )
+        if self.od_config.cache_backend not in ("none", None):
+            raise NotImplementedError(
+                f"Cache backend '{self.od_config.cache_backend}' is not supported "
+                "with the diffusers backend. TeaCache/Cache-DiT require hooking "
+                "into individual transformer blocks."
+            )
+        if self.od_config.enforce_eager:
+            raise NotImplementedError(
+                "Eager execution is not supported with the diffusers backend. "
+                "Use a native pipeline for continuous batching mode."
+            )
+        if self.od_config.quantization_config is not None:
+            raise NotImplementedError(
+                "Quantization is not supported with the diffusers backend. Use a native pipeline for quantization."
+            )
+
+    # ------------------------------------------------------------------
+    # Wrap settings, inputs, and outputs
+    # ------------------------------------------------------------------
+
+    def _set_attention_backend(self) -> None:
+        """Set the attention backend.
+
+        Roughly follow the logic in vllm_omni/diffusion/attention/backends/utils/fa.py,
+        But also consider the available attention backends in diffusers.
+        (See: https://huggingface.co/docs/diffusers/optimization/attention_backends)
+        """
+        if not hasattr(self._pipeline, "transformer"):
+            logging.info("No transformer found in diffusers pipeline. Skipping attention backend setting.")
+            return
+
+        attention_backend_config = self.od_config.attention_backend or os.environ.get("DIFFUSION_ATTENTION_BACKEND")
+        attention_backend_attempts: list[str] = []
+        match attention_backend_config:
+            case "FLASH_ATTN" | None:
+                if current_omni_platform.is_rocm():
+                    attention_backend_attempts.append("aiter")
+                elif current_omni_platform.is_xpu():
+                    attention_backend_attempts.append("_native_xla")
+                elif current_omni_platform.is_musa():
+                    logger.warning(
+                        "Unknown diffusers attention backend option for MUSA platform. Falling back to SDPA."
+                    )
+                    attention_backend_attempts.append("native")
+                else:
+                    attention_backend_attempts.extend(
+                        [
+                            "_flash_3_hub",
+                            "_flash_3_varlen_hub",
+                            "_flash_3",
+                            "_flash_varlen_3",
+                            "flash_hub",
+                            "flash_varlen_hub",
+                            "flash",
+                            "flash_varlen",
+                            "_native_flash",
+                        ]
+                    )
+            case "SAGE_ATTN":
+                attention_backend_attempts.extend(["sage_hub", "sage", "sage", "sage_varlen"])
+            case "ASCEND":
+                attention_backend_attempts.append("_native_npu")
+            case "TORCH_SDPA":
+                attention_backend_attempts.append("native")
+            case _:
+                logger.warning(f"Invalid attention backend: {attention_backend_config}. Falling back to SDPA.")
+                attention_backend_attempts.append("native")
+
+        attempt_errors: list[str] = []
+        set_backend: str | None = None
+        for backend in attention_backend_attempts:
+            try:
+                self._pipeline.transformer.set_attention_backend(backend)
+                set_backend = backend
+                break
+            except Exception as e:
+                attempt_errors.append(str(e))
+
+        # If all attempts fail, fallback to SDPA and warn the user about the failures
+        if len(attempt_errors) == len(attention_backend_attempts):
+            self._pipeline.transformer.set_attention_backend("native")
+            logger.warning(
+                f"Failed to set attention backend '{attention_backend_config}' for "
+                f"diffusers pipeline {self._pipeline.__class__.__name__}. "
+                "Falling back to SDPA. "
+                f"The following attempts were made: {dict(zip(attention_backend_attempts, attempt_errors))}"
+            )
+            return
+
+        # If some attempts fail, only warn the user about the failures
+        logger.info(
+            f"Set diffusers attention backend to '{set_backend}', adapted from "
+            f"user config value '{attention_backend_config}'."
+        )
+        if len(attempt_errors) > 0:
+            logger.warning(
+                f"The following failed attempts were made before choosing this diffusers backend: "
+                f"{dict(zip(attention_backend_attempts, attempt_errors))}"
+            )
+
+    def _build_call_kwargs(self, req: OmniDiffusionRequest) -> dict[str, Any]:
+        """Translate ``OmniDiffusionRequest`` into diffusers ``__call__`` kwargs."""
+        sampling = req.sampling_params
+        prompt, neg_prompt = self._extract_prompt(req.prompts)
+
+        # Merge user-provided call kwargs from stage/CLI defaults.
+        # Request-time parameters take precedence over stage-config defaults
+        call_kwargs = self.od_config.diffusers_call_kwargs
+        kwargs: dict[str, Any] = {
+            **call_kwargs,
+            "prompt": prompt,
+            "num_inference_steps": sampling.num_inference_steps,
+            "guidance_scale": sampling.guidance_scale,
+            "output_type": sampling.output_type or self.od_config.output_type,
+        }
+
+        if sampling.height is not None:
+            kwargs["height"] = sampling.height
+        if sampling.width is not None:
+            kwargs["width"] = sampling.width
+        if sampling.num_frames is not None and sampling.num_frames > 1:
+            kwargs["num_frames"] = sampling.num_frames
+        if sampling.num_outputs_per_prompt is not None and sampling.num_outputs_per_prompt > 1:
+            kwargs["num_images_per_prompt"] = sampling.num_outputs_per_prompt
+
+        if neg_prompt is not None:
+            kwargs["negative_prompt"] = neg_prompt
+
+        if sampling.generator is not None:
+            kwargs["generator"] = sampling.generator
+        elif sampling.seed is not None:
+            kwargs["generator"] = torch.Generator(device=sampling.generator_device).manual_seed(sampling.seed)
+        else:
+            kwargs["generator"] = torch.Generator(device=sampling.generator_device)
+
+        if sampling.latents is not None:
+            kwargs["latents"] = sampling.latents
+
+        return kwargs
+
+    @staticmethod
+    def _extract_prompt(prompt_obj: list[OmniPromptType]) -> tuple[str | list[str], str | list[str] | None]:
+        """Extract the text prompts and negative prompts from a list of prompt objects."""
+        if len(prompt_obj) == 1:
+            if isinstance(prompt_obj[0], str):
+                return prompt_obj[0], None
+            else:
+                return prompt_obj[0].get("prompt", ""), prompt_obj[0].get("negative_prompt", None)
+
+        prompts = []
+        negative_prompts: list[str] | None = []
+        for prompt in prompt_obj:
+            if isinstance(prompt, str):
+                prompts.append(prompt)
+            else:
+                prompts.append(prompt.get("prompt", ""))
+                negative_prompts.append(prompt.get("negative_prompt", ""))
+        if all(not np for np in negative_prompts):
+            negative_prompts = None
+        return prompts, negative_prompts
+
+    @staticmethod
+    def _extract_negative_prompt(prompt_obj: Any) -> str | None:
+        """Extract the negative prompt from a prompt object, if present."""
+        if isinstance(prompt_obj, dict):
+            return prompt_obj.get("negative_prompt")
+        return getattr(prompt_obj, "negative_prompt", None)
+
+    def _wrap_output(self, output: Any) -> DiffusionOutput:
+        """Convert diffusers pipeline output to ``DiffusionOutput``.
+
+        Diffusers output types:
+        - ``ImagePipelineOutput(images=...)`` — text2img, img2img
+        - ``VideoPipelineOutput(frames=...)`` — text2vid, img2vid
+        """
+        from vllm_omni.diffusion.data import DiffusionOutput
+
+        if hasattr(output, "images"):
+            # Preserve diffusers image format (`output_type`)
+            return DiffusionOutput(output=output.images)
+
+        if hasattr(output, "frames"):
+            # Preserve diffusers video format (`output_type`)
+            return DiffusionOutput(output=output.frames)
+
+        if hasattr(output, "audios"):
+            return DiffusionOutput(output=output.audios)
+
+        return DiffusionOutput(output=output)

vllm_omni/diffusion/registry.py

@@ -208,6 +208,11 @@ _DIFFUSION_MODELS = {
         "pipeline_omnivoice",
         "OmniVoicePipeline",
     ),
+    "DiffusersAdapterPipeline": (
+        "diffusers_adapter",
+        "pipeline_diffusers_adapter",
+        "DiffusersAdapterPipeline",
+    ),
 }