[Model] Add LTX-2.5 Diffusers backend support

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py

@@ -23,7 +23,14 @@ 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 (
+    LTX25_DIFFUSERS_COMMIT,
+    LTX25_DIFFUSERS_MODEL_ID,
+    LTX25_ORIGINAL_MODEL_ID,
+    BasePipelineUtils,
+    get_pipeline_utils,
+    is_ltx25_diffusers_model,
+)
 from vllm_omni.diffusion.models.diffusers_adapter.quantization_utils import (
     apply_diffusers_quantization_config,
     convert_diffusers_quantization_config,
@@ -67,6 +74,23 @@ _DIFFUSERS_CONFIG_LOAD_KWARGS = {
 }
 
 
+def _supports_ltx25_diffusers() -> bool:
+    """Check the standard-pipeline capabilities added with LTX-2.5."""
+    try:
+        import diffusers
+        import transformers
+
+        pipeline_class = diffusers.LTX2Pipeline
+        init_parameters = inspect.signature(pipeline_class.__init__).parameters
+        return (
+            "duration_head" in init_parameters
+            and "prompt_enhancer" in init_parameters
+            and hasattr(transformers, "Gemma4UnifiedForConditionalGeneration")
+        )
+    except (AttributeError, ImportError, RuntimeError):
+        return False
+
+
 class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
     """Black-box adapter that delegates full pipeline execution to a diffusers pipeline.
 
@@ -110,6 +134,7 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
 
         model_id = self.od_config.model
         dtype = self.od_config.dtype
+        self._validate_model_compatibility()
 
         load_kwargs = {
             "torch_dtype": dtype,
@@ -130,6 +155,11 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
         logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}")
 
         self._pipeline = DiffusionPipeline.from_pretrained(model_id, **load_kwargs)
+        # Gated repositories may be loadable with an explicit token in
+        # diffusers_load_kwargs even when config enrichment could not resolve
+        # their pipeline class. In that case the loaded class is authoritative.
+        if pipeline_class_name is None:
+            self._pipeline_utils = get_pipeline_utils(self._pipeline.__class__.__name__)
         self._pipeline_utils.apply_post_load_updates(self._pipeline, self.od_config)
 
         self._pipeline.to(self.device)
@@ -146,14 +176,14 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
         # 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()
+                self._enable_vae_optimization("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()
+                self._enable_vae_optimization("tiling")
             except Exception as e:
                 logger.warning(
                     f"Failed to enable VAE tiling for diffusers pipeline {self._pipeline.__class__.__name__}: {e}"
@@ -162,6 +192,18 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
         # Attention backend
         self._set_attention_backend()
 
+    def _enable_vae_optimization(self, optimization: str) -> None:
+        """Enable a VAE optimization through the pipeline or VAE component API."""
+        pipeline_method = getattr(self._pipeline, f"enable_vae_{optimization}", None)
+        if callable(pipeline_method):
+            pipeline_method()
+            return
+
+        vae_method = getattr(getattr(self._pipeline, "vae", None), f"enable_{optimization}", None)
+        if not callable(vae_method):
+            raise AttributeError(f"Neither the pipeline nor its VAE supports {optimization}")
+        vae_method()
+
     # ------------------------------------------------------------------
     # Step-wise execution — explicitly rejected
     # ------------------------------------------------------------------
@@ -247,6 +289,23 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
                     "quantization conversion."
                 )
 
+    def _validate_model_compatibility(self) -> None:
+        normalized_model = str(self.od_config.model).rstrip("/\\").replace("\\", "/").lower()
+        if normalized_model == LTX25_ORIGINAL_MODEL_ID.lower():
+            raise ValueError(
+                f"{LTX25_ORIGINAL_MODEL_ID} is a component-weight repository without a Diffusers "
+                f"model_index.json. Use the official converted repository {LTX25_DIFFUSERS_MODEL_ID} "
+                "with --diffusion-load-format diffusers."
+            )
+
+        if is_ltx25_diffusers_model(self.od_config) and not _supports_ltx25_diffusers():
+            raise ImportError(
+                f"{LTX25_DIFFUSERS_MODEL_ID} requires unreleased LTX-2.5 support from Diffusers and "
+                "Gemma 4 Unified support from Transformers. Install the source-preview dependencies with "
+                "`python -m pip install --upgrade 'transformers>=5.10.1' "
+                f"'diffusers @ git+https://github.com/huggingface/diffusers.git@{LTX25_DIFFUSERS_COMMIT}'`."
+            )
+
     def _load_diffusers_component_names(
         self,
         model_id: str,
@@ -370,11 +429,22 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
                 )
 
         # Request-time sampling params
-        for key, value in sampling.__dict__.items():
+        for key, value in sampling.extra_args.items():
             if value is None:
                 continue
             if self._accept_call_kwargs is None or key in self._accept_call_kwargs:
                 kwargs[key] = value
+            else:
+                logger.warning(
+                    f"Skipping unsupported diffusers pipeline __call__ argument `{key}` from extra_args. "
+                    f"Check out the documentation of {self._pipeline.__class__.__name__}."
+                )
+
+        for key, value in sampling.__dict__.items():
+            if key == "extra_args" or value is None:
+                continue
+            if self._accept_call_kwargs is None or key in self._accept_call_kwargs:
+                kwargs[key] = value
 
         # Special format fields in sampling params
         if output_type := sampling.output_type or self.od_config.output_type:
@@ -393,6 +463,8 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
         else:
             kwargs["generator"] = torch.Generator(device=sampling.generator_device)
 
+        self._pipeline_utils.update_call_kwargs(self.od_config, kwargs)
+
         logger.info(
             "Calling diffusers pipeline with kwargs: %s", DiffusersAdapterPipeline._summarize_call_kwargs_value(kwargs)
         )
@@ -502,6 +574,21 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
             return DiffusionOutput(output=output.images)
 
         if hasattr(output, "frames"):
+            audio = getattr(output, "audio", None)
+            if audio is not None:
+                metadata: dict[str, Any] = {}
+                if (sample_rate := self._get_output_audio_sample_rate()) is not None:
+                    metadata["audio"] = {"sample_rate": sample_rate}
+                envelope: dict[str, Any] = {
+                    "payload": {
+                        "video": output.frames,
+                        "audio": audio,
+                    }
+                }
+                if metadata:
+                    envelope["metadata"] = metadata
+                return DiffusionOutput(output=envelope)
+
             # Preserve diffusers video format (`output_type`)
             return DiffusionOutput(output=output.frames)
 
@@ -510,6 +597,24 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
 
         return DiffusionOutput(output=output)
 
+    def _get_output_audio_sample_rate(self) -> int | None:
+        vocoder = getattr(self._pipeline, "vocoder", None)
+        config = getattr(vocoder, "config", None)
+        if config is None:
+            return None
+
+        if hasattr(config, "get"):
+            value = config.get("output_sampling_rate")
+        else:
+            value = getattr(config, "output_sampling_rate", None)
+        if isinstance(value, bool):
+            return None
+        try:
+            sample_rate = int(value)
+        except (TypeError, ValueError):
+            return None
+        return sample_rate if sample_rate > 0 else None
+
     @staticmethod
     def _summarize_call_kwargs_value(value: Any) -> Any:
         """Return a sanitized summary of diffusers call kwargs for logging."""