diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index 5826f273e..718756fe4 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/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, diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py index 9e682ffef..21f44a683 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py +++ b/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: