diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index a59e9a524..3c2322940 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -265,44 +265,6 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): self._num_timesteps = None self._current_timestep = None - @staticmethod - def _pack_text_embeds( - text_hidden_states: torch.Tensor, - sequence_lengths: torch.Tensor, - device: str | torch.device, - padding_side: str = "left", - scale_factor: int = 8, - eps: float = 1e-6, - ) -> torch.Tensor: - batch_size, seq_len, hidden_dim, num_layers = text_hidden_states.shape - original_dtype = text_hidden_states.dtype - - token_indices = torch.arange(seq_len, device=device).unsqueeze(0) - if padding_side == "right": - mask = token_indices < sequence_lengths[:, None] - elif padding_side == "left": - start_indices = seq_len - sequence_lengths[:, None] - mask = token_indices >= start_indices - else: - raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}") - mask = mask[:, :, None, None] - - masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0) - num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1) - masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / (num_valid_positions + eps) - - x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True) - x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True) - - normalized_hidden_states = (text_hidden_states - masked_mean) / (x_max - x_min + eps) - normalized_hidden_states = normalized_hidden_states * scale_factor - - normalized_hidden_states = normalized_hidden_states.flatten(2) - mask_flat = mask.squeeze(-1).expand(-1, -1, hidden_dim * num_layers) - normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0) - normalized_hidden_states = normalized_hidden_states.to(dtype=original_dtype) - return normalized_hidden_states - def _get_gemma_prompt_embeds( self, prompt: str | list[str], @@ -342,16 +304,9 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): ) text_encoder_hidden_states = text_encoder_outputs.hidden_states text_encoder_hidden_states = torch.stack(text_encoder_hidden_states, dim=-1) - sequence_lengths = prompt_attention_mask.sum(dim=-1) - - prompt_embeds = self._pack_text_embeds( - text_encoder_hidden_states, - sequence_lengths, - device=device, - padding_side=self.tokenizer.padding_side, - scale_factor=scale_factor, - ) - prompt_embeds = prompt_embeds.to(dtype=dtype) + # `diffusers>=0.38` moved per_layer_masked_mean_norm inside the connector, + # so pack to 3D here and let the connector handle normalization. + prompt_embeds = text_encoder_hidden_states.flatten(2, 3).to(dtype=dtype) _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) @@ -893,9 +848,9 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): device=device, ) # Compute positive prompt connectors - additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0 + tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left") connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors( - prompt_embeds, additive_attention_mask, additive_mask=True + prompt_embeds, prompt_attention_mask, padding_side=tokenizer_padding_side ) # Compute negative prompt connectors when CFG is enabled @@ -903,17 +858,14 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): negative_connector_audio_prompt_embeds = None negative_connector_attention_mask = None if self.do_classifier_free_guidance: - negative_additive_attention_mask = ( - 1 - negative_prompt_attention_mask.to(negative_prompt_embeds.dtype) - ) * -1000000.0 ( negative_connector_prompt_embeds, negative_connector_audio_prompt_embeds, negative_connector_attention_mask, ) = self.connectors( negative_prompt_embeds, - negative_additive_attention_mask, - additive_mask=True, + negative_prompt_attention_mask, + padding_side=tokenizer_padding_side, ) latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py index c3bbf3aa0..84ef29305 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py @@ -5,9 +5,9 @@ Fully independent LTX-2.3 pipeline for vLLM-Omni. This pipeline does NOT inherit from LTX2Pipeline because: -- LTX-2.3 uses a different text encoding strategy (flatten ALL 49 hidden states - vs. LTX-2's _pack_text_embeds with per-layer normalization and pooling) -- LTX-2.3 connectors expect the padding_side API (not additive_mask) +- LTX-2.3 connectors run per_token_rms_norm + per-modality video/audio + projection internally (per_modality_projections=True), + 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) @@ -230,10 +230,10 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin): ): """Encode prompts using Gemma-3-12B, returning ALL 49 hidden states flattened. - LTX-2.3 differs from LTX-2 in text encoding: - - LTX-2: uses _pack_text_embeds (layer selection + pooling) - - LTX-2.3: stacks ALL 49 hidden states and flattens to [B, seq, 188160] - The connectors unflatten, apply per_token_rms_norm, and project internally. + Stacks all 49 hidden states and flattens to [B, seq, hidden * 49]. The + connectors unflatten, apply per_token_rms_norm, and project internally + (same shape contract as LTX-2 since the `diffusers==0.38` connector + migration; the two differ only in the connector's internal norm path). """ device = device or self.device dtype = dtype or self.text_encoder.dtype diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py index 057c144e9..af4135d65 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py @@ -460,9 +460,9 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline): device=device, ) # Compute positive prompt connectors - additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0 + tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left") connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors( - prompt_embeds, additive_attention_mask, additive_mask=True + prompt_embeds, prompt_attention_mask, padding_side=tokenizer_padding_side ) # Compute negative prompt connectors when CFG is enabled @@ -470,17 +470,14 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline): negative_connector_audio_prompt_embeds = None negative_connector_attention_mask = None if self.do_classifier_free_guidance: - negative_additive_attention_mask = ( - 1 - negative_prompt_attention_mask.to(negative_prompt_embeds.dtype) - ) * -1000000.0 ( negative_connector_prompt_embeds, negative_connector_audio_prompt_embeds, negative_connector_attention_mask, ) = self.connectors( negative_prompt_embeds, - negative_additive_attention_mask, - additive_mask=True, + negative_prompt_attention_mask, + padding_side=tokenizer_padding_side, ) latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1