diff --git a/simpletuner/helpers/models/ltxvideo2/autoencoder.py b/simpletuner/helpers/models/ltxvideo2/autoencoder.py index c53141151..553c0fcd7 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,12 @@ 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 + conv_in_out_channels = upsampler_in_channels if spatio_temporal_scale else out_channels + if in_channels != conv_in_out_channels: self.conv_in = LTX2VideoResnetBlock3d( in_channels=in_channels, - out_channels=out_channels, + out_channels=conv_in_out_channels, dropout=dropout, eps=resnet_eps, non_linearity=resnet_act_fn, @@ -986,7 +988,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..1256011e7 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,74 @@ 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_up_block_without_upsampling_projects_to_resnet_width(self): + up_block = LTX2VideoUpBlock3d( + in_channels=64, + out_channels=32, + num_layers=1, + spatio_temporal_scale=False, + upscale_factor=2, + ) + + self.assertIsNotNone(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, 4, 8, 8)) + + 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() diff --git a/tests/test_ltxvideo2_dynamic_shift.py b/tests/test_ltxvideo2_dynamic_shift.py new file mode 100644 index 000000000..02293dfba --- /dev/null +++ b/tests/test_ltxvideo2_dynamic_shift.py @@ -0,0 +1,52 @@ +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): + def is_calculate_shift(node): + return (isinstance(node, ast.Name) and node.id == "calculate_shift") or ( + isinstance(node, ast.Attribute) and node.attr == "calculate_shift" + ) + + return [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) and is_calculate_shift(node.func) + ] + + +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.assertGreaterEqual(len(calls), 1) + self.assertTrue( + any( + any(isinstance(arg, ast.Name) and arg.id == "video_sequence_length" for arg in call.args) + or any( + kw.arg == "video_sequence_length" and isinstance(kw.value, ast.Name) and kw.value.id == "video_sequence_length" + for kw in call.keywords + ) + for call in calls + ) + ) + + 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()