[Model] Add native SANA-Video 2B T2V and I2V support (#5508)

Signed-off-by: leo <92263211+napleon-liu@users.noreply.github.com>
Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
Signed-off-by: liuyao0322 <yaoliu548926@gmail.com>
Co-authored-by: leo <92263211+napleon-liu@users.noreply.github.com>

vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py

@@ -23,7 +23,11 @@ from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from torch import nn
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import BasePipelineUtils, get_pipeline_utils
+from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import (
+    BasePipelineUtils,
+    get_pipeline_utils_for_config,
+    resolve_diffusers_pipeline_class,
+)
 from vllm_omni.diffusion.models.diffusers_adapter.quantization_utils import (
     apply_diffusers_quantization_config,
     convert_diffusers_quantization_config,
@@ -117,9 +121,7 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
         }
         convert_diffusers_quantization_config(load_kwargs)
 
-        pipeline_class = self.od_config.diffusers_pipeline_cls
-        pipeline_class_name = pipeline_class.__name__ if pipeline_class is not None else None
-        self._pipeline_utils = get_pipeline_utils(pipeline_class_name)
+        self._pipeline_utils = get_pipeline_utils_for_config(self.od_config)
         self._pipeline_utils.update_load_kwargs(self.od_config, load_kwargs)
         component_names = (
             self._load_diffusers_component_names(model_id, load_kwargs)
@@ -129,7 +131,7 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
         apply_diffusers_quantization_config(self.od_config, load_kwargs, component_names)
         logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}")
 
-        self._pipeline = DiffusionPipeline.from_pretrained(model_id, **load_kwargs)
+        self._pipeline = self._load_pipeline_from_pretrained(model_id, load_kwargs)
         self._pipeline_utils.apply_post_load_updates(self._pipeline, self.od_config)
 
         self._pipeline.to(self.device)
@@ -162,6 +164,23 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
         # Attention backend
         self._set_attention_backend()
 
+    def _load_pipeline_from_pretrained(
+        self,
+        model_id: str,
+        load_kwargs: dict[str, Any],
+    ) -> DiffusionPipeline:
+        pipeline_class = resolve_diffusers_pipeline_class(self.od_config)
+        configured_pipeline_class = self.od_config.diffusers_pipeline_cls
+        # Preserve DiffusionPipeline's normal dynamic class resolution unless
+        # a model-specific adapter hook explicitly selected another class.
+        loader_class = (
+            pipeline_class
+            if pipeline_class is not None and pipeline_class is not configured_pipeline_class
+            else DiffusionPipeline
+        )
+        logger.info("Loading Diffusers pipeline class %s", loader_class.__name__)
+        return loader_class.from_pretrained(model_id, **load_kwargs)
+
     # ------------------------------------------------------------------
     # Step-wise execution — explicitly rejected
     # ------------------------------------------------------------------
@@ -376,6 +395,13 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
             if self._accept_call_kwargs is None or key in self._accept_call_kwargs:
                 kwargs[key] = value
 
+        self._pipeline_utils.update_call_kwargs(
+            req,
+            sampling,
+            self._accept_call_kwargs,
+            kwargs,
+        )
+
         # Special format fields in sampling params
         if output_type := sampling.output_type or self.od_config.output_type:
             kwargs["output_type"] = output_type

vllm_omni/diffusion/models/diffusers_adapter/pipeline_utils.py

@@ -1,17 +1,27 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from typing import Any
 
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.request import resolve_video_num_frames
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 
 class BasePipelineUtils:
     """No-op hooks for pipeline-specific diffusers adapter behavior."""
 
+    def resolve_pipeline_class(
+        self,
+        od_config: OmniDiffusionConfig,
+        configured_pipeline_class: type[DiffusionPipeline] | None,
+    ) -> type[DiffusionPipeline] | None:
+        """Resolve the concrete Diffusers class used to load a checkpoint."""
+        return configured_pipeline_class
+
     def update_load_kwargs(self, od_config: OmniDiffusionConfig, load_kwargs: dict[str, Any]) -> None:
         pass
 
@@ -21,6 +31,59 @@ class BasePipelineUtils:
     def validate_runtime_sampling_params(self, sampling: OmniDiffusionSamplingParams) -> None:
         pass
 
+    def update_call_kwargs(
+        self,
+        req: DiffusionRequestBatch,
+        sampling: OmniDiffusionSamplingParams,
+        accepted_call_kwargs: set[str] | None,
+        call_kwargs: dict[str, Any],
+    ) -> None:
+        pass
+
+
+class SanaVideoPipelineUtils(BasePipelineUtils):
+    def resolve_pipeline_class(
+        self,
+        od_config: OmniDiffusionConfig,
+        configured_pipeline_class: type[DiffusionPipeline] | None,
+    ) -> type[DiffusionPipeline] | None:
+        # Both released SANA-Video checkpoints declare SanaVideoPipeline in
+        # model_index.json. Honour an explicit Omni I2V class selection by
+        # loading the same components into Diffusers' I2V pipeline instead.
+        if od_config.model_class_name == "SanaImageToVideoPipeline":
+            from diffusers import SanaImageToVideoPipeline
+
+            return SanaImageToVideoPipeline
+        return configured_pipeline_class
+
+    def update_call_kwargs(
+        self,
+        req: DiffusionRequestBatch,
+        sampling: OmniDiffusionSamplingParams,
+        accepted_call_kwargs: set[str] | None,
+        call_kwargs: dict[str, Any],
+    ) -> None:
+        # Diffusers SANA-Video calls the shared request's `num_frames`
+        # parameter `frames`.
+        if (
+            accepted_call_kwargs is not None
+            and "frames" in accepted_call_kwargs
+            and "num_frames" not in accepted_call_kwargs
+        ):
+            configured_default = call_kwargs.get("frames", 81)
+            if configured_default is None:
+                configured_default = 81
+            call_kwargs["frames"] = resolve_video_num_frames(
+                sampling.num_frames,
+                default_num_frames=configured_default,
+                is_dummy_run=req.is_dummy_run(),
+            )
+
+        # SANA's startup dummy dimensions can map to a resolution bucket that
+        # conflicts with the upstream pipeline's divisibility validation.
+        if req.is_dummy_run() and accepted_call_kwargs is not None and "use_resolution_binning" in accepted_call_kwargs:
+            call_kwargs["use_resolution_binning"] = False
+
 
 class WanPipelineUtils(BasePipelineUtils):
     def update_load_kwargs(self, od_config: OmniDiffusionConfig, load_kwargs: dict[str, Any]) -> None:
@@ -49,6 +112,8 @@ class WanPipelineUtils(BasePipelineUtils):
 
 
 PIPELINE_UTILS_REGISTRY: dict[str, type[BasePipelineUtils]] = {
+    "SanaVideoPipeline": SanaVideoPipelineUtils,
+    "SanaImageToVideoPipeline": SanaVideoPipelineUtils,
     "WanPipeline": WanPipelineUtils,
     "WanImageToVideoPipeline": WanPipelineUtils,
     "WanVACEPipeline": WanPipelineUtils,
@@ -62,3 +127,20 @@ def get_pipeline_utils(pipeline_class_name: str | None) -> BasePipelineUtils:
         return BasePipelineUtils()
     pipeline_utils_cls = PIPELINE_UTILS_REGISTRY.get(pipeline_class_name, BasePipelineUtils)
     return pipeline_utils_cls()
+
+
+def get_pipeline_utils_for_config(od_config: OmniDiffusionConfig) -> BasePipelineUtils:
+    """Return model-specific adapter hooks for the requested or configured class."""
+    requested_class_name = od_config.model_class_name
+    if requested_class_name in PIPELINE_UTILS_REGISTRY:
+        return get_pipeline_utils(requested_class_name)
+
+    configured_pipeline_class = od_config.diffusers_pipeline_cls
+    configured_class_name = configured_pipeline_class.__name__ if configured_pipeline_class is not None else None
+    return get_pipeline_utils(configured_class_name)
+
+
+def resolve_diffusers_pipeline_class(od_config: OmniDiffusionConfig) -> type[DiffusionPipeline] | None:
+    """Resolve the concrete Diffusers pipeline used for loading and capability checks."""
+    pipeline_utils = get_pipeline_utils_for_config(od_config)
+    return pipeline_utils.resolve_pipeline_class(od_config, od_config.diffusers_pipeline_cls)