From 7464d230fa2cba83b3fd2eeb9e040a1da3f1f19d Mon Sep 17 00:00:00 2001 From: bghira Date: Thu, 6 Aug 2026 13:16:51 -0600 Subject: [PATCH 1/2] Fix LTX2 VAE up-block channel widths Compute LTX2 decoder up-block input channels from the current decoder width and project mismatched inputs to the upsampler width before upsampling. This keeps non-nominal constant-width decoder configurations from sending tensors with the wrong channel count into the 3D upsampler path. Add tests for projection, projection skipping, and a compact decoder shape that previously failed with mismatched widths. --- .../helpers/models/ltxvideo2/autoencoder.py | 9 +-- tests/test_ltxvideo2_autoencoder.py | 55 ++++++++++++++++++- 2 files changed, 59 insertions(+), 5 deletions(-) diff --git a/simpletuner/helpers/models/ltxvideo2/autoencoder.py b/simpletuner/helpers/models/ltxvideo2/autoencoder.py index c53141151..262464f65 100644 --- a/simpletuner/helpers/models/ltxvideo2/autoencoder.py +++ b/simpletuner/helpers/models/ltxvideo2/autoencoder.py @@ -611,7 +611,7 @@ def forward( return hidden_states -# Like LTXVideoUpBlock3d but with no conv_in and the updated LTX2VideoResnetBlock3d +# Like LTXVideoUpBlock3d but with the updated LTX2VideoResnetBlock3d class LTX2VideoUpBlock3d(nn.Module): r""" Up block used in the LTXVideo model. @@ -663,10 +663,11 @@ def __init__( self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(in_channels * 4, 0) self.conv_in = None - if in_channels != out_channels: + upsampler_in_channels = out_channels * upscale_factor + if in_channels != upsampler_in_channels: self.conv_in = LTX2VideoResnetBlock3d( in_channels=in_channels, - out_channels=out_channels, + out_channels=upsampler_in_channels, dropout=dropout, eps=resnet_eps, non_linearity=resnet_act_fn, @@ -986,7 +987,7 @@ def __init__( num_block_out_channels = len(block_out_channels) self.up_blocks = nn.ModuleList([]) for i in range(num_block_out_channels): - input_channel = output_channel // upsample_factor[i] + input_channel = output_channel output_channel = block_out_channels[i] // upsample_factor[i] up_block = LTX2VideoUpBlock3d( diff --git a/tests/test_ltxvideo2_autoencoder.py b/tests/test_ltxvideo2_autoencoder.py index 1433afacb..7b89f93b0 100644 --- a/tests/test_ltxvideo2_autoencoder.py +++ b/tests/test_ltxvideo2_autoencoder.py @@ -1,6 +1,8 @@ import unittest -from simpletuner.helpers.models.ltxvideo2.autoencoder import LTX2VideoUpBlock3d, LTX2VideoUpsampler3d +import torch + +from simpletuner.helpers.models.ltxvideo2.autoencoder import LTX2VideoDecoder3d, LTX2VideoUpBlock3d, LTX2VideoUpsampler3d class TestLTX2VideoAutoencoder(unittest.TestCase): @@ -27,6 +29,57 @@ def test_up_block_rejects_unknown_upsample_type(self): upsample_type="invalid", ) + def test_up_block_conv_in_projects_to_upsampler_width(self): + up_block = LTX2VideoUpBlock3d( + in_channels=48, + out_channels=32, + num_layers=1, + spatio_temporal_scale=True, + upscale_factor=2, + ) + + self.assertIsNotNone(up_block.conv_in) + + sample = torch.randn(1, 48, 4, 8, 8) + with torch.no_grad(): + output = up_block(sample, causal=False) + + self.assertEqual(output.shape, (1, 32, 7, 16, 16)) + + def test_up_block_skips_conv_in_when_input_matches_upsampler_width(self): + up_block = LTX2VideoUpBlock3d( + in_channels=64, + out_channels=32, + num_layers=1, + spatio_temporal_scale=True, + upscale_factor=2, + ) + + self.assertIsNone(up_block.conv_in) + + sample = torch.randn(1, 64, 4, 8, 8) + with torch.no_grad(): + output = up_block(sample, causal=False) + + self.assertEqual(output.shape, (1, 32, 7, 16, 16)) + + def test_decoder_accepts_non_nominal_constant_width_blocks(self): + decoder = LTX2VideoDecoder3d( + in_channels=4, + out_channels=3, + block_out_channels=(16, 16, 16), + layers_per_block=(1, 1, 1, 1), + patch_size=1, + patch_size_t=1, + inject_noise=(False, False, False, False), + ) + + latents = torch.randn(1, 4, 3, 4, 4) + with torch.no_grad(): + output = decoder(latents, causal=False) + + self.assertEqual(output.shape, (1, 3, 17, 32, 32)) + if __name__ == "__main__": unittest.main() From 5f610b1b36127ffd2f7c4c5625d1367a74760550 Mon Sep 17 00:00:00 2001 From: bghira Date: Thu, 6 Aug 2026 13:17:04 -0600 Subject: [PATCH 2/2] Cover LTX2 dynamic shift sequence length Add AST-based regression coverage that both LTX Video 2 text-to-video and image-to-video pipelines pass the actual video_sequence_length into calculate_shift. This guards against reverting to shape assumptions that undercount the packed video sequence when dynamic timestep shifting is enabled. --- tests/test_ltxvideo2_dynamic_shift.py | 38 +++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 tests/test_ltxvideo2_dynamic_shift.py diff --git a/tests/test_ltxvideo2_dynamic_shift.py b/tests/test_ltxvideo2_dynamic_shift.py new file mode 100644 index 000000000..a849ce247 --- /dev/null +++ b/tests/test_ltxvideo2_dynamic_shift.py @@ -0,0 +1,38 @@ +import ast +import inspect +import textwrap +import unittest + +from simpletuner.helpers.models.ltxvideo2 import pipeline_ltx2, pipeline_ltx2_image2video + + +def _ltx2_pipeline_call_ast(pipeline_cls): + source = inspect.getsource(pipeline_cls.__call__) + return ast.parse(textwrap.dedent(source)) + + +def _calculate_shift_calls(tree): + return [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "calculate_shift" + ] + + +class LTXVideo2DynamicShiftTests(unittest.TestCase): + def assert_uses_actual_video_sequence_length_for_shift(self, pipeline_cls): + tree = _ltx2_pipeline_call_ast(pipeline_cls) + calls = _calculate_shift_calls(tree) + + self.assertEqual(len(calls), 1) + self.assertEqual(ast.unparse(calls[0].args[0]), "video_sequence_length") + + def test_text_to_video_dynamic_shift_uses_actual_sequence_length(self): + self.assert_uses_actual_video_sequence_length_for_shift(pipeline_ltx2.LTX2Pipeline) + + def test_image_to_video_dynamic_shift_uses_actual_sequence_length(self): + self.assert_uses_actual_video_sequence_length_for_shift(pipeline_ltx2_image2video.LTX2ImageToVideoPipeline) + + +if __name__ == "__main__": + unittest.main()