Function get_ltx2_audio_vae_config

scripts/convert_ltx2_to_diffusers.py:656–707

(version: str)

def get_ltx2_audio_vae_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
    if version == "2.0":
        config = {
            "model_id": "Lightricks/LTX-2",
            "diffusers_config": {
                "base_channels": 128,
                "output_channels": 2,
                "ch_mult": (1, 2, 4),
                "num_res_blocks": 2,
                "attn_resolutions": None,
                "in_channels": 2,
                "resolution": 256,
                "latent_channels": 8,
                "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": 64,
                "double_z": True,
            },
        }
        rename_dict = LTX_2_0_AUDIO_VAE_RENAME_DICT
        special_keys_remap = LTX_2_0_AUDIO_VAE_SPECIAL_KEYS_REMAP
    elif version == "2.3":
        config = {
            "model_id": "Lightricks/LTX-2.3",
            "diffusers_config": {
                "base_channels": 128,
                "output_channels": 2,
                "ch_mult": (1, 2, 4),
                "num_res_blocks": 2,
                "attn_resolutions": None,
                "in_channels": 2,
                "resolution": 256,
                "latent_channels": 8,
                "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": 64,
                "double_z": True,
            },  # Same config as LTX-2.0
        }
        rename_dict = LTX_2_0_AUDIO_VAE_RENAME_DICT
        special_keys_remap = LTX_2_0_AUDIO_VAE_SPECIAL_KEYS_REMAP
    return config, rename_dict, special_keys_remap


def convert_ltx2_audio_vae(original_state_dict: dict[str, Any], version: str) -> dict[str, Any]: