diff --git a/tests/pipelines/ltx2/dfr_dummies.py b/tests/pipelines/ltx2/dfr_dummies.py deleted file mode 100644 index bd395ededfec..000000000000 --- a/tests/pipelines/ltx2/dfr_dummies.py +++ /dev/null @@ -1,196 +0,0 @@ -# Copyright 2026 The HuggingFace Team. -# -# 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 torch -from transformers import AutoTokenizer, Gemma3ForConditionalGeneration - -from diffusers import ( - AutoencoderKLLTX2Audio, - AutoencoderKLLTX2Video, - FlowMatchEulerDiscreteScheduler, - LTX2VideoTransformer3DModel, - LTXEulerAncestralRFScheduler, -) -from diffusers.pipelines.ltx2 import LTX2LatentUpsamplerModel, LTX2TextConnectors -from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder - - -BASE_TEXT_ENCODER_CKPT_ID = "hf-internal-testing/tiny-gemma3" - - -def get_dfr_dummy_components(*, spatial_upsampler: bool = False, temporal_upsampler: bool = False): - tokenizer = AutoTokenizer.from_pretrained(BASE_TEXT_ENCODER_CKPT_ID) - text_encoder = Gemma3ForConditionalGeneration.from_pretrained(BASE_TEXT_ENCODER_CKPT_ID) - - torch.manual_seed(0) - transformer = LTX2VideoTransformer3DModel( - in_channels=4, - out_channels=4, - patch_size=1, - patch_size_t=1, - num_attention_heads=2, - attention_head_dim=8, - cross_attention_dim=16, - audio_in_channels=4, - audio_out_channels=4, - audio_num_attention_heads=2, - audio_attention_head_dim=4, - audio_cross_attention_dim=8, - num_layers=2, - qk_norm="rms_norm_across_heads", - caption_channels=text_encoder.config.text_config.hidden_size, - rope_double_precision=False, - rope_type="split", - vae_scale_factors=(2, 2, 2), - use_keyframes_abs_pos_embedding=True, - ) - torch.nn.init.normal_(transformer.keyframes_abs_pos_embedding, std=0.1) - - torch.manual_seed(0) - connectors = LTX2TextConnectors( - caption_channels=text_encoder.config.text_config.hidden_size, - text_proj_in_factor=text_encoder.config.text_config.num_hidden_layers + 1, - video_connector_num_attention_heads=4, - video_connector_attention_head_dim=8, - video_connector_num_layers=1, - video_connector_num_learnable_registers=None, - audio_connector_num_attention_heads=4, - audio_connector_attention_head_dim=8, - audio_connector_num_layers=1, - audio_connector_num_learnable_registers=None, - connector_rope_base_seq_len=32, - rope_theta=10000.0, - rope_double_precision=False, - causal_temporal_positioning=False, - rope_type="split", - ) - - torch.manual_seed(0) - vae = AutoencoderKLLTX2Video( - in_channels=3, - out_channels=3, - latent_channels=4, - block_out_channels=(8,), - decoder_block_out_channels=(8,), - layers_per_block=(1,), - decoder_layers_per_block=(1, 1), - spatio_temporal_scaling=(True,), - decoder_spatio_temporal_scaling=(True,), - decoder_inject_noise=(False, False), - downsample_type=("spatial",), - upsample_residual=(False,), - upsample_factor=(1,), - timestep_conditioning=False, - patch_size=1, - patch_size_t=1, - encoder_causal=True, - decoder_causal=False, - ) - vae.use_framewise_encoding = False - vae.use_framewise_decoding = False - - torch.manual_seed(0) - audio_vae = AutoencoderKLLTX2Audio( - base_channels=4, - output_channels=2, - ch_mult=(1,), - num_res_blocks=1, - attn_resolutions=None, - in_channels=2, - resolution=32, - latent_channels=2, - norm_type="pixel", - causality_axis="height", - dropout=0.0, - mid_block_add_attention=False, - sample_rate=16000, - mel_hop_length=160, - is_causal=True, - mel_bins=8, - ) - - torch.manual_seed(0) - vocoder = LTX2Vocoder( - in_channels=audio_vae.config.output_channels * audio_vae.config.mel_bins, - hidden_channels=32, - out_channels=2, - upsample_kernel_sizes=[4, 4], - upsample_factors=[2, 2], - resnet_kernel_sizes=[3], - resnet_dilations=[[1, 3, 5]], - leaky_relu_negative_slope=0.1, - output_sampling_rate=16000, - ) - - components = { - "transformer": transformer, - "vae": vae, - "audio_vae": audio_vae, - "scheduler": FlowMatchEulerDiscreteScheduler(), - "text_encoder": text_encoder, - "tokenizer": tokenizer, - "connectors": connectors, - "vocoder": vocoder, - "processor": None, - "prompt_enhancer": None, - "duration_head": None, - } - if spatial_upsampler: - torch.manual_seed(0) - components["latent_upsampler"] = LTX2LatentUpsamplerModel( - in_channels=4, - mid_channels=32, - num_blocks_per_stage=1, - dims=3, - spatial_upsample=True, - temporal_upsample=False, - use_rational_resampler=False, - ) - if temporal_upsampler: - torch.manual_seed(0) - components["temporal_latent_upsampler"] = LTX2LatentUpsamplerModel( - in_channels=4, - mid_channels=32, - num_blocks_per_stage=1, - dims=3, - spatial_upsample=False, - temporal_upsample=True, - ) - return components - - -def get_temporal_dummy_components(): - components = get_dfr_dummy_components(temporal_upsampler=True) - components["scheduler"] = LTXEulerAncestralRFScheduler(eta=0.5) - # A refine round never enhances a prompt -- stage 1 already did, and re-enhancing would - # denoise the canvas under a different prompt than the one that generated it. - for name in ("duration_head", "processor", "prompt_enhancer"): - components.pop(name) - return components - - -def get_dfr_dummy_inputs(**overrides): - inputs = { - "prompt": "a robot dancing", - "height": 32, - "width": 32, - "num_frames": 9, - "frame_rate": 25.0, - "sigmas": [1.0, 0.5], - "use_cross_timestep": False, - "max_sequence_length": 16, - "output_type": "pt", - } - inputs.update(overrides) - return inputs diff --git a/tests/pipelines/ltx2/test_ltx2_diffusion_decode.py b/tests/pipelines/ltx2/test_ltx2_diffusion_decode.py index 46dba671b255..89806c78f946 100644 --- a/tests/pipelines/ltx2/test_ltx2_diffusion_decode.py +++ b/tests/pipelines/ltx2/test_ltx2_diffusion_decode.py @@ -16,13 +16,13 @@ import torch from diffusers import ( - AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler, LTX2VideoDiffusionDecodePipeline, LTX2VideoDiffusionDecoderModel, ) from ...testing_utils import enable_full_determinism, torch_device +from .testing_utils import get_dummy_vae enable_full_determinism() @@ -56,25 +56,10 @@ def _build(with_vae: bool = False): vae = None if with_vae: torch.manual_seed(0) + # Wider latents than the shared default, to match the decoder, and the causal decoder this pipeline + # is exercised with. vae = ( - AutoencoderKLLTX2Video( - in_channels=3, - out_channels=3, - latent_channels=DECODER_CONFIG["latent_channels"], - block_out_channels=(8,), - decoder_block_out_channels=(8,), - layers_per_block=(1,), - decoder_layers_per_block=(1, 1), - spatio_temporal_scaling=(True,), - decoder_spatio_temporal_scaling=(True,), - decoder_inject_noise=(False, False), - downsample_type=("spatial",), - upsample_residual=(False,), - upsample_factor=(1,), - timestep_conditioning=False, - patch_size=1, - patch_size_t=1, - ) + get_dummy_vae(latent_channels=DECODER_CONFIG["latent_channels"], decoder_causal=True) .to(torch_device) .eval() ) diff --git a/tests/pipelines/ltx2/test_ltx2_image2video.py b/tests/pipelines/ltx2/test_ltx2_image2video.py index 2a640d9a2516..599555384140 100644 --- a/tests/pipelines/ltx2/test_ltx2_image2video.py +++ b/tests/pipelines/ltx2/test_ltx2_image2video.py @@ -17,7 +17,6 @@ from diffusers import LTX2ImageToVideoPipeline from diffusers.pipelines.ltx2 import LTX2LatentUpsamplePipeline -from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel from ...testing_utils import assert_tensors_close, enable_full_determinism, torch_device from ..testing_utils import PipelineTesterMixin @@ -26,6 +25,7 @@ LTX2LoraMemoryTesterMixin, LTX2LoraTesterMixin, LTX2MemoryTesterMixin, + get_dummy_latent_upsampler, ) @@ -48,13 +48,6 @@ class LTX2ImageToVideoPipelineTesterConfig(LTX2BaseTesterConfig): ) batch_input_params = frozenset(["prompt", "negative_prompt", "image"]) - def get_dummy_upsample_component(self, in_channels=4, mid_channels=32, num_blocks_per_stage=1): - return LTX2LatentUpsamplerModel( - in_channels=in_channels, - mid_channels=mid_channels, - num_blocks_per_stage=num_blocks_per_stage, - ) - def get_dummy_inputs(self): generator = self.get_generator(0) image = torch.rand((1, 3, 32, 32), generator=generator) @@ -183,7 +176,7 @@ def test_two_stages_inference_with_upsampler(self): assert audio_latent.shape == (1, 2, 5, 2) assert audio_latent.shape[1] == pipe.vocoder.config.out_channels - upsampler = self.get_dummy_upsample_component(in_channels=video_latent.shape[1]) + upsampler = get_dummy_latent_upsampler(in_channels=video_latent.shape[1]) upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=upsampler) upscaled_video_latent = upsample_pipe(latents=video_latent, output_type="latent", return_dict=False)[0] assert upscaled_video_latent.shape == (1, 4, 3, 32, 32) diff --git a/tests/pipelines/ltx2/test_pipeline_ltx2_dfr.py b/tests/pipelines/ltx2/test_pipeline_ltx2_dfr.py index 792436b07f3e..f33a72d9bff6 100644 --- a/tests/pipelines/ltx2/test_pipeline_ltx2_dfr.py +++ b/tests/pipelines/ltx2/test_pipeline_ltx2_dfr.py @@ -21,8 +21,8 @@ from diffusers.pipelines.ltx2.utils import EPILOGUE_KEYFRAME_STRENGTH, MAX_CONDITIONING_FPS, trim_canvas from ...testing_utils import enable_full_determinism, torch_device -from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin -from .dfr_dummies import get_dfr_dummy_components, get_dfr_dummy_inputs +from ..testing_utils import BasePipelineTesterConfig, PipelineTesterMixin +from .testing_utils import LTX2MemoryTesterMixin, get_dfr_dummy_components, get_dfr_dummy_inputs enable_full_determinism() @@ -35,11 +35,18 @@ class LTX2DFRPipelineTesterConfig(BasePipelineTesterConfig): ) batch_input_params = frozenset(["prompt"]) output_shape = (9, 3, 32, 32) - optional_input_params = BasePipelineTesterConfig.optional_input_params - { - "num_inference_steps", - "num_images_per_prompt", - "latents", - } + # DFR is a video pipeline (`num_videos_per_prompt`, not `num_images_per_prompt`), takes a second latent input + # for the audio stream, and schedules its passes with an explicit `sigmas` list instead of `num_inference_steps`. + optional_input_params = frozenset( + [ + "num_videos_per_prompt", + "generator", + "latents", + "audio_latents", + "output_type", + "return_dict", + ] + ) def get_dummy_components(self): return get_dfr_dummy_components() @@ -59,13 +66,13 @@ def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff= ) def test_padded_canvas_is_trimmed_back_to_the_request(self): - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() inputs["num_frames"] = 11 assert pipe(**inputs).frames.shape[1] == 11 def test_latent_output_keeps_the_padded_canvas(self): - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() inputs["num_frames"] = 11 inputs["output_type"] = "latent" @@ -74,7 +81,7 @@ def test_latent_output_keeps_the_padded_canvas(self): assert trimmed.shape[2] < output.frames.shape[2] def test_the_pass_conditions_at_the_snapped_fps(self): - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) captured = [] original = pipe.prepare_latents @@ -104,7 +111,7 @@ def capture(**kwargs): assert torch.allclose(coords[:1, :, : expected.shape[2]], expected) def test_partially_conditioned_keyframe_starts_from_its_clean_content(self): - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) keyframe = torch.randn(1, 4, 1, 16, 16, device=torch_device) latents, conditioning_mask, clean_latents, _, _, _ = pipe.prepare_latents( keyframe_latents=[(8, keyframe, 0.95)], @@ -123,7 +130,7 @@ def test_partially_conditioned_keyframe_starts_from_its_clean_content(self): assert torch.allclose(latents[:, block], packed_keyframe * 0.95, atol=1e-6) def test_public_latents_are_normalized_on_the_way_in(self): - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) raw = torch.randn(1, 4, 5, 16, 16, device=torch_device) packed_raw = pipe._pack_latents(raw) _, _, clean, _, _, _ = pipe.prepare_latents( @@ -141,11 +148,10 @@ def test_public_latents_are_normalized_on_the_way_in(self): assert not torch.allclose(clean[:, : packed_raw.shape[1]], packed_raw) def test_keyframe_marker_reaches_the_transformer(self): - components = self.get_dummy_components() - pipe = self.get_pipeline(**components).to(torch_device) + pipe = self.get_pipeline().to(torch_device) marked = pipe(**self.get_dummy_inputs()).frames with torch.no_grad(): - components["transformer"].keyframes_abs_pos_embedding.zero_() + pipe.transformer.keyframes_abs_pos_embedding.zero_() unmarked = pipe(**self.get_dummy_inputs()).frames assert not torch.allclose(marked, unmarked) @@ -205,7 +211,7 @@ def _epilogue_plan(self, pipe, keyframe_positions=(8, 16, 24, 32, 40, 48, 56, 64 return tiles, video_coords, plan, len(keyframe_positions), latent_frames * latent_height * latent_width def test_a_tiled_epilogue_pass_routes_every_token_with_unit_total_weight(self): - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) _, video_coords, plan, _, _ = self._epilogue_plan(pipe) totals = torch.zeros(video_coords.shape[2], device=video_coords.device) for tile in plan: @@ -213,7 +219,7 @@ def test_a_tiled_epilogue_pass_routes_every_token_with_unit_total_weight(self): assert torch.allclose(totals, torch.ones_like(totals), atol=1e-6) def test_a_keyframe_two_epilogue_windows_share_is_a_single_token(self): - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) tiles, _, plan, num_keyframes, first_keyframe_token = self._epilogue_plan(pipe) assert len({(frames.start, frames.stop) for frames, _, _, _ in tiles}) > 1 tokens_per_keyframe = (32 // pipe.vae_spatial_compression_ratio) ** 2 @@ -230,7 +236,7 @@ def test_a_keyframe_two_epilogue_windows_share_is_a_single_token(self): assert shared, "no keyframe token is shared across windows" def test_the_epilogue_is_given_its_keyframes_rather_than_regenerating_them(self): - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) keyframe = torch.randn(1, 4, 1, 16, 16, device=torch_device) positions = [8, 16, 24, 32, 40, 48, 56, 64] guidance = torch.cat([keyframe] * len(positions), dim=2) @@ -258,7 +264,7 @@ def test_the_epilogue_is_given_its_keyframes_rather_than_regenerating_them(self) def test_composed_epilogue_pins_guidance_and_tiles(self): # `__call__` takes the tiling, not a resolved token plan: the plan needs the RoPE coordinates # `prepare_latents` builds inside the call, so a caller cannot produce one that is guaranteed to match. - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) keyframe = torch.randn(1, 4, 1, 16, 16, device=torch_device) output = pipe( **self.get_dummy_inputs(), @@ -281,7 +287,7 @@ def test_composed_epilogue_pins_guidance_and_tiles(self): def test_a_single_call_tile_reproduces_the_untiled_call(self): # The whole canvas as one tile with unit weights must be a no-op, which is what pins the token plan # `__call__` resolves against its own coordinates. - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) keyframe = torch.randn(1, 4, 1, 16, 16, device=torch_device) shared = { "generate_slots": False, @@ -321,7 +327,7 @@ def test_tile_audio_is_the_stage_1_window_on_the_playback_clock(self, pixel_star assert torch.allclose(window.flatten(), expected) def test_a_single_tile_plan_reproduces_the_untiled_pass(self): - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) pipe._interrupt = False pipe._current_timestep = None pipe._attention_kwargs = None @@ -396,7 +402,7 @@ def run(video_tile_plan): def test_distilled_euler_keeps_the_scheduler_step(self): # Stage 1 / 2 stay on FlowMatch Euler. Re-pinning after every step is ancestral-only; doing it here # would snap IC-LoRA reference tokens and first-frame anchors every step and change the canvas. - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) pipe._interrupt = False pipe._current_timestep = None pipe._attention_kwargs = None @@ -446,7 +452,7 @@ def zeros_step(model_output, timestep, sample, **kwargs): assert torch.allclose(out, torch.zeros_like(out)) def test_ancestral_step_does_not_erode_conditioning(self): - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) pipe.scheduler = LTXEulerAncestralRFScheduler(eta=0.5) pipe._interrupt = False pipe._current_timestep = None @@ -498,7 +504,7 @@ def test_ancestral_step_does_not_erode_conditioning(self): assert cos > 0.9, f"anchor tokens drifted from their conditioned content (cos={cos:.3f})" def test_the_epilogue_keeps_every_batch_element_distinct(self): - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) rebuilt = pipe.rebuild_epilogue_keyframes( torch.randn(2, 4, 2, 8, 8, device=torch_device), decode_timestep=0.0, @@ -511,12 +517,5 @@ def test_the_epilogue_keeps_every_batch_element_distinct(self): assert not torch.allclose(rebuilt[0], rebuilt[1]) -class TestLTX2DFRPipelineMemory(LTX2DFRPipelineTesterConfig, MemoryTesterMixin): - @pytest.mark.skip( - "Pre-existing for the whole LTX-2 family, not DFR-specific: the shared harness group-offloads only " - "`text_encoder` / `transformer` and moves `vae`, leaving the LTX-2-specific `connectors` on the CPU while it " - "receives accelerator tensors from the offloaded text encoder. Verified to fail identically on the stock " - "`LTX2Pipeline`. `test_pipeline_level_group_offloading_inference`, which offloads every component, passes." - ) - def test_group_offloading_inference(self): - pass +class TestLTX2DFRPipelineMemory(LTX2DFRPipelineTesterConfig, LTX2MemoryTesterMixin): + pass diff --git a/tests/pipelines/ltx2/test_pipeline_ltx2_dfr_temporal_refine.py b/tests/pipelines/ltx2/test_pipeline_ltx2_dfr_temporal_refine.py index 213ff47e345b..5cdeeb9b548e 100644 --- a/tests/pipelines/ltx2/test_pipeline_ltx2_dfr_temporal_refine.py +++ b/tests/pipelines/ltx2/test_pipeline_ltx2_dfr_temporal_refine.py @@ -28,7 +28,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import BasePipelineTesterConfig, PipelineTesterMixin -from .dfr_dummies import get_dfr_dummy_components, get_dfr_dummy_inputs, get_temporal_dummy_components +from .testing_utils import get_dfr_dummy_components, get_dfr_dummy_inputs, get_temporal_dummy_components enable_full_determinism() @@ -41,11 +41,9 @@ class LTX2DFRTemporalRefinePipelineTesterConfig(BasePipelineTesterConfig): ) batch_input_params = frozenset(["prompt"]) output_shape = (17, 3, 32, 32) - optional_input_params = BasePipelineTesterConfig.optional_input_params - { - "num_inference_steps", - "num_images_per_prompt", - "latents", - } + # A refine round takes one canvas at a time, so there is no per-prompt multiplier; it takes a second latent + # input for the audio stream and schedules its passes with an explicit `sigmas` list, not `num_inference_steps`. + optional_input_params = frozenset(["generator", "latents", "audio_latents", "output_type", "return_dict"]) def get_dummy_components(self): return get_temporal_dummy_components() @@ -80,7 +78,7 @@ def test_inference_batch_single_identical(self, *args, **kwargs): pass def test_temporal_upsample_round_doubles_the_frame_count(self): - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) output = pipe(**self.get_dummy_inputs()) assert output.frames.shape[1] == (9 - 1) * 2 + 1 assert pipe.num_timesteps == 2 @@ -104,7 +102,7 @@ def test_a_non_ancestral_scheduler_is_refused(self, scheduler, message): def test_temporal_round_tiles_get_distinct_ancestral_noise(self): dfr = LTX2DFRPipeline(**get_dfr_dummy_components()).to(torch_device) stage = dfr(**get_dfr_dummy_inputs(generator=self.get_generator(0), num_frames=17, output_type="latent")) - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) seeds = [] original = pipe.denoise @@ -147,7 +145,7 @@ def test_a_condition_keeps_its_moment_through_the_refine_round(self): conditions=[LTX2VideoCondition(frames=image, index=4, strength=1.0, crf=0)], ) ) - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) coords = [] original = pipe.prepare_latents @@ -207,7 +205,7 @@ def test_last_frame_condition_stays_at_the_end_across_a_temporal_round(self): conditions=LTX2VideoCondition(frames=frame, index=0, strength=1.0, crf=0), ) ) - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) shared = { "prompt": "a robot dancing", "height": 32, @@ -244,7 +242,7 @@ def test_last_frame_condition_stays_at_the_end_across_a_temporal_round(self): def test_a_carried_slot_is_the_copy_the_stitched_canvas_kept(self): dfr = LTX2DFRPipeline(**get_dfr_dummy_components()).to(torch_device) stage = dfr(**get_dfr_dummy_inputs(generator=self.get_generator(0), num_frames=17, output_type="latent")) - pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) slot_slices, anchors, denoised = [], [], [] original_prepare, original_denoise = pipe.prepare_latents, pipe.denoise diff --git a/tests/pipelines/ltx2/testing_utils.py b/tests/pipelines/ltx2/testing_utils.py index e9fd67c86cf5..89f544fdf0d3 100644 --- a/tests/pipelines/ltx2/testing_utils.py +++ b/tests/pipelines/ltx2/testing_utils.py @@ -15,8 +15,10 @@ Every LTX2 pipeline in this directory takes the same set of dummy sub-modules — they differ only in which optional components they accept and in what `__call__` takes — so the component builder lives here and the -per-pipeline configs subclass `LTX2BaseTesterConfig`. The two scoped tester mixins below carry the -pipeline-family-wide skips so each test file does not restate them. +per-pipeline configs subclass `LTX2BaseTesterConfig`. The DFR pipelines drive their fixtures through the +module-level `get_dfr_*` helpers instead, because their tests build extra pipelines mid-test rather than only +through a tester config. The scoped tester mixins below carry the pipeline-family-wide skips so each test file +does not restate them. """ import pytest @@ -28,8 +30,9 @@ AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler, LTX2VideoTransformer3DModel, + LTXEulerAncestralRFScheduler, ) -from diffusers.pipelines.ltx2 import LTX2DurationHead, LTX2TextConnectors +from diffusers.pipelines.ltx2 import LTX2DurationHead, LTX2LatentUpsamplerModel, LTX2TextConnectors from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder from ...testing_utils import torch_device @@ -41,6 +44,222 @@ ) +BASE_TEXT_ENCODER_CKPT_ID = "hf-internal-testing/tiny-gemma3" + +# Components the LTX2 pipelines accept but that these fast tests leave unset. Which ones exist differs per +# pipeline — only the pipelines that can predict a duration take a `duration_head`, and only some take an +# `audio_scheduler` — so each config lists its own set and `get_ltx2_dummy_components` fills them with `None`. +DEFAULT_UNSET_COMPONENTS = ("processor", "prompt_enhancer", "duration_head") + +# The DFR transformer places its keyframes with an absolute positional embedding the other LTX2 pipelines do not +# use, and it reads the VAE strides off its own config to map keyframe positions onto latent frames. +DFR_TRANSFORMER_KWARGS = {"vae_scale_factors": (2, 2, 2), "use_keyframes_abs_pos_embedding": True} + + +def get_dummy_vae(**overrides): + """The tiny `AutoencoderKLLTX2Video` every LTX2 test uses; `overrides` go straight to the constructor. + + Seeding is left to the caller so that each fixture keeps the RNG stream the hardcoded slices were recorded + against. + """ + kwargs = { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 4, + "block_out_channels": (8,), + "decoder_block_out_channels": (8,), + "layers_per_block": (1,), + "decoder_layers_per_block": (1, 1), + "spatio_temporal_scaling": (True,), + "decoder_spatio_temporal_scaling": (True,), + "decoder_inject_noise": (False, False), + "downsample_type": ("spatial",), + "upsample_residual": (False,), + "upsample_factor": (1,), + "timestep_conditioning": False, + "patch_size": 1, + "patch_size_t": 1, + "encoder_causal": True, + "decoder_causal": False, + } + kwargs.update(overrides) + + vae = AutoencoderKLLTX2Video(**kwargs) + vae.use_framewise_encoding = False + vae.use_framewise_decoding = False + return vae + + +def get_dummy_latent_upsampler(**overrides): + """The tiny `LTX2LatentUpsamplerModel`; `overrides` go straight to the constructor. + + Seeding is left to the caller, as in `get_dummy_vae`. + """ + kwargs = {"in_channels": 4, "mid_channels": 32, "num_blocks_per_stage": 1} + kwargs.update(overrides) + + return LTX2LatentUpsamplerModel(**kwargs) + + +def get_ltx2_dummy_components(*, unset_components=DEFAULT_UNSET_COMPONENTS, transformer_kwargs=None): + """The dummy sub-modules shared by every LTX2 pipeline. + + `transformer_kwargs` overrides the transformer config for the variants that need a different one (see + `DFR_TRANSFORMER_KWARGS`); `unset_components` names the optional components to hand back as `None`. + """ + tokenizer = AutoTokenizer.from_pretrained(BASE_TEXT_ENCODER_CKPT_ID) + text_encoder = Gemma3ForConditionalGeneration.from_pretrained(BASE_TEXT_ENCODER_CKPT_ID) + + torch.manual_seed(0) + transformer = LTX2VideoTransformer3DModel( + in_channels=4, + out_channels=4, + patch_size=1, + patch_size_t=1, + num_attention_heads=2, + attention_head_dim=8, + cross_attention_dim=16, + audio_in_channels=4, + audio_out_channels=4, + audio_num_attention_heads=2, + audio_attention_head_dim=4, + audio_cross_attention_dim=8, + num_layers=2, + qk_norm="rms_norm_across_heads", + caption_channels=text_encoder.config.text_config.hidden_size, + rope_double_precision=False, + rope_type="split", + **(transformer_kwargs or {}), + ) + if transformer.config.use_keyframes_abs_pos_embedding: + # Registered zero-initialized, so without this the keyframe marker would be invisible to the transformer + # and the tests that assert it reaches the forward pass could not tell it apart from an unmarked run. + torch.nn.init.normal_(transformer.keyframes_abs_pos_embedding, std=0.1) + + torch.manual_seed(0) + connectors = LTX2TextConnectors( + caption_channels=text_encoder.config.text_config.hidden_size, + text_proj_in_factor=text_encoder.config.text_config.num_hidden_layers + 1, + video_connector_num_attention_heads=4, + video_connector_attention_head_dim=8, + video_connector_num_layers=1, + video_connector_num_learnable_registers=None, + audio_connector_num_attention_heads=4, + audio_connector_attention_head_dim=8, + audio_connector_num_layers=1, + audio_connector_num_learnable_registers=None, + connector_rope_base_seq_len=32, + rope_theta=10000.0, + rope_double_precision=False, + causal_temporal_positioning=False, + rope_type="split", + ) + + torch.manual_seed(0) + vae = get_dummy_vae() + + torch.manual_seed(0) + audio_vae = AutoencoderKLLTX2Audio( + base_channels=4, + output_channels=2, + ch_mult=(1,), + num_res_blocks=1, + attn_resolutions=None, + in_channels=2, + resolution=32, + latent_channels=2, + norm_type="pixel", + causality_axis="height", + dropout=0.0, + mid_block_add_attention=False, + sample_rate=16000, + mel_hop_length=160, + is_causal=True, + mel_bins=8, + ) + + torch.manual_seed(0) + vocoder = LTX2Vocoder( + in_channels=audio_vae.config.output_channels * audio_vae.config.mel_bins, + hidden_channels=32, + out_channels=2, + upsample_kernel_sizes=[4, 4], + upsample_factors=[2, 2], + resnet_kernel_sizes=[3], + resnet_dilations=[[1, 3, 5]], + leaky_relu_negative_slope=0.1, + output_sampling_rate=16000, + ) + + return { + "transformer": transformer, + "vae": vae, + "audio_vae": audio_vae, + "scheduler": FlowMatchEulerDiscreteScheduler(), + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "connectors": connectors, + "vocoder": vocoder, + **dict.fromkeys(unset_components), + } + + +def get_dummy_duration_head(): + """A tiny `LTX2DurationHead`, for the pipelines that accept one (`duration_head` in `unset_components`).""" + torch.manual_seed(0) + # The dummy connectors emit 4 heads * 8 head_dim = 32 wide output for both streams. + return LTX2DurationHead( + video_cross_attention_dim=32, + audio_cross_attention_dim=32, + pooler_hidden_dim=8, + num_queries=1, + num_pooler_heads=2, + mlp_hidden_dim=8, + ) + + +def get_dfr_dummy_components(*, spatial_upsampler: bool = False, temporal_upsampler: bool = False): + """The shared LTX2 components with the DFR transformer, plus the upsamplers a multi-stage run needs.""" + components = get_ltx2_dummy_components(transformer_kwargs=DFR_TRANSFORMER_KWARGS) + if spatial_upsampler: + torch.manual_seed(0) + components["latent_upsampler"] = get_dummy_latent_upsampler( + dims=3, spatial_upsample=True, temporal_upsample=False, use_rational_resampler=False + ) + if temporal_upsampler: + torch.manual_seed(0) + components["temporal_latent_upsampler"] = get_dummy_latent_upsampler( + dims=3, spatial_upsample=False, temporal_upsample=True + ) + return components + + +def get_temporal_dummy_components(): + components = get_dfr_dummy_components(temporal_upsampler=True) + components["scheduler"] = LTXEulerAncestralRFScheduler(eta=0.5) + # A refine round never enhances a prompt -- stage 1 already did, and re-enhancing would + # denoise the canvas under a different prompt than the one that generated it. + for name in ("duration_head", "processor", "prompt_enhancer"): + components.pop(name) + return components + + +def get_dfr_dummy_inputs(**overrides): + inputs = { + "prompt": "a robot dancing", + "height": 32, + "width": 32, + "num_frames": 9, + "frame_rate": 25.0, + "sigmas": [1.0, 0.5], + "use_cross_timestep": False, + "max_sequence_length": 16, + "output_type": "pt", + } + inputs.update(overrides) + return inputs + + class LTX2BaseTesterConfig(BasePipelineTesterConfig): """Dummy component set shared by every LTX2 pipeline in this directory.""" @@ -59,140 +278,17 @@ class LTX2BaseTesterConfig(BasePipelineTesterConfig): ) output_shape = (5, 3, 32, 32) - base_text_encoder_ckpt_id = "hf-internal-testing/tiny-gemma3" + base_text_encoder_ckpt_id = BASE_TEXT_ENCODER_CKPT_ID - # Components the pipeline accepts but that these fast tests leave unset. Which ones exist differs per - # pipeline — only the pipelines that can predict a duration take a `duration_head`, and only some take an - # `audio_scheduler` — so each config lists its own set and `get_dummy_components` fills them with `None`. - unset_components = ("processor", "prompt_enhancer", "duration_head") + # See `DEFAULT_UNSET_COMPONENTS`; each config narrows or widens this to the components its pipeline accepts. + unset_components = DEFAULT_UNSET_COMPONENTS def get_dummy_components(self): - tokenizer = AutoTokenizer.from_pretrained(self.base_text_encoder_ckpt_id) - text_encoder = Gemma3ForConditionalGeneration.from_pretrained(self.base_text_encoder_ckpt_id) - - torch.manual_seed(0) - transformer = LTX2VideoTransformer3DModel( - in_channels=4, - out_channels=4, - patch_size=1, - patch_size_t=1, - num_attention_heads=2, - attention_head_dim=8, - cross_attention_dim=16, - audio_in_channels=4, - audio_out_channels=4, - audio_num_attention_heads=2, - audio_attention_head_dim=4, - audio_cross_attention_dim=8, - num_layers=2, - qk_norm="rms_norm_across_heads", - caption_channels=text_encoder.config.text_config.hidden_size, - rope_double_precision=False, - rope_type="split", - ) - - torch.manual_seed(0) - connectors = LTX2TextConnectors( - caption_channels=text_encoder.config.text_config.hidden_size, - text_proj_in_factor=text_encoder.config.text_config.num_hidden_layers + 1, - video_connector_num_attention_heads=4, - video_connector_attention_head_dim=8, - video_connector_num_layers=1, - video_connector_num_learnable_registers=None, - audio_connector_num_attention_heads=4, - audio_connector_attention_head_dim=8, - audio_connector_num_layers=1, - audio_connector_num_learnable_registers=None, - connector_rope_base_seq_len=32, - rope_theta=10000.0, - rope_double_precision=False, - causal_temporal_positioning=False, - rope_type="split", - ) - - torch.manual_seed(0) - vae = AutoencoderKLLTX2Video( - in_channels=3, - out_channels=3, - latent_channels=4, - block_out_channels=(8,), - decoder_block_out_channels=(8,), - layers_per_block=(1,), - decoder_layers_per_block=(1, 1), - spatio_temporal_scaling=(True,), - decoder_spatio_temporal_scaling=(True,), - decoder_inject_noise=(False, False), - downsample_type=("spatial",), - upsample_residual=(False,), - upsample_factor=(1,), - timestep_conditioning=False, - patch_size=1, - patch_size_t=1, - encoder_causal=True, - decoder_causal=False, - ) - vae.use_framewise_encoding = False - vae.use_framewise_decoding = False - - torch.manual_seed(0) - audio_vae = AutoencoderKLLTX2Audio( - base_channels=4, - output_channels=2, - ch_mult=(1,), - num_res_blocks=1, - attn_resolutions=None, - in_channels=2, - resolution=32, - latent_channels=2, - norm_type="pixel", - causality_axis="height", - dropout=0.0, - mid_block_add_attention=False, - sample_rate=16000, - mel_hop_length=160, - is_causal=True, - mel_bins=8, - ) - - torch.manual_seed(0) - vocoder = LTX2Vocoder( - in_channels=audio_vae.config.output_channels * audio_vae.config.mel_bins, - hidden_channels=32, - out_channels=2, - upsample_kernel_sizes=[4, 4], - upsample_factors=[2, 2], - resnet_kernel_sizes=[3], - resnet_dilations=[[1, 3, 5]], - leaky_relu_negative_slope=0.1, - output_sampling_rate=16000, - ) - - scheduler = FlowMatchEulerDiscreteScheduler() - - return { - "transformer": transformer, - "vae": vae, - "audio_vae": audio_vae, - "scheduler": scheduler, - "text_encoder": text_encoder, - "tokenizer": tokenizer, - "connectors": connectors, - "vocoder": vocoder, - **dict.fromkeys(self.unset_components), - } + return get_ltx2_dummy_components(unset_components=self.unset_components) def get_dummy_duration_head(self): """A tiny `LTX2DurationHead`, for the pipelines that accept one (`duration_head` in `unset_components`).""" - torch.manual_seed(0) - # The dummy connectors emit 4 heads * 8 head_dim = 32 wide output for both streams. - return LTX2DurationHead( - video_cross_attention_dim=32, - audio_cross_attention_dim=32, - pooler_hidden_dim=8, - num_queries=1, - num_pooler_heads=2, - mlp_hidden_dim=8, - ) + return get_dummy_duration_head() def get_pipeline_with_duration_head(self): components = self.get_dummy_components() @@ -203,9 +299,11 @@ def get_pipeline_with_duration_head(self): class LTX2MemoryTesterMixin(MemoryTesterMixin): """`MemoryTesterMixin` for the LTX2 pipelines in this directory.""" - # The shared helper only offloads a fixed set of component names and leaves LTX2's extra module components - # (`connectors`, `audio_vae`, `vocoder`) on CPU, so the forward pass mixes devices. Pipeline-level - # offloading, which walks every component, is exercised by `test_pipeline_level_group_offloading_inference`. + # The shared helper only group-offloads `text_encoder` / `transformer` and moves `vae`, leaving LTX2's extra + # module components (`connectors`, `audio_vae`, `vocoder`) on the CPU while they receive accelerator tensors + # from the offloaded text encoder, so the forward pass mixes devices. This is pre-existing for the whole LTX2 + # family rather than specific to any one pipeline. Pipeline-level offloading, which walks every component, is + # exercised by `test_pipeline_level_group_offloading_inference`. @pytest.mark.skip("Using test_pipeline_level_group_offloading_inference instead") def test_group_offloading_inference(self): pass