Update pipeline.py

pipeline.py

@@ -1,240 +1,1272 @@
-import spaces
-import gradio as gr
-import torch
-from diffusers import DiffusionPipeline
-from diffusers.utils import load_image, export_to_video
-import os
-import random
+# Copyright 2025 Lightricks and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
 import numpy as np
-from moviepy import ImageSequenceClip, AudioFileClip, VideoFileClip
-from PIL import Image, ImageOps
-
-# --- 1. Model Setup & Configuration ---
-
-# Define the specific distilled sigmas (from LTX-2 documentation)
-DISTILLED_SIGMA_VALUES = [
-    1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875
-]
-
-print("Loading LTX-2 Distilled Pipeline...")
-pipe = DiffusionPipeline.from_pretrained(
-    "rootonchair/LTX-2-19b-distilled",
-    custom_pipeline="multimodalart/ltx2-audio-to-video",
-    torch_dtype=torch.bfloat16
-)
-pipe.to("cuda")
-
-# --- Restored LoRA Loading ---
-# We load the Camera Control LoRA and fuse it.
-# This stabilizes the video (static camera), which is usually preferred for talking heads.
-print("Loading and Fusing Camera Control LoRA...")
-pipe.load_lora_weights("Lightricks/LTX-2-19b-LoRA-Camera-Control-Static", adapter_name="camera_control")
-# We fuse with a scale of 1.0. This merges the weights into the base model permanently for this session.
-pipe.fuse_lora(lora_scale=1.0)
-
-# --- 2. Helper Functions ---
-
-def save_video_with_audio(video_frames, audio_path, fps=24):
-    """
-    Combines the generated video frames with the original input audio.
-    """
-    output_filename = f"output_{random.randint(0, 100000)}.mp4"
-    
-    # 1. Handle Diffusers Output Formats
-    if isinstance(video_frames, list):
-        if video_frames and isinstance(video_frames[0], list):
-            frames_to_process = video_frames[0]
-        else:
-            frames_to_process = video_frames
-        np_frames = [np.array(img) for img in frames_to_process]
-        clip = ImageSequenceClip(np_frames, fps=fps)
-        
-    elif isinstance(video_frames, str):
-        clip = VideoFileClip(video_frames)
+import torch
+import torchaudio
+import torchaudio.transforms as T
+from transformers import Gemma3ForConditionalGeneration, GemmaTokenizer, GemmaTokenizerFast
+
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.image_processor import PipelineImageInput
+from diffusers.loaders import FromSingleFileMixin, LTXVideoLoraLoaderMixin
+from diffusers.models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video
+from diffusers.models.transformers import LTX2VideoTransformer3DModel
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.video_processor import VideoProcessor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors
+from diffusers.pipelines.ltx2.pipeline_output import LTX2PipelineOutput
+from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import DiffusionPipeline
+        >>> from diffusers.pipelines.ltx2.export_utils import encode_video
+        >>> from diffusers.utils import load_image
+
+        >>> pipe = DiffusionPipeline.from_pretrained(
+        ...     "Lightricks/LTX-2",
+        ...     custom_pipeline="pipeline_ltx2_audio2video",
+        ...     torch_dtype=torch.bfloat16
+        ... )
+        >>> pipe.enable_model_cpu_offload()
+
+        >>> image = load_image(
+        ...     "https://huggingface.co/datasets/a-r-r-o-w/tiny-meme-dataset-captioned/resolve/main/images/8.png"
+        ... )
+        >>> prompt = "A young girl stands calmly in the foreground, looking directly at the camera, as a house fire rages in the background."
+        >>> negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+        >>> frame_rate = 24.0
+        >>> video, audio = pipe(
+        ...     image=image,
+        ...     audio="path/to/audio.wav",
+        ...     prompt=prompt,
+        ...     negative_prompt=negative_prompt,
+        ...     width=768,
+        ...     height=512,
+        ...     num_frames=121,
+        ...     frame_rate=frame_rate,
+        ...     num_inference_steps=40,
+        ...     guidance_scale=4.0,
+        ...     output_type="np",
+        ...     return_dict=False,
+        ... )
+        >>> video = (video * 255).round().astype("uint8")
+        >>> video = torch.from_numpy(video)
+
+        >>> encode_video(
+        ...     video[0],
+        ...     fps=frame_rate,
+        ...     audio=audio[0].float().cpu(),
+        ...     audio_sample_rate=pipe.vocoder.config.output_sampling_rate,  # should be 24000
+        ...     output_path="video.mp4",
+        ... )
+        ```
+"""
+
+
+def retrieve_latents(
+    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
     else:
-        temp_path = "temp_video_no_audio.mp4"
-        export_to_video(video_frames, temp_path, fps=fps)
-        clip = VideoFileClip(temp_path)
-
-    # 2. Load and Process Audio
-    audio_clip = AudioFileClip(audio_path)
-    
-    if audio_clip.duration > clip.duration:
-        audio_clip = audio_clip.subclipped(0, clip.duration)
-    
-    # 3. Combine and Save
-    final_clip = clip.with_audio(audio_clip)
-    
-    final_clip.write_videofile(
-        output_filename, 
-        fps=fps, 
-        codec="libx264", 
-        audio_codec="aac",
-        logger="bar"
-    )
-    
-    final_clip.close()
-    audio_clip.close()
-    if 'clip' in locals(): clip.close()
-    
-    return output_filename
-
-def process_image_for_aspect_ratio(image, aspect_ratio_str):
-    """
-    Crops and resizes the image to match the target resolution based on aspect ratio.
-    """
-    # Define resolutions (W, H)
-    resolutions = {
-        "1:1 (Square)": (512, 512),
-        "16:9 (Cinematic)": (768, 512),
-        "9:16 (Vertical)": (512, 768)
-    }
-    
-    target_width, target_height = resolutions.get(aspect_ratio_str, (768, 512))
-    
-    # Use ImageOps.fit to center crop and resize automatically
-    # This preserves aspect ratio of the content while filling the target dimensions
-    processed_img = ImageOps.fit(
-        image, 
-        (target_width, target_height), 
-        method=Image.LANCZOS, 
-        centering=(0.5, 0.5)
-    )
-    
-    return processed_img, target_width, target_height
-
-# --- 3. Inference Function ---
-@spaces.GPU(duration=120, size='xlarge')
-def generate(
-    image_path, 
-    audio_path, 
-    prompt, 
-    negative_prompt,
-    aspect_ratio,
-    video_duration,
-    seed
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
 ):
-    if not image_path or not audio_path:
-        raise gr.Error("Please provide both an image and an audio file.")
-
-    # Set reproducibility
-    if seed == -1:
-        seed = random.randint(0, 1000000)
-    generator = torch.Generator(device="cuda").manual_seed(seed)
-
-    # 1. Load and Preprocess Image
-    original_image = load_image(image_path)
-    # Crop/Resize logic
-    image, width, height = process_image_for_aspect_ratio(original_image, aspect_ratio)
-
-    print(f"Generating with seed: {seed}, Resolution: {width}x{height}, Duration: {video_duration}s")
-
-    # 2. Calculate Frames
-    fps = 24.0
-    # LTX-2 constraint: (num_frames - 1) % 8 == 0
-    total_frames = int(video_duration * fps)
-    
-    # Round to nearest valid block of 8, plus 1
-    # Example: 4 seconds * 24 = 96 frames. 
-    # 96 is divisible by 8. So we take 96 + 1 = 97 frames.
-    base_block = round(total_frames / 8) * 8
-    num_frames = base_block + 1
-    
-    # Ensure sane minimum
-    if num_frames < 9: num_frames = 9
-    
-    print(f"Calculated frames: {num_frames}")
-
-    # 3. Run Inference
-    video_output, _ = pipe(
-        image=image,
-        audio=audio_path,
-        prompt=prompt,
-        negative_prompt=negative_prompt, 
-        width=width,
-        height=height,
-        num_frames=num_frames,
-        frame_rate=fps,
-        num_inference_steps=8, # Distilled uses 8 steps
-        sigmas=DISTILLED_SIGMA_VALUES, 
-        guidance_scale=1.0, 
-        generator=generator,
-        return_dict=False,
-    )
-
-    # 4. Post-process: Add audio
-    output_video_path = save_video_with_audio(video_output, audio_path, fps=fps)
-    
-    return output_video_path, seed
-
-# --- 4. Gradio Interface Definition ---
-
-css = """
-#col-container { max-width: 800px; margin: 0 auto; }
-"""
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
 
-with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown("# ⚡ LTX-2 Distilled Audio-to-Video")
-        gr.Markdown("Generate lip-synced or audio-reactive video from a single image using the distilled 8-step LTX-2 model.")
-
-        with gr.Row():
-            with gr.Column():
-                input_image = gr.Image(label="Input Image", type="filepath", height=300)
-                input_audio = gr.Audio(label="Input Audio", type="filepath")
-            
-            with gr.Column():
-                result_video = gr.Video(label="Generated Video")
-        
-        prompt = gr.Textbox(
-            label="Prompt", 
-            value="A person speaking, lips moving in sync with the words, talking head",
-            lines=2
+
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+    return noise_cfg
+
+
+class LTX2AudioToVideoPipeline(DiffusionPipeline, FromSingleFileMixin, LTXVideoLoraLoaderMixin):
+    r"""
+    Pipeline for audio-to-video generation with optional image conditioning.
+
+    This pipeline generates video conditioned on input audio, forcing the video generation
+    to attend to specific audio cues. It also supports image conditioning for image-to-video
+    generation with audio.
+
+    Reference: https://github.com/Lightricks/LTX-Video
+    """
+
+    model_cpu_offload_seq = "text_encoder->connectors->transformer->vae->audio_vae->vocoder"
+    _optional_components = []
+    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKLLTX2Video,
+        audio_vae: AutoencoderKLLTX2Audio,
+        text_encoder: Gemma3ForConditionalGeneration,
+        tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
+        connectors: LTX2TextConnectors,
+        transformer: LTX2VideoTransformer3DModel,
+        vocoder: LTX2Vocoder,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            audio_vae=audio_vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            connectors=connectors,
+            transformer=transformer,
+            vocoder=vocoder,
+            scheduler=scheduler,
         )
-        
-        # New Controls
-        with gr.Row():
-            aspect_ratio = gr.Radio(
-                choices=["1:1 (Square)", "16:9 (Cinematic)", "9:16 (Vertical)"],
-                value="16:9 (Cinematic)",
-                label="Output Aspect Ratio (Auto-Crops Image)"
+
+        self.vae_spatial_compression_ratio = (
+            self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32
+        )
+        self.vae_temporal_compression_ratio = (
+            self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8
+        )
+        self.audio_vae_mel_compression_ratio = (
+            self.audio_vae.mel_compression_ratio if getattr(self, "audio_vae", None) is not None else 4
+        )
+        self.audio_vae_temporal_compression_ratio = (
+            self.audio_vae.temporal_compression_ratio if getattr(self, "audio_vae", None) is not None else 4
+        )
+        self.transformer_spatial_patch_size = (
+            self.transformer.config.patch_size if getattr(self, "transformer", None) is not None else 1
+        )
+        self.transformer_temporal_patch_size = (
+            self.transformer.config.patch_size_t if getattr(self, "transformer") is not None else 1
+        )
+
+        self.audio_sampling_rate = (
+            self.audio_vae.config.sample_rate if getattr(self, "audio_vae", None) is not None else 16000
+        )
+        self.audio_hop_length = (
+            self.audio_vae.config.mel_hop_length if getattr(self, "audio_vae", None) is not None else 160
+        )
+
+        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio, resample="bilinear")
+        self.tokenizer_max_length = (
+            self.tokenizer.model_max_length if getattr(self, "tokenizer", None) is not None else 1024
+        )
+
+    @staticmethod
+    def _pack_text_embeds(
+        text_hidden_states: torch.Tensor,
+        sequence_lengths: torch.Tensor,
+        device: Union[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: Union[str, List[str]],
+        num_videos_per_prompt: int = 1,
+        max_sequence_length: int = 1024,
+        scale_factor: int = 8,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        if getattr(self, "tokenizer", None) is not None:
+            self.tokenizer.padding_side = "left"
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        prompt = [p.strip() for p in prompt]
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        prompt_attention_mask = text_inputs.attention_mask
+        text_input_ids = text_input_ids.to(device)
+        prompt_attention_mask = prompt_attention_mask.to(device)
+
+        text_encoder_outputs = self.text_encoder(
+            input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True
+        )
+        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)
+
+        _, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
+
+        prompt_attention_mask = prompt_attention_mask.view(batch_size, -1)
+        prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1)
+
+        return prompt_embeds, prompt_attention_mask
+
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        do_classifier_free_guidance: bool = True,
+        num_videos_per_prompt: int = 1,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_attention_mask: Optional[torch.Tensor] = None,
+        negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+        max_sequence_length: int = 1024,
+        scale_factor: int = 8,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds(
+                prompt=prompt,
+                num_videos_per_prompt=num_videos_per_prompt,
+                max_sequence_length=max_sequence_length,
+                scale_factor=scale_factor,
+                device=device,
+                dtype=dtype,
+            )
+
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+            if prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+
+            negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds(
+                prompt=negative_prompt,
+                num_videos_per_prompt=num_videos_per_prompt,
+                max_sequence_length=max_sequence_length,
+                scale_factor=scale_factor,
+                device=device,
+                dtype=dtype,
+            )
+
+        return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
+
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        callback_on_step_end_tensor_inputs=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+    ):
+        if height % 32 != 0 or width % 32 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.")
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
             )
-            video_duration = gr.Slider(
-                label="Video Duration (Seconds)",
-                minimum=1.0,
-                maximum=6.0,
-                step=0.5,
-                value=4.0,
-                info="Approximate length. Longer videos require more GPU memory."
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
             )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if prompt_embeds is not None and prompt_attention_mask is None:
+            raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
+
+        if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
+            raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
 
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Textbox(
-                label="Negative Prompt", 
-                value="low quality, worst quality, deformed, distorted",
-                placeholder="Usually ignored by distilled models with guidance 1.0"
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+            if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
+                raise ValueError(
+                    "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
+                    f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
+                    f" {negative_prompt_attention_mask.shape}."
+                )
+
+    @staticmethod
+    def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = latents.shape
+        post_patch_num_frames = num_frames // patch_size_t
+        post_patch_height = height // patch_size
+        post_patch_width = width // patch_size
+        latents = latents.reshape(
+            batch_size,
+            -1,
+            post_patch_num_frames,
+            patch_size_t,
+            post_patch_height,
+            patch_size,
+            post_patch_width,
+            patch_size,
+        )
+        latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3)
+        return latents
+
+    @staticmethod
+    def _unpack_latents(
+        latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1
+    ) -> torch.Tensor:
+        batch_size = latents.size(0)
+        latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size)
+        latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
+        return latents
+
+    @staticmethod
+    def _normalize_latents(
+        latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
+    ) -> torch.Tensor:
+        latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+        latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+        latents = (latents - latents_mean) * scaling_factor / latents_std
+        return latents
+
+    @staticmethod
+    def _denormalize_latents(
+        latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
+    ) -> torch.Tensor:
+        latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+        latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+        latents = latents * latents_std / scaling_factor + latents_mean
+        return latents
+
+    @staticmethod
+    def _pack_audio_latents(
+        latents: torch.Tensor, patch_size: Optional[int] = None, patch_size_t: Optional[int] = None
+    ) -> torch.Tensor:
+        if patch_size is not None and patch_size_t is not None:
+            batch_size, num_channels, latent_length, latent_mel_bins = latents.shape
+            post_patch_latent_length = latent_length / patch_size_t
+            post_patch_mel_bins = latent_mel_bins / patch_size
+            latents = latents.reshape(
+                batch_size, -1, post_patch_latent_length, patch_size_t, post_patch_mel_bins, patch_size
             )
-            seed = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
+            latents = latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2)
+        else:
+            latents = latents.transpose(1, 2).flatten(2, 3)
+        return latents
 
-        run_btn = gr.Button("Generate Video", variant="primary")
+    @staticmethod
+    def _unpack_audio_latents(
+        latents: torch.Tensor,
+        latent_length: int,
+        num_mel_bins: int,
+        patch_size: Optional[int] = None,
+        patch_size_t: Optional[int] = None,
+    ) -> torch.Tensor:
+        if patch_size is not None and patch_size_t is not None:
+            batch_size = latents.size(0)
+            latents = latents.reshape(batch_size, latent_length, num_mel_bins, -1, patch_size_t, patch_size)
+            latents = latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3)
+        else:
+            latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2)
+        return latents
+
+    @staticmethod
+    def _denormalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor):
+        """
+        Denormalize audio latents. The latents should be in patchified form [B, T, C*F]
+        where the last dimension matches the size of latents_mean/latents_std.
+        """
+        latents_mean = latents_mean.to(latents.device, latents.dtype)
+        latents_std = latents_std.to(latents.device, latents.dtype)
+        return (latents * latents_std) + latents_mean
+
+    @staticmethod
+    def _normalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor):
+        """
+        Normalize audio latents. The latents should be in patchified form [B, T, C*F]
+        where the last dimension matches the size of latents_mean/latents_std.
+        """
+        latents_mean = latents_mean.to(latents.device, latents.dtype)
+        latents_std = latents_std.to(latents.device, latents.dtype)
+        return (latents - latents_mean) / latents_std
+
+    @staticmethod
+    def _patchify_audio_latents(latents: torch.Tensor) -> torch.Tensor:
+        """
+        Patchify audio latents from [B, C, T, F] to [B, T, C*F].
+        This is needed for normalization which operates on the flattened channel*freq dimension.
+        """
+        # latents shape: [B, C, T, F] -> [B, T, C*F]
+        batch, channels, time, freq = latents.shape
+        return latents.permute(0, 2, 1, 3).reshape(batch, time, channels * freq)
+
+    @staticmethod
+    def _unpatchify_audio_latents(latents: torch.Tensor, channels: int, freq: int) -> torch.Tensor:
+        """
+        Unpatchify audio latents from [B, T, C*F] to [B, C, T, F].
+        """
+        # latents shape: [B, T, C*F] -> [B, C, T, F]
+        batch, time, _ = latents.shape
+        return latents.reshape(batch, time, channels, freq).permute(0, 2, 1, 3)
+
+    def _preprocess_audio(self, audio: Union[str, torch.Tensor], target_sample_rate: int) -> torch.Tensor:
+        """
+        Reads audio and converts to Mel Spectrogram matching Audio VAE expectations.
+
+        The Audio VAE encoder expects input shape: (batch_size, in_channels, time, mel_bins)
+        where in_channels=2 (stereo) and mel_bins=64 by default.
+
+        Uses the same mel spectrogram parameters as Wan2GP's AudioProcessor for compatibility.
+        """
+        if isinstance(audio, str):
+            waveform, sr = torchaudio.load(audio)
+        else:
+            waveform = audio
+            sr = target_sample_rate
+
+        if sr != target_sample_rate:
+            waveform = torchaudio.functional.resample(waveform, sr, target_sample_rate)
+
+        # Handle mono/stereo: VAE expects 2 channels
+        if waveform.shape[0] == 1:
+            # Duplicate mono to stereo
+            waveform = waveform.repeat(2, 1)
+        elif waveform.shape[0] > 2:
+            # Take first 2 channels if more than stereo
+            waveform = waveform[:2, :]
+
+        # Add batch dimension: [channels, samples] -> [batch, channels, samples]
+        waveform = waveform.unsqueeze(0)
+
+        n_fft = 1024
+        # Mel spectrogram parameters matching Wan2GP's AudioProcessor exactly
+        mel_transform = T.MelSpectrogram(
+            sample_rate=target_sample_rate,
+            n_fft=n_fft,
+            win_length=n_fft,
+            hop_length=self.audio_hop_length,
+            f_min=0.0,
+            f_max=target_sample_rate / 2.0,
+            n_mels=self.audio_vae.config.mel_bins,
+            window_fn=torch.hann_window,
+            center=True,
+            pad_mode="reflect",
+            power=1.0,  # Important: power=1.0, not 2.0
+            mel_scale="slaney",
+            norm="slaney",
+        )
+
+        # waveform shape: [batch, channels, samples]
+        # mel_spec shape after transform: [batch, channels, mel_bins, time]
+        mel_spec = mel_transform(waveform)
+
+        # Log scaling
+        mel_spec = torch.log(torch.clamp(mel_spec, min=1e-5))
+
+        # Permute to [batch, channels, time, mel_bins] as expected by VAE
+        mel_spec = mel_spec.permute(0, 1, 3, 2).contiguous()
+
+        return mel_spec
+
+    def prepare_latents(
+        self,
+        image: Optional[torch.Tensor] = None,
+        batch_size: int = 1,
+        num_channels_latents: int = 128,
+        height: int = 512,
+        width: int = 704,
+        num_frames: int = 161,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+        generator: Optional[torch.Generator] = None,
+        latents: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        height = height // self.vae_spatial_compression_ratio
+        width = width // self.vae_spatial_compression_ratio
+        num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
+
+        shape = (batch_size, num_channels_latents, num_frames, height, width)
+        mask_shape = (batch_size, 1, num_frames, height, width)
+
+        if latents is not None:
+            conditioning_mask = latents.new_zeros(mask_shape)
+            conditioning_mask[:, :, 0] = 1.0
+            conditioning_mask = self._pack_latents(
+                conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size
+            ).squeeze(-1)
+            if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape:
+                raise ValueError(
+                    f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}."
+                )
+            return latents.to(device=device, dtype=dtype), conditioning_mask
+
+        if isinstance(generator, list):
+            if len(generator) != batch_size:
+                raise ValueError(
+                    f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                    f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+                )
+
+            init_latents = [
+                retrieve_latents(self.vae.encode(image[i].unsqueeze(0).unsqueeze(2)), generator[i], "argmax")
+                for i in range(batch_size)
+            ]
+        else:
+            init_latents = [
+                retrieve_latents(self.vae.encode(img.unsqueeze(0).unsqueeze(2)), generator, "argmax") for img in image
+            ]
+
+        init_latents = torch.cat(init_latents, dim=0).to(dtype)
+        init_latents = self._normalize_latents(init_latents, self.vae.latents_mean, self.vae.latents_std)
+        init_latents = init_latents.repeat(1, 1, num_frames, 1, 1)
+
+        conditioning_mask = torch.zeros(mask_shape, device=device, dtype=dtype)
+        conditioning_mask[:, :, 0] = 1.0
+
+        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = init_latents * conditioning_mask + noise * (1 - conditioning_mask)
+
+        conditioning_mask = self._pack_latents(
+            conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size
+        ).squeeze(-1)
+        latents = self._pack_latents(
+            latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size
+        )
+
+        return latents, conditioning_mask
+
+    def prepare_audio_latents(
+        self,
+        batch_size: int = 1,
+        num_channels_latents: int = 8,
+        num_mel_bins: int = 64,
+        num_frames: int = 121,
+        frame_rate: float = 25.0,
+        sampling_rate: int = 16000,
+        hop_length: int = 160,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+        generator: Optional[torch.Generator] = None,
+        audio_input: Optional[Union[str, torch.Tensor]] = None,
+        latents: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, int, Optional[torch.Tensor]]:
+        duration_s = num_frames / frame_rate
+        latents_per_second = (
+            float(sampling_rate) / float(hop_length) / float(self.audio_vae_temporal_compression_ratio)
+        )
+        target_length = round(duration_s * latents_per_second)
+
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype), target_length, None
+
+        latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
+
+        if audio_input is not None:
+            mel_spec = self._preprocess_audio(audio_input, sampling_rate).to(device=device)
+
+            # Encode the mel spectrogram to latents (use VAE's dtype for encoding)
+            mel_spec = mel_spec.to(dtype=self.audio_vae.dtype)
+            init_latents = self.audio_vae.encode(mel_spec).latent_dist.sample(generator)
+            init_latents = init_latents.to(dtype=dtype)
+
+            # Normalize: patchify -> normalize -> unpatchify
+            # init_latents shape: [B, C, T, F] where C=latent_channels, F=latent_mel_bins
+            latent_channels = init_latents.shape[1]
+            latent_freq = init_latents.shape[3]
+            init_latents_patched = self._patchify_audio_latents(init_latents)  # [B, T, C*F]
+            init_latents_patched = self._normalize_audio_latents(
+                init_latents_patched, self.audio_vae.latents_mean, self.audio_vae.latents_std
+            )
+            init_latents = self._unpatchify_audio_latents(init_latents_patched, latent_channels, latent_freq)  # [B, C, T, F]
+
+            current_len = init_latents.shape[2]
+            if current_len < target_length:
+                padding = target_length - current_len
+                init_latents = torch.nn.functional.pad(init_latents, (0, 0, 0, padding))
+            elif current_len > target_length:
+                init_latents = init_latents[:, :, :target_length, :]
+
+            noise = randn_tensor(init_latents.shape, generator=generator, device=device, dtype=dtype)
+
+            if init_latents.shape[0] != batch_size:
+                init_latents = init_latents.repeat(batch_size, 1, 1, 1)
+                noise = noise.repeat(batch_size, 1, 1, 1)
+
+            packed_noise = self._pack_audio_latents(noise)
+
+            return packed_noise, target_length, init_latents
+
+        shape = (batch_size, num_channels_latents, target_length, latent_mel_bins)
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError("Generator size mismatch")
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = self._pack_audio_latents(latents)
+
+        return latents, target_length, None
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def guidance_rescale(self):
+        return self._guidance_rescale
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1.0
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def current_timestep(self):
+        return self._current_timestep
+
+    @property
+    def attention_kwargs(self):
+        return self._attention_kwargs
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    def _get_audio_duration(self, audio: Union[str, torch.Tensor], sample_rate: int) -> float:
+        """Get duration of audio in seconds."""
+        if isinstance(audio, str):
+            info = torchaudio.info(audio)
+            return info.num_frames / info.sample_rate
+        else:
+            # audio is a tensor with shape [channels, samples] or [samples]
+            num_samples = audio.shape[-1]
+            return num_samples / sample_rate
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        image: PipelineImageInput = None,
+        audio: Optional[Union[str, torch.Tensor]] = None,
+        prompt: Union[str, List[str]] = None,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        height: int = 512,
+        width: int = 768,
+        num_frames: Optional[int] = None,
+        max_frames: int = 257,
+        frame_rate: float = 24.0,
+        num_inference_steps: int = 40,
+        timesteps: List[int] = None,
+        sigmas: Optional[List[float]] = None,
+        guidance_scale: float = 4.0,
+        guidance_rescale: float = 0.0,
+        num_videos_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        audio_latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_attention_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+        decode_timestep: Union[float, List[float]] = 0.0,
+        decode_noise_scale: Optional[Union[float, List[float]]] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 1024,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            image (`PipelineImageInput`):
+                The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`.
+            audio (`str` or `torch.Tensor`, *optional*):
+                The input audio to condition the generation on. Can be a path to an audio file or a waveform tensor.
+                When provided, the generated video will be synchronized to this audio input.
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation.
+            height (`int`, *optional*, defaults to `512`):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to `768`):
+                The width in pixels of the generated image.
+            num_frames (`int`, *optional*):
+                The number of video frames to generate. If not provided and audio is given,
+                it will be calculated from the audio duration. Otherwise defaults to 121.
+            max_frames (`int`, *optional*, defaults to `257`):
+                Maximum number of frames to generate. Used to cap the calculated frames
+                when deriving from audio duration. 257 frames at 25fps is ~10 seconds.
+            frame_rate (`float`, *optional*, defaults to `24.0`):
+                The frames per second (FPS) of the generated video.
+            num_inference_steps (`int`, *optional*, defaults to 40):
+                The number of denoising steps.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            guidance_scale (`float`, *optional*, defaults to `4.0`):
+                Guidance scale for classifier-free guidance.
+            guidance_rescale (`float`, *optional*, defaults to 0.0):
+                Guidance rescale factor.
+            num_videos_per_prompt (`int`, *optional*, defaults to 1):
+                The number of videos to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                Random generator(s) for reproducibility.
+            latents (`torch.Tensor`, *optional*):
+                Pre-generated noisy latents for video generation.
+            audio_latents (`torch.Tensor`, *optional*):
+                Pre-generated noisy latents for audio generation.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings.
+            prompt_attention_mask (`torch.Tensor`, *optional*):
+                Pre-generated attention mask for text embeddings.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings.
+            negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):
+                Pre-generated attention mask for negative text embeddings.
+            decode_timestep (`float`, defaults to `0.0`):
+                The timestep at which generated video is decoded.
+            decode_noise_scale (`float`, defaults to `None`):
+                The interpolation factor between random noise and denoised latents.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generated image.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a `LTX2PipelineOutput` instead of a plain tuple.
+            attention_kwargs (`dict`, *optional*):
+                Kwargs dictionary for the attention processor.
+            callback_on_step_end (`Callable`, *optional*):
+                A function called at the end of each denoising step.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function.
+            max_sequence_length (`int`, *optional*, defaults to `1024`):
+                Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`LTX2PipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, `LTX2PipelineOutput` is returned, otherwise a `tuple`.
+        """
+
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        # Calculate num_frames from audio duration if not provided
+        if num_frames is None:
+            if audio is not None:
+                # Get audio duration and calculate frames
+                audio_duration = self._get_audio_duration(audio, self.audio_sampling_rate)
+                # Calculate frames needed for this audio duration
+                # Add 1 because frames = duration * fps + 1 (first frame at t=0)
+                calculated_frames = int(audio_duration * frame_rate) + 1
+                # Cap at max_frames and ensure it's valid for the model
+                # LTX2 requires (num_frames - 1) to be divisible by temporal_compression_ratio (8)
+                num_frames = min(calculated_frames, max_frames)
+                # Adjust to valid frame count: (num_frames - 1) % 8 == 0
+                num_frames = ((num_frames - 1) // self.vae_temporal_compression_ratio) * self.vae_temporal_compression_ratio + 1
+                num_frames = max(num_frames, 9)  # Minimum valid frame count
+                logger.info(f"Audio duration: {audio_duration:.2f}s -> num_frames: {num_frames}")
+            else:
+                num_frames = 121  # Default
+
+        self.check_inputs(
+            prompt=prompt,
+            height=height,
+            width=width,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._guidance_rescale = guidance_rescale
+        self._attention_kwargs = attention_kwargs
+        self._interrupt = False
+        self._current_timestep = None
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        (
+            prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_embeds,
+            negative_prompt_attention_mask,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            num_videos_per_prompt=num_videos_per_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            max_sequence_length=max_sequence_length,
+            device=device,
+        )
+        if self.do_classifier_free_guidance:
+            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)
+
+        additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0
+        connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors(
+            prompt_embeds, additive_attention_mask, additive_mask=True
+        )
+
+        if latents is None:
+            image = self.video_processor.preprocess(image, height=height, width=width)
+            image = image.to(device=device, dtype=prompt_embeds.dtype)
+
+        num_channels_latents = self.transformer.config.in_channels
+        latents, conditioning_mask = self.prepare_latents(
+            image,
+            batch_size * num_videos_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            num_frames,
+            torch.float32,
+            device,
+            generator,
+            latents,
+        )
+        if self.do_classifier_free_guidance:
+            conditioning_mask = torch.cat([conditioning_mask, conditioning_mask])
+
+        num_mel_bins = self.audio_vae.config.mel_bins if getattr(self, "audio_vae", None) is not None else 64
+        latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
+
+        num_channels_latents_audio = (
+            self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8
+        )
+
+        audio_latents, audio_num_frames, clean_audio_latents = self.prepare_audio_latents(
+            batch_size * num_videos_per_prompt,
+            num_channels_latents=num_channels_latents_audio,
+            num_mel_bins=num_mel_bins,
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            sampling_rate=self.audio_sampling_rate,
+            hop_length=self.audio_hop_length,
+            dtype=torch.float32,
+            device=device,
+            generator=generator,
+            latents=audio_latents,
+            audio_input=audio,
+        )
+
+        # If clean audio latents are provided, pack them for use in the transformer
+        # This is the key fix: we pass clean (not noisy) audio latents to the transformer
+        packed_clean_audio_latents = None
+        if clean_audio_latents is not None:
+            packed_clean_audio_latents = self._pack_audio_latents(clean_audio_latents)
+
+        latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
+        latent_height = height // self.vae_spatial_compression_ratio
+        latent_width = width // self.vae_spatial_compression_ratio
+        video_sequence_length = latent_num_frames * latent_height * latent_width
+
+        if sigmas is None:
+            sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
         
-        # Output info
-        used_seed = gr.Number(label="Used Seed", visible=False)
-
-    # Event Logic
-    run_btn.click(
-        fn=generate,
-        inputs=[
-            input_image, 
-            input_audio, 
-            prompt, 
-            negative_prompt, 
-            aspect_ratio,
-            video_duration,
-            seed
-        ],
-        outputs=[result_video, used_seed]
-    )
-
-if __name__ == "__main__":
-    demo.queue().launch()
+        mu = calculate_shift(
+            video_sequence_length,
+            self.scheduler.config.get("base_image_seq_len", 1024),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.95),
+            self.scheduler.config.get("max_shift", 2.05),
+        )
+
+        audio_scheduler = copy.deepcopy(self.scheduler)
+        _, _ = retrieve_timesteps(
+            audio_scheduler,
+            num_inference_steps,
+            device,
+            timesteps,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            timesteps,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        rope_interpolation_scale = (
+            self.vae_temporal_compression_ratio / frame_rate,
+            self.vae_spatial_compression_ratio,
+            self.vae_spatial_compression_ratio,
+        )
+        video_coords = self.transformer.rope.prepare_video_coords(
+            latents.shape[0], latent_num_frames, latent_height, latent_width, latents.device, fps=frame_rate
+        )
+        audio_coords = self.transformer.audio_rope.prepare_audio_coords(
+            audio_latents.shape[0], audio_num_frames, audio_latents.device
+        )
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                self._current_timestep = t
+
+                # When audio conditioning is provided, use clean audio latents directly
+                # (not noisy). This matches Wan2GP's approach where audio with denoise_mask=0
+                # stays constant throughout denoising.
+                if packed_clean_audio_latents is not None:
+                    audio_latents_input = packed_clean_audio_latents.to(dtype=prompt_embeds.dtype)
+                else:
+                    audio_latents_input = audio_latents.to(dtype=prompt_embeds.dtype)
+
+                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+                latent_model_input = latent_model_input.to(prompt_embeds.dtype)
+                audio_latent_model_input = (
+                    torch.cat([audio_latents_input] * 2) if self.do_classifier_free_guidance else audio_latents_input
+                )
+                audio_latent_model_input = audio_latent_model_input.to(prompt_embeds.dtype)
+
+                timestep = t.expand(latent_model_input.shape[0])
+                video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask)
+
+                # When audio conditioning is provided, set audio_timestep to 0.
+                # This tells the transformer the audio is already "denoised" (clean),
+                # matching Wan2GP's approach where denoise_mask=0 results in timestep=0.
+                if packed_clean_audio_latents is not None:
+                    audio_timestep = torch.zeros_like(timestep)
+                else:
+                    audio_timestep = timestep
+
+                with self.transformer.cache_context("cond_uncond"):
+                    noise_pred_video, noise_pred_audio = self.transformer(
+                        hidden_states=latent_model_input,
+                        audio_hidden_states=audio_latent_model_input,
+                        encoder_hidden_states=connector_prompt_embeds,
+                        audio_encoder_hidden_states=connector_audio_prompt_embeds,
+                        timestep=video_timestep,
+                        audio_timestep=audio_timestep,
+                        encoder_attention_mask=connector_attention_mask,
+                        audio_encoder_attention_mask=connector_attention_mask,
+                        num_frames=latent_num_frames,
+                        height=latent_height,
+                        width=latent_width,
+                        fps=frame_rate,
+                        audio_num_frames=audio_num_frames,
+                        video_coords=video_coords,
+                        audio_coords=audio_coords,
+                        attention_kwargs=attention_kwargs,
+                        return_dict=False,
+                    )
+                noise_pred_video = noise_pred_video.float()
+                noise_pred_audio = noise_pred_audio.float()
+
+                if self.do_classifier_free_guidance:
+                    noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2)
+                    noise_pred_video = noise_pred_video_uncond + self.guidance_scale * (
+                        noise_pred_video_text - noise_pred_video_uncond
+                    )
+
+                    noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2)
+                    noise_pred_audio = noise_pred_audio_uncond + self.guidance_scale * (
+                        noise_pred_audio_text - noise_pred_audio_uncond
+                    )
+
+                    if self.guidance_rescale > 0:
+                        noise_pred_video = rescale_noise_cfg(
+                            noise_pred_video, noise_pred_video_text, guidance_rescale=self.guidance_rescale
+                        )
+                        noise_pred_audio = rescale_noise_cfg(
+                            noise_pred_audio, noise_pred_audio_text, guidance_rescale=self.guidance_rescale
+                        )
+
+                noise_pred_video = self._unpack_latents(
+                    noise_pred_video,
+                    latent_num_frames,
+                    latent_height,
+                    latent_width,
+                    self.transformer_spatial_patch_size,
+                    self.transformer_temporal_patch_size,
+                )
+                latents = self._unpack_latents(
+                    latents,
+                    latent_num_frames,
+                    latent_height,
+                    latent_width,
+                    self.transformer_spatial_patch_size,
+                    self.transformer_temporal_patch_size,
+                )
+
+                noise_pred_video = noise_pred_video[:, :, 1:]
+                noise_latents = latents[:, :, 1:]
+                pred_latents = self.scheduler.step(noise_pred_video, t, noise_latents, return_dict=False)[0]
+
+                latents = torch.cat([latents[:, :, :1], pred_latents], dim=2)
+                latents = self._pack_latents(
+                    latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size
+                )
+
+                # Only step audio latents when not using audio conditioning.
+                # When audio conditioning is provided, we keep the clean latents constant.
+                if packed_clean_audio_latents is None:
+                    audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
+                # else: audio stays constant (packed_clean_audio_latents)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        latents = self._unpack_latents(
+            latents,
+            latent_num_frames,
+            latent_height,
+            latent_width,
+            self.transformer_spatial_patch_size,
+            self.transformer_temporal_patch_size,
+        )
+        latents = self._denormalize_latents(
+            latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
+        )
+
+        # Denormalize audio latents for decoding
+        # Need to: patchify -> denormalize -> unpatchify (inverse of normalization)
+        if clean_audio_latents is not None:
+            # clean_audio_latents is in 4D format [B, C, T, F]
+            latent_channels = clean_audio_latents.shape[1]
+            latent_freq = clean_audio_latents.shape[3]
+            audio_patched = self._patchify_audio_latents(clean_audio_latents)  # [B, T, C*F]
+            audio_patched = self._denormalize_audio_latents(
+                audio_patched, self.audio_vae.latents_mean, self.audio_vae.latents_std
+            )
+            audio_latents_for_decode = self._unpatchify_audio_latents(audio_patched, latent_channels, latent_freq)
+        else:
+            # audio_latents is in packed format [B, T, C*F] from the denoising loop
+            audio_latents_for_decode = self._denormalize_audio_latents(
+                audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std
+            )
+            # Unpack to 4D format [B, C, T, F]
+            audio_latents_for_decode = self._unpack_audio_latents(
+                audio_latents_for_decode, audio_num_frames, num_mel_bins=latent_mel_bins
+            )
+
+        if output_type == "latent":
+            video = latents
+            audio = audio_latents_for_decode
+        else:
+            latents = latents.to(prompt_embeds.dtype)
+
+            if not self.vae.config.timestep_conditioning:
+                timestep = None
+            else:
+                noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype)
+                if not isinstance(decode_timestep, list):
+                    decode_timestep = [decode_timestep] * batch_size
+                if decode_noise_scale is None:
+                    decode_noise_scale = decode_timestep
+                elif not isinstance(decode_noise_scale, list):
+                    decode_noise_scale = [decode_noise_scale] * batch_size
+
+                timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype)
+                decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[
+                    :, None, None, None, None
+                ]
+                latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise
+
+            latents = latents.to(self.vae.dtype)
+            video = self.vae.decode(latents, timestep, return_dict=False)[0]
+            video = self.video_processor.postprocess_video(video, output_type=output_type)
+
+            audio_latents_for_decode = audio_latents_for_decode.to(self.audio_vae.dtype)
+            generated_mel_spectrograms = self.audio_vae.decode(audio_latents_for_decode, return_dict=False)[0]
+            audio = self.vocoder(generated_mel_spectrograms)
+
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (video, audio)
+
+        return LTX2PipelineOutput(frames=video, audio=audio)