Entire other VAE in LTX-2.3

Hello! I know you are a vibe-coder, but this is something The VAE in the checkpoint loader code is denoted as ```python if version == "2.3": return { "in_channels": 3, "out_channels": 3, "latent_channels": 128, "block_out_channels": (256, 512, 1024, 1024), "down_block_types": ( "LTX2VideoDownBlock3D", "LTX2VideoDownBlock3D", "LTX2VideoDownBlock3D", "LTX2VideoDownBlock3D", ), "decoder_block_out_channels": (256, 512, 512, 1024), "layers_per_block": (4, 6, 4, 2, 2), "decoder_layers_per_block": (4, 6, 4, 2, 2), "spatio_temporal_scaling": (True, True, True, True), "decoder_spatio_temporal_scaling": (True, True, True, True), "decoder_inject_noise": (False, False, False, False, False), "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), "upsample_type": ("spatiotemporal", "spatiotemporal", "temporal", "spatial"), "upsample_residual": (False, False, False, False), "upsample_factor": (2, 2, 1, 2), "timestep_conditioning": False, "patch_size": 4, "patch_size_t": 1, "resnet_norm_eps": 1e-6, "encoder_causal": True, "decoder_causal": False, "encoder_spatial_padding_mode": "zeros", "decoder_spatial_padding_mode": "zeros", "spatial_compression_ratio": 32, "temporal_compression_ratio": 8, } ``` Then it fails with the error: ``` File "/media/kabachuha/xiangliu/simpletunerx/SimpleTunerX-exp/simpletuner/helpers/models/ltxvideo2/checkpoint_loader.py", line 701, in convert_ltx2_video_vae vae.load_state_dict(original_state_dict, strict=True, assign=True) File "/media/kabachuha/holodok01/miniconda3/envs/simpletuner/lib/python3.12/site-packages/torch/nn/modules/module.py", line 2635, in load_state_dict raise RuntimeError( RuntimeError: Error(s) in loading state_dict for AutoencoderKLLTX2Video: size mismatch for decoder.up_blocks.0.upsamplers.0.conv.conv.weight: copying a param with shape torch.Size([4096, 1024, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([2048, 1024, 3, 3, 3]). size mismatch for decoder.up_blocks.0.upsamplers.0.conv.conv.bias: copying a param with shape torch.Size([4096]) from checkpoint, the shape in current model is torch.Size([2048]). size mismatch for decoder.up_blocks.1.upsamplers.0.conv.conv.weight: copying a param with shape torch.Size([4096, 512, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([1024, 512, 3, 3, 3]). size mismatch for decoder.up_blocks.1.upsamplers.0.conv.conv.bias: copying a param with shape torch.Size([4096]) from checkpoint, the shape in current model is torch.Size([1024]). size mismatch for decoder.up_blocks.2.upsamplers.0.conv.conv.weight: copying a param with shape torch.Size([512, 512, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([2048, 512, 3, 3, 3]). size mismatch for decoder.up_blocks.2.upsamplers.0.conv.conv.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([2048]). size mismatch for decoder.up_blocks.3.upsamplers.0.conv.conv.weight: copying a param with shape torch.Size([512, 256, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([1024, 256, 3, 3, 3]). size mismatch for decoder.up_blocks.3.upsamplers.0.conv.conv.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([1024]). ``` It gaslighted me for like an hour, fortunately there is the actual diffusers config for LTX-2.3: https://huggingface.co/CalamitousFelicitousness/LTX-2.3-dev-Diffusers/blob/main/vae/config.json And the VAE definition fixes to: ``` return { "in_channels": 3, "out_channels": 3, "latent_channels": 128, "block_out_channels": (256, 512, 1024, 1024), "down_block_types": ( "LTX2VideoDownBlock3D", "LTX2VideoDownBlock3D", "LTX2VideoDownBlock3D", "LTX2VideoDownBlock3D", ), "decoder_block_out_channels": (256, 512, 512, 1024), "layers_per_block": (4, 6, 4, 2, 2), "decoder_layers_per_block": (4, 6, 4, 2, 2), "spatio_temporal_scaling": (True, True, True, True), "decoder_spatio_temporal_scaling": (True, True, True, True), "decoder_inject_noise": (False, False, False, False, False), "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), "upsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), # Fixed: match Diffusers config "upsample_residual": (True, True, True, True), # Fixed: all True "upsample_factor": (2, 2, 1, 2), "timestep_conditioning": False, "patch_size": 4, "patch_size_t": 1, "resnet_norm_eps": 1e-6, "encoder_causal": True, "decoder_causal": False, "encoder_spatial_padding_mode": "zeros", "decoder_spatial_padding_mode": "reflect", # Fixed: reflect instead of zeros "spatial_compression_ratio": 32, "temporal_compression_ratio": 8, } ``` Which is may not be ideal because it's compared with LLM, but then the VAE loads fine from the 2.3 dev checkpoint, caches the dataset and launches training successfully (well, not counting the device mismatch error I will fix later...) Did you launch the update without any real-life tests on LTX-2.3? 🙂