[Bugfix][XPU] Enable LTX-2 family on XPU (#7961)

Signed-off-by: Joshna Medisetty <joshna.medisetty@intel.com>

vllm_omni/diffusion/models/ltx2/ltx2_components.py

@@ -30,6 +30,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
 from vllm_omni.transformers_utils.repo_utils import hf_api
 
 if TYPE_CHECKING:
@@ -633,6 +634,13 @@ def _place_aux_components(pipeline: Any) -> None:
         module.to(pipeline.device)
 
 
+def _declare_text_encoder_offload_plan(pipeline: Any) -> None:
+    """Make the text encoder streamable; offloading the DiT alone leaves it resident."""
+    language_model = getattr(getattr(pipeline.text_encoder, "model", None), "language_model", None)
+    if language_model is not None and hasattr(language_model, "layers"):
+        pipeline._offload_plan = OffloadPlan(encoder_block_attrs={"text_encoder": ("model.language_model.layers",)})
+
+
 def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
     """Build the common LTX component graph selected by ``component_profile``."""
     profile: LTXComponentProfile = pipeline.component_profile
@@ -680,6 +688,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
         dtype=dtype,
         revision=revision,
     )
+    _declare_text_encoder_offload_plan(pipeline)
     _install_connector_attention(
         pipeline.connectors,
         preserve_learned_register_mask=profile.preserve_connector_attention_mask,

vllm_omni/diffusion/models/ltx2/ltx2_runtime.py

@@ -85,7 +85,7 @@ def _run_ltx_vocoder(vocoder: nn.Module, generated_mel: torch.Tensor) -> torch.T
 
         input_dtype = generated_mel.dtype
         module_dtype = next(vocoder.parameters()).dtype
-        if device_type == "mps":
+        if device_type in ("mps", "xpu", "cpu"):
             if module_dtype != torch.float32:
                 vocoder.float()
             try: