support ltx2 a2v

Author: mi804

diffsynth/pipelines/ltx2_audio_video.py

@@ -67,8 +67,9 @@ def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16):
             LTX2AudioVideoUnit_SwitchStage2(),
             LTX2AudioVideoUnit_NoiseInitializer(),
             LTX2AudioVideoUnit_LatentsUpsampler(),
-            LTX2AudioVideoUnit_SetScheduleStage2(),
             LTX2AudioVideoUnit_InputImagesEmbedder(),
+            LTX2AudioVideoUnit_InputAudioEmbedder(),
+            LTX2AudioVideoUnit_SetScheduleStage2(),
         ]
         self.model_fn = model_fn_ltx2
 
@@ -155,8 +156,9 @@ def denoise_stage(self, inputs_shared, inputs_posi, inputs_nega, units, cfg_scal
                 **models, timestep=timestep, progress_id=progress_id
             )
             inputs_shared["video_latents"] = self.step(self.scheduler, inputs_shared["video_latents"], progress_id=progress_id, noise_pred=noise_pred_video,
-                                                       inpaint_mask=inputs_shared.get("denoise_mask_video", None), input_latents=inputs_shared.get("input_latents_video", None), **inputs_shared)
-            inputs_shared["audio_latents"] = self.step(self.scheduler, inputs_shared["audio_latents"], progress_id=progress_id, noise_pred=noise_pred_audio, **inputs_shared)
+                                                       inpaint_mask=inputs_shared.get("video_denoise_mask", None), input_latents=inputs_shared.get("video_input_latents", None), **inputs_shared)
+            inputs_shared["audio_latents"] = self.step(self.scheduler, inputs_shared["audio_latents"], progress_id=progress_id, noise_pred=noise_pred_audio,
+                                                       inpaint_mask=inputs_shared.get("audio_denoise_mask", None), input_latents=inputs_shared.get("audio_input_latents", None), **inputs_shared)
         return inputs_shared, inputs_posi, inputs_nega
 
     @torch.no_grad()
@@ -173,6 +175,9 @@ def __call__(
         # In-Context Video Control
         in_context_videos: Optional[list[list[Image.Image]]] = None,
         in_context_downsample_factor: Optional[int] = 2,
+        # Audio-to-video
+        input_audio: Optional[torch.Tensor] = None,
+        audio_sample_rate: Optional[int] = 48000,
         # Randomness
         seed: Optional[int] = None,
         rand_device: Optional[str] = "cpu",
@@ -210,6 +215,7 @@ def __call__(
         }
         inputs_shared = {
             "input_images": input_images, "input_images_indexes": input_images_indexes, "input_images_strength": input_images_strength,
+            "input_audio": (input_audio, audio_sample_rate) if input_audio is not None else None,
             "in_context_videos": in_context_videos, "in_context_downsample_factor": in_context_downsample_factor,
             "seed": seed, "rand_device": rand_device,
             "height": height, "width": width, "num_frames": num_frames, "frame_rate": frame_rate,
@@ -361,7 +367,8 @@ def process(self, pipe: LTX2AudioVideoPipeline, input_video, video_noise, tiled,
             input_video = pipe.preprocess_video(input_video)
             input_latents = pipe.video_vae_encoder.encode(input_video, tiled, tile_size_in_pixels, tile_overlap_in_pixels).to(dtype=pipe.torch_dtype, device=pipe.device)
             if pipe.scheduler.training:
-                return {"video_latents": input_latents, "input_latents": input_latents}
+                # input_latents key is for training to add noise. with no prefix "video" to keep loss function keyword consistent.
+                return {"video_latents": video_noise, "input_latents": input_latents}
             else:
                 raise NotImplementedError("Video-to-video not implemented yet.")
 
@@ -370,7 +377,7 @@ class LTX2AudioVideoUnit_InputAudioEmbedder(PipelineUnit):
     def __init__(self):
         super().__init__(
             input_params=("input_audio", "audio_noise"),
-            output_params=("audio_latents", "audio_input_latents", "audio_positions", "audio_latent_shape"),
+            output_params=("audio_latents", "audio_input_latents", "audio_noise", "audio_positions", "audio_latent_shape", "audio_denoise_mask"),
             onload_model_names=("audio_vae_encoder",)
         )
 
@@ -380,21 +387,37 @@ def process(self, pipe: LTX2AudioVideoPipeline, input_audio, audio_noise):
         else:
             input_audio, sample_rate = input_audio
             pipe.load_models_to_device(self.onload_model_names)
-            input_audio = pipe.audio_processor.waveform_to_mel(input_audio.unsqueeze(0), waveform_sample_rate=sample_rate).to(dtype=pipe.torch_dtype)
+            input_audio = pipe.audio_processor.waveform_to_mel(input_audio.unsqueeze(0), waveform_sample_rate=sample_rate).to(dtype=pipe.torch_dtype, device=pipe.device)
             audio_input_latents = pipe.audio_vae_encoder(input_audio)
             audio_latent_shape = AudioLatentShape.from_torch_shape(audio_input_latents.shape)
             audio_positions = pipe.audio_patchifier.get_patch_grid_bounds(audio_latent_shape, device=pipe.device)
             if pipe.scheduler.training:
-                return {"audio_latents": audio_input_latents, "audio_input_latents": audio_input_latents, "audio_positions": audio_positions, "audio_latent_shape": audio_latent_shape}
+                return {
+                    "audio_latents": audio_input_latents,
+                    "audio_input_latents": audio_input_latents,
+                    "audio_positions": audio_positions,
+                    "audio_latent_shape": audio_latent_shape,
+                }
             else:
-                raise NotImplementedError("Audio-to-video not supported.")
+                b, c, t, f = audio_input_latents.shape
+                audio_denoise_mask = torch.zeros((b, 1, t, 1), device=audio_input_latents.device, dtype=audio_input_latents.dtype)
+                audio_noise = torch.rand_like(audio_input_latents)
+                audio_latents = pipe.scheduler.add_noise(audio_input_latents, audio_noise, pipe.scheduler.timesteps[0])
+                return {
+                    "audio_latents": audio_latents,
+                    "audio_input_latents": audio_input_latents,
+                    "audio_noise": audio_noise,
+                    "audio_positions": audio_positions,
+                    "audio_latent_shape": audio_latent_shape,
+                    "audio_denoise_mask": audio_denoise_mask,
+                }
 
 
 class LTX2AudioVideoUnit_InputImagesEmbedder(PipelineUnit):
     def __init__(self):
         super().__init__(
             input_params=("input_images", "input_images_indexes", "input_images_strength", "video_latents", "height", "width", "frame_rate", "tiled", "tile_size_in_pixels", "tile_overlap_in_pixels", "initial_latents"),
-            output_params=("denoise_mask_video", "input_latents_video", "ref_frames_latents", "ref_frames_positions"),
+            output_params=("video_denoise_mask", "video_input_latents", "ref_frames_latents", "ref_frames_positions"),
             onload_model_names=("video_vae_encoder")
         )
 
@@ -406,9 +429,9 @@ def get_image_latent(self, pipe, input_image, height, width, tiled, tile_size_in
         latents = pipe.video_vae_encoder.encode(image, tiled, tile_size_in_pixels, tile_overlap_in_pixels).to(pipe.device)
         return latents
 
-    def apply_input_images_to_latents(self, latents, input_latents, input_indexes, input_strength=1.0, initial_latents=None, denoise_mask_video=None):
+    def apply_input_images_to_latents(self, latents, input_latents, input_indexes, input_strength=1.0, initial_latents=None, video_denoise_mask=None):
         b, _, f, h, w = latents.shape
-        denoise_mask = torch.ones((b, 1, f, h, w), dtype=latents.dtype, device=latents.device) if denoise_mask_video is None else denoise_mask_video
+        denoise_mask = torch.ones((b, 1, f, h, w), dtype=latents.dtype, device=latents.device) if video_denoise_mask is None else video_denoise_mask
         initial_latents = torch.zeros_like(latents) if initial_latents is None else initial_latents
         for idx, input_latent in zip(input_indexes, input_latents):
             idx = min(max(1 + (idx-1) // 8, 0), f - 1)
@@ -424,13 +447,13 @@ def process(self, pipe: LTX2AudioVideoPipeline, video_latents, input_images, hei
             if len(input_images_indexes) != len(set(input_images_indexes)):
                 raise ValueError("Input images must have unique indexes.")
             pipe.load_models_to_device(self.onload_model_names)
-            frame_conditions = {"input_latents_video": None, "denoise_mask_video": None, "ref_frames_latents": [], "ref_frames_positions": []}
+            frame_conditions = {"video_input_latents": None, "video_denoise_mask": None, "ref_frames_latents": [], "ref_frames_positions": []}
             for img, index in zip(input_images, input_images_indexes):
                 latents = self.get_image_latent(pipe, img, height, width, tiled, tile_size_in_pixels, tile_overlap_in_pixels)
                 # first_frame by replacing latents
                 if index == 0:
-                    input_latents_video, denoise_mask_video = self.apply_input_images_to_latents(video_latents, [latents], [0], input_images_strength, initial_latents)
-                    frame_conditions.update({"input_latents_video": input_latents_video, "denoise_mask_video": denoise_mask_video})
+                    video_input_latents, video_denoise_mask = self.apply_input_images_to_latents(video_latents, [latents], [0], input_images_strength, initial_latents)
+                    frame_conditions.update({"video_input_latents": video_input_latents, "video_denoise_mask": video_denoise_mask})
                 # other frames by adding reference latents
                 else:
                     latent_coords = pipe.video_patchifier.get_patch_grid_bounds(output_shape=VideoLatentShape.from_torch_shape(latents.shape), device=pipe.device)
@@ -560,14 +583,17 @@ def model_fn_ltx2(
     audio_patchifier=None,
     timestep=None,
     # First Frame Conditioning
-    input_latents_video=None,
-    denoise_mask_video=None,
+    video_input_latents=None,
+    video_denoise_mask=None,
     # Other Frames Conditioning
     ref_frames_latents=None,
     ref_frames_positions=None,
     # In-Context Conditioning
     in_context_video_latents=None,
     in_context_video_positions=None,
+    # Audio Inputs
+    audio_input_latents=None,
+    audio_denoise_mask=None,
     # Gradient Checkpointing
     use_gradient_checkpointing=False,
     use_gradient_checkpointing_offload=False,
@@ -581,12 +607,12 @@ def model_fn_ltx2(
     seq_len_video = video_latents.shape[1]
     video_timesteps = timestep.repeat(1, video_latents.shape[1], 1)
     # Frist frame conditioning by replacing the video latents
-    if input_latents_video is not None:
-        denoise_mask_video = video_patchifier.patchify(denoise_mask_video)
-        video_latents = video_latents * denoise_mask_video + video_patchifier.patchify(input_latents_video) * (1.0 - denoise_mask_video)
-        video_timesteps = denoise_mask_video * video_timesteps
-    
-    # Conditioning by replacing the video latents
+    if video_input_latents is not None:
+        video_denoise_mask = video_patchifier.patchify(video_denoise_mask)
+        video_latents = video_latents * video_denoise_mask + video_patchifier.patchify(video_input_latents) * (1.0 - video_denoise_mask)
+        video_timesteps = video_denoise_mask * video_timesteps
+
+    # Reference conditioning by appending the reference video or frame latents
     total_ref_latents = ref_frames_latents if ref_frames_latents is not None else []
     total_ref_positions = ref_frames_positions if ref_frames_positions is not None else []
     total_ref_latents += [in_context_video_latents] if in_context_video_latents is not None else []
@@ -605,6 +631,10 @@ def model_fn_ltx2(
         audio_timesteps = timestep.repeat(1, audio_latents.shape[1], 1)
     else:
         audio_timesteps = None
+    if audio_input_latents is not None:
+        audio_denoise_mask = audio_patchifier.patchify(audio_denoise_mask)
+        audio_latents = audio_latents * audio_denoise_mask + audio_patchifier.patchify(audio_input_latents) * (1.0 - audio_denoise_mask)
+        audio_timesteps = audio_denoise_mask * audio_timesteps
 
     vx, ax = dit(
         video_latents=video_latents,

diffsynth/utils/data/media_io_ltx2.py

@@ -1,13 +1,11 @@
-
 from fractions import Fraction
 import torch
+import torchaudio
 import av
 from tqdm import tqdm
 from PIL import Image
 import numpy as np
 from io import BytesIO
-from collections.abc import Generator, Iterator
-import torchaudio
 
 
 def _resample_audio(
@@ -70,9 +68,9 @@ def _prepare_audio_stream(container: av.container.Container, audio_sample_rate:
     audio_stream = container.add_stream("aac")
     supported_sample_rates = audio_stream.codec_context.codec.audio_rates
     if supported_sample_rates:
-            best_rate = min(supported_sample_rates, key=lambda x: abs(x - audio_sample_rate))
-            if best_rate != audio_sample_rate:
-                print(f"Using closest supported audio sample rate: {best_rate}")
+        best_rate = min(supported_sample_rates, key=lambda x: abs(x - audio_sample_rate))
+        if best_rate != audio_sample_rate:
+            print(f"Using closest supported audio sample rate: {best_rate}")
     else:
         best_rate = audio_sample_rate
     audio_stream.codec_context.sample_rate = best_rate
@@ -117,7 +115,7 @@ def write_video_audio_ltx2(
     stream.width = width
     stream.height = height
     stream.pix_fmt = "yuv420p"
-    
+
     if audio is not None:
         if audio_sample_rate is None:
             raise ValueError("audio_sample_rate is required when audio is provided")
@@ -138,13 +136,24 @@ def write_video_audio_ltx2(
     container.close()
 
 
-def read_audio_with_torchaudio(path: str, start_time: float = 0, duration: float | None = None) -> torch.Tensor:
+def resample_waveform(waveform: torch.Tensor, source_rate: int, target_rate: int) -> torch.Tensor:
+    """Resample waveform to target sample rate if needed."""
+    if source_rate == target_rate:
+        return waveform
+    resampled = torchaudio.functional.resample(waveform, source_rate, target_rate)
+    return resampled.to(dtype=waveform.dtype)
+
+
+def read_audio_with_torchaudio(path: str, start_time: float = 0, duration: float | None = None, resample: bool = False, resample_rate: int = 48000) -> torch.Tensor:
     waveform, sample_rate = torchaudio.load(path, channels_first=True)
+    if resample:
+        waveform = resample_waveform(waveform, sample_rate, resample_rate)
+        sample_rate = resample_rate
     start_frame = int(start_time * sample_rate)
     if start_frame > waveform.shape[-1]:
         raise ValueError(f"start_time of {start_time} exceeds max duration of {waveform.shape[-1] / sample_rate:.2f}")
-    end_frame = -1 if duration is None else int(duration * sample_rate)
-    return waveform[..., start_frame:end_frame]
+    end_frame = -1 if duration is None else int(duration * sample_rate + start_frame)
+    return waveform[..., start_frame:end_frame], sample_rate
 
 
 def encode_single_frame(output_file: str, image_array: np.ndarray, crf: float) -> None:

examples/ltx2/model_inference/LTX-2.3-A2V-TwoStage.py

@@ -1,8 +1,7 @@
 import torch
 from diffsynth.pipelines.ltx2_audio_video import LTX2AudioVideoPipeline, ModelConfig
 from diffsynth.utils.data.media_io_ltx2 import read_audio_with_torchaudio, write_video_audio_ltx2
-
-audio = read_audio_with_torchaudio("data/example_video_dataset/ltx2/sing.MP3")
+from modelscope import dataset_snapshot_download
 
 vram_config = {
     "offload_dtype": torch.bfloat16,
@@ -25,7 +24,9 @@
     tokenizer_config=ModelConfig(model_id="google/gemma-3-12b-it-qat-q4_0-unquantized"),
     stage2_lora_config=ModelConfig(model_id="Lightricks/LTX-2.3", origin_file_pattern="ltx-2.3-22b-distilled-lora-384.safetensors"),
 )
-prompt = "A girl is very happy, she is speaking: “I enjoy working with Diffsynth-Studio, it's a perfect framework.”"
+
+dataset_snapshot_download("DiffSynth-Studio/example_video_dataset", allow_file_pattern="ltx2/*", local_dir="data/example_video_dataset")
+prompt = "A beautiful woman with a flower crown is singing happily under a blooming cherry tree."
 negative_prompt = (
     "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, "
     "grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, "
@@ -39,21 +40,26 @@
     "pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, "
     "inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts."
 )
-height, width, num_frames = 512 * 2, 768 * 2, 121
+height, width, num_frames, frame_rate = 512 * 2, 768 * 2, 121, 24
+duration = num_frames / frame_rate
+audio, audio_sample_rate = read_audio_with_torchaudio("data/example_video_dataset/ltx2/sing.MP3", start_time=1, duration=duration)
 video, audio = pipe(
     prompt=prompt,
     negative_prompt=negative_prompt,
+    input_audio=audio,
+    audio_sample_rate=audio_sample_rate,
     seed=43,
     height=height,
     width=width,
     num_frames=num_frames,
+    frame_rate=frame_rate,
     tiled=True,
     use_two_stage_pipeline=True,
 )
 write_video_audio_ltx2(
     video=video,
     audio=audio,
-    output_path='ltx2.3_twostage.mp4',
-    fps=24,
+    output_path='ltx2.3_twostage_a2v.mp4',
+    fps=frame_rate,
     audio_sample_rate=pipe.audio_vocoder.output_sampling_rate,
 )