diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 945f83caa..d4b0c9177 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -148,7 +148,7 @@ The following tables show which models support each feature: | **Wan2.2-S2V** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ | | **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | -| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py index 3930677d3..9943d092b 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py @@ -72,6 +72,31 @@ class TestPipelineIndependence: assert issubclass(LTX23Pipeline, CFGParallelMixin) + def test_ltx23_pipeline_declares_offload_components(self): + """LTX23Pipeline must expose LTX-2.3-specific modules to offload discovery.""" + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery + + pipe = object.__new__(LTX23Pipeline) + torch.nn.Module.__init__(pipe) + pipe.transformer = torch.nn.Linear(1, 1) + pipe.text_encoder = torch.nn.Linear(1, 1) + pipe.connectors = torch.nn.Linear(1, 1) + pipe.vae = torch.nn.Linear(1, 1) + pipe.audio_vae = torch.nn.Linear(1, 1) + pipe.vocoder = torch.nn.Linear(1, 1) + + modules = ModuleDiscovery.discover(pipe) + + assert LTX23Pipeline._dit_modules == ["transformer"] + assert LTX23Pipeline._encoder_modules == ["text_encoder", "connectors"] + assert LTX23Pipeline._vae_modules == ["vae", "audio_vae"] + assert LTX23Pipeline._resident_modules == ["vocoder"] + assert modules.dit_names == ["transformer"] + assert modules.encoder_names == ["text_encoder", "connectors"] + assert modules.resident_names == ["vocoder"] + assert len(modules.vaes) == 2 + class TestCFGParallelHelpers: """Test LTX-2.3 CFG helper math without loading model weights.""" diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py index d28ff69d9..caa40c93f 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py @@ -10,7 +10,6 @@ This pipeline does NOT inherit from LTX2Pipeline because: versus LTX-2's per_layer_masked_mean_norm + shared projection path - LTX-2.3 uses a BWE vocoder outputting 48kHz audio (not 16kHz) - LTX-2.3 transformer requires the sigma parameter for prompt modulation -- CPU offloading is required for the 22B transformer (~44GB VRAM) """ from __future__ import annotations @@ -20,7 +19,7 @@ import json import os from collections.abc import Iterable from contextlib import nullcontext -from typing import Any +from typing import Any, ClassVar import numpy as np import torch @@ -46,7 +45,9 @@ from vllm_omni.diffusion.distributed.parallel_state import ( 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.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery from vllm_omni.diffusion.request import OmniDiffusionRequest from .pipeline_ltx2 import ( @@ -112,7 +113,7 @@ def get_ltx2_post_process_func(od_config: OmniDiffusionConfig): return post_process_func -class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): +class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsComponentDiscovery): """Fully independent LTX-2.3 pipeline. Key differences from LTX2Pipeline: @@ -120,11 +121,14 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): - Connectors: uses padding_side API (not additive_mask) - Vocoder: uses LTX2VocoderWithBWE (48kHz output) - Transformer: passes sigma for prompt_adaln - - CPU offloading: text encoder, connectors, VAE, vocoder stay on CPU """ # Audio is diffused jointly with video; warmup must size audio tokens. dummy_run_num_frames = 2 + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "connectors"] + _vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"] + _resident_modules: ClassVar[list[str]] = ["vocoder"] def __init__( self, @@ -166,7 +170,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): # --- Tokenizer (lightweight, stays wherever) --- self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) - # --- Text encoder: load on CPU, move to GPU only during encoding --- + # --- Text encoder --- with torch.device("cpu"): self.text_encoder = from_pretrained_with_prefetch( Gemma3ForConditionalGeneration.from_pretrained, @@ -177,7 +181,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): torch_dtype=dtype, ) - # --- Connectors: CPU (LTX-2.3 connectors include caption projection) --- + # --- Connectors (LTX-2.3 connectors include caption projection) --- self.connectors = from_pretrained_with_prefetch( LTX2TextConnectors.from_pretrained, model, @@ -187,7 +191,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): torch_dtype=dtype, ) - # --- VAE, Audio VAE: CPU --- + # --- VAE, Audio VAE --- self.vae = from_pretrained_with_prefetch( AutoencoderKLLTX2Video.from_pretrained, model, @@ -220,6 +224,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): transformer_config = load_transformer_config(model, "transformer", local_files_only) quant_config = getattr(self.od_config, "quantization_config", None) self.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config) + self._place_aux_components() # --- Scheduler --- self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( @@ -261,6 +266,20 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): self._num_timesteps = None self._current_timestep = None + def _place_aux_components(self) -> None: + parallel_config = getattr(self.od_config, "parallel_config", None) + use_managed_placement = bool( + getattr(self.od_config, "enable_cpu_offload", False) + or getattr(self.od_config, "enable_layerwise_offload", False) + or getattr(parallel_config, "use_hsdp", False) + ) + if use_managed_placement: + return + + modules = ModuleDiscovery.discover(self) + for module in (*modules.encoders, *modules.vaes, *modules.resident_modules): + module.to(self.device) + # ------------------------------------------------------------------ # Text Encoding (LTX-2.3 specific) # ------------------------------------------------------------------ @@ -303,16 +322,11 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): text_input_ids = text_inputs.input_ids.to(device) prompt_attention_mask = text_inputs.attention_mask.to(device) - # Move text encoder to GPU for encoding - self.text_encoder.to(device) text_encoder_outputs = self.text_encoder( input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True, ) - # Move text encoder back to CPU immediately - self.text_encoder.to("cpu") - torch.accelerator.empty_cache() hidden_states = text_encoder_outputs.hidden_states @@ -958,14 +972,10 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) - self.connectors.to(device) tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left") connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors( prompt_embeds, prompt_attention_mask, padding_side=tokenizer_padding_side ) - self.connectors.to("cpu") - if torch.cuda.is_available(): - torch.accelerator.empty_cache() positive_connector_prompt_embeds = connector_prompt_embeds positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds @@ -1241,22 +1251,14 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): ] latents = (1 - decode_noise_scale_t) * latents + decode_noise_scale_t * noise - # Move VAE, audio_vae, vocoder to GPU for decoding - self.vae.to(device) latents = latents.to(self.vae.dtype) video = self.vae.decode(latents, timestep_decode, return_dict=False)[0] video = self.video_processor.postprocess_video(video, output_type=output_type) - self.vae.to("cpu") - self.audio_vae.to(device) audio_latents = audio_latents.to(self.audio_vae.dtype) generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0] - self.audio_vae.to("cpu") - self.vocoder.to(device) audio = self.vocoder(generated_mel_spectrograms) - self.vocoder.to("cpu") - torch.accelerator.empty_cache() return DiffusionOutput(output=(video, audio))