[tests] refactor ltx2 dfr pipeline tests

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()
         )

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)