Fix LTX2 VAE widths and dynamic shift coverage

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(

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

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