AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-24 20:59:27 +03:00
Code Activity

[tests] refactor ltx2 dfr pipeline tests (#14655)

Browse Source 
* Add LTX-2.5 DFR pipeline with generated keyframe slots

Ports DFRPipeline from the Lightricks reference. Stage 1 generates video plus
single-pixel-frame keyframe slots at a fraction of the requested resolution on
a VAE-aligned segment grid; both are spatially latent-upsampled and stage 2
re-denoises at twice that resolution with the slots re-attached and an optional
spatial detailing IC-LoRA active for that stage only. Optional temporal x2/x4
refine rounds tile the canvas at shared keyframes and densify with ancestral
Euler. With spatial_upscalings=2 a full-resolution detailing epilogue follows
the rounds.

The transformer already stored keyframes_abs_pos_embedding for load/save; this
wires it into the forward through a new video_keyframes_mask argument, which
only a keyframes-aware pipeline passes, so other pipelines are unaffected.

The epilogue denoises the whole canvas in one loop and tiles the transformer
call inside it, so every Euler step steps a canvas whose tiles have already
agreed on their overlaps. Spatial tiles blend under a trapezoidal mask, since
neither side of a height or width border holds a known answer. Temporal tiles
are cut on the keyframe seams the last refine round stitched on: both windows
reproduce a shared keyframe there, so the later one drops its run-up under a
rectangular mask rather than averaging it. Conditionings are attached once on
the whole canvas and filtered per tile at the token level, and a keyframe two
windows share is one token they both read.

The epilogue is handed its keyframes rather than asked to generate them. Each
carry plane is decoded on its own -- the VAE is causal, so a stacked decode
would bleed neighbours -- then Lanczos-stretched x2 in RGB and encoded again at
the output resolution, and pinned fully clean. Only the video latent is
spatially upsampled.

Conditioning fps is snapped to 60 above 30 rather than merely capped there, at
every stage. RoPE time is pixel_frame / fps, and the transformer is trained
around 24/25/30 and 60; a temporal round taking 24 fps to 48 lands between
those, and it shows as stutter at the latent borders. Playback fps is
unchanged, so 24 fps with one round still ships 48 fps.

A condition's index is read on the canvas num_frames asks for, and the moment
it names is carried onto each refine round's longer canvas by scaling its pixel
position by 2**round. The scaled position does not generally land on a latent
boundary, so it travels as a pixel index rather than through the public latent
index; a keyframe conditioning is appended as extra tokens instead of being
spliced into the base grid, so it does not need to.

height and width must be divisible by 2**spatial_upscalings times the VAE's
spatial compression ratio, which makes 4K 3840x2176 rather than 3840x2160. That
rule is checked ahead of the looser one every LTX-2 pipeline applies, so the
error names the divisor a DFR caller actually has to satisfy.

Four details are easy to get wrong, and each is covered by a test after showing
up as a visible seam at a tile handover:

- The ancestral step injects noise into every token, so the conditioning blend
  has to be re-applied afterwards. Skipping it lets the strength-0.95 anchor
  keyframes erode over the schedule, and those anchors are the only thing
  pinning adjacent tiles onto the same content.
- Velocity is converted to x0 with each token's own noise level, not the scalar
  schedule sigma: a token held at strength s sits at (1 - s) * sigma.
- Each tile draws its ancestral noise from a generator seeded
  seed + 1000 * round + tile, kept separate from the main generator so the
  draws do not consume state the next tile's initial noising reads.
- Two tiles invent the slot that falls in the later one's dropped lead-in. The
  stitch keeps the earlier tile's frames there, so the earlier tile's copy is
  the one the canvas holds, and the one the next round must anchor on.

* Split the DFR pipeline into composable stages

Addresses the review on #14567: use the same compose pattern as the other LTX
two-stage pipelines, get ancestral Euler from an existing scheduler, and put the
temporal rounds in their own pipeline so the schedule is not switched mid-call.

`LTX2DFRPipeline.__call__` is now one denoise pass at `height` x `width`. Callers
compose stage 1, `LTX2LatentUpsamplePipeline`, stage 2 and each temporal round,
and a documented recipe is the copy-paste 1080p path. The recipe knobs
(`spatial_upscalings`, `temporal_upscalings`, `detailing_lora_adapter_name`) and
the required upsampler components are gone; `height`/`width` are this pass, not
the final output.

`ancestral_euler_step` is replaced by `LTXEulerAncestralRFScheduler.step` plus a
re-application of the conditioning blend, which ancestral noise would otherwise
erode on the strength-0.95 seam anchors. The new `LTX2DFRTemporalRefinePipeline`
owns that scheduler and one round; stage 1, stage 2 and the epilogue stay on
`FlowMatchEulerDiscreteScheduler`. It refuses any other scheduler rather than
silently taking a deterministic step and returning a softer canvas.

Pack/unpack, `prepare_latents`, `denoise` and `encode_conditions` move to
`LTX2DFRCoreMixin` so neither pipeline subclasses the other. Public latents are
raw on both sides of every boundary, `output_type="latent"` returns the untrimmed
canvas so a slot on the pad is not dropped, and `trim_canvas` does the trim
before decode. `LTX2DFRPipelineOutput` carries `keyframes` and
`keyframe_positions`, which cannot be re-derived after a round.

Verified against the pre-split implementation: bit-exact on dummy components in
fp32 and bf16, with and without the IC-LoRA reference, over one and two rounds;
and within one bf16 ulp on the real checkpoint, where the only difference is that
the split normalizes upsampled latents in fp32 rather than bf16.

Also in this pass:

- `__call__` takes `video_tiles` (the `epilogue_tiles` layout) instead of a
  resolved token plan. Resolving one needs the RoPE coordinates that only exist
  once `prepare_latents` has run, so a caller could not build the plan at all.
- `rebuild_epilogue_keyframes` is public and takes and returns raw latents. The
  composed epilogue needs it, so it was public API in practice while named
  private, and its normalized return forced callers into `_denormalize_latents`.
- Drop prompt enhancement and `num_videos_per_prompt` from the temporal pipeline.
  Enhancement belongs to stage 1 -- re-running it would denoise the canvas under
  a different prompt than the one that generated it -- and the batch is set by
  the incoming latent canvas, so `num_videos_per_prompt > 1` only ever raised.
- Fix the docs recipe: `requested_frames` counted latent frames where
  `trim_canvas` wants pixel frames, truncating a 241-frame render to 25; the
  three pipelines share components, so place them together instead of offloading
  one and leaving `temporal_latent_upsampler` off the device; and the detailing
  IC-LoRA is applied at 0.5, the strength the reference hardcodes.

- Pass `crf=0` on the conditions in the refine-round test. The default CRF sends
  the image through H.264 re-compression, which needs PyAV, so the test failed on
  any environment without it while testing nothing about re-compression.

* Drop LTX2DFRCoreMixin in favor of # Copied from

Each DFR pipeline now owns pack/unpack, prepare_latents, and denoise so the
classes stay self-contained. Shared helpers copy from LTX2Pipeline /
LTX2ConditionPipeline, and DFR-specific methods copy from LTX2DFRPipeline.
dfr_core.py keeps only the constants and canvas helpers.

* inline _finalize_output

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* remove dfr_core.py

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* make return_dict=False return a tuple of 4, consistent with diffusers convention; update docs and pipeline output accordingly

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* small refactors to the ltx2 dfr testing suite

* save

* remove dfr_dummies.py.

* remove tests/pipelines/ltx2/test_ltx2_dfr_layout.py

---------

Co-authored-by: aartyomov <aartyomov@lightricks.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Co-authored-by: YiYi Xu <yixu310@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Sayak Paul
2026-09-01 09:17:10 +05:30
committed by GitHub
parent 61a11f2191
commit 77c13f8a49
6 changed files with 279 additions and 402 deletions
Download Patch File Download Diff File Expand all files Collapse all files

196
tests/pipelines/ltx2/dfr_dummies.py
View File

@@ -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

23
tests/pipelines/ltx2/test_ltx2_diffusion_decode.py
View File

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

11
tests/pipelines/ltx2/test_ltx2_image2video.py
View File

@@ -17,7 +17,6 @@ import torch
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 @@ from .testing_utils import (
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 @@ class TestLTX2ImageToVideoPipeline(LTX2ImageToVideoPipelineTesterConfig, Pipelin
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)

65
tests/pipelines/ltx2/test_pipeline_ltx2_dfr.py
View File

@@ -21,8 +21,8 @@ from diffusers.pipelines.ltx2.pipeline_ltx2_dfr_temporal_refine import _audio_wi
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 @@ class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin):
)
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 @@ class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin):
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 @@ class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin):
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 @@ class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin):
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 @@ class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin):
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 @@ class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin):
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 @@ class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin):
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 @@ class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin):
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 @@ class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin):
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 @@ class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin):
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 @@ class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin):
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 @@ class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin):
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 @@ class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin):
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 @@ class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin):
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 @@ class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin):
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

20
tests/pipelines/ltx2/test_pipeline_ltx2_dfr_temporal_refine.py
View File

@@ -28,7 +28,7 @@ from diffusers.pipelines.ltx2.utils import ANCHOR_KEYFRAME_STRENGTH
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 @@ class TestLTX2DFRTemporalRefinePipeline(LTX2DFRTemporalRefinePipelineTesterConfi
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 @@ class TestLTX2DFRTemporalRefinePipeline(LTX2DFRTemporalRefinePipelineTesterConfi
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 @@ class TestLTX2DFRTemporalRefinePipeline(LTX2DFRTemporalRefinePipelineTesterConfi
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 @@ class TestLTX2DFRTemporalRefinePipeline(LTX2DFRTemporalRefinePipelineTesterConfi
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 @@ class TestLTX2DFRTemporalRefinePipeline(LTX2DFRTemporalRefinePipelineTesterConfi
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

366
tests/pipelines/ltx2/testing_utils.py
View File

@@ -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 @@ from diffusers import (
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 @@ from ..testing_utils import (
)
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 @@ class LTX2BaseTesterConfig(BasePipelineTesterConfig):
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