| 1 | import torch |
| 2 | |
| 3 | from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps |
| 4 | from ltx_core.model.audio_vae.attention import AttentionType |
| 5 | from ltx_core.model.audio_vae.audio_vae import AudioDecoder, AudioEncoder |
| 6 | from ltx_core.model.audio_vae.causality_axis import CausalityAxis |
| 7 | from ltx_core.model.audio_vae.vocoder import MelSTFT, Vocoder, VocoderWithBWE |
| 8 | from ltx_core.model.common.normalization import NormType |
| 9 | from ltx_core.model.model_protocol import ModelConfigurator |
| 10 | from ltx_core.utils import check_config_value |
| 11 | |
| 12 | |
| 13 | def _vocoder_from_config( |
| 14 | cfg: dict, |
| 15 | apply_final_activation: bool = True, |
| 16 | output_sampling_rate: int | None = None, |
| 17 | ) -> Vocoder: |
| 18 | """Instantiate a Vocoder from a flat config dict. |
| 19 | Args: |
| 20 | cfg: Vocoder config dict (keys match Vocoder constructor args). |
| 21 | apply_final_activation: Whether to apply tanh/clamp at the output. |
| 22 | output_sampling_rate: Explicit override for the output sample rate. |
| 23 | When None, reads from cfg["output_sampling_rate"] (default 24000). |
| 24 | """ |
| 25 | return Vocoder( |
| 26 | resblock_kernel_sizes=cfg.get("resblock_kernel_sizes", [3, 7, 11]), |
| 27 | upsample_rates=cfg.get("upsample_rates", [6, 5, 2, 2, 2]), |
| 28 | upsample_kernel_sizes=cfg.get("upsample_kernel_sizes", [16, 15, 8, 4, 4]), |
| 29 | resblock_dilation_sizes=cfg.get("resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]]), |
| 30 | upsample_initial_channel=cfg.get("upsample_initial_channel", 1024), |
| 31 | resblock=cfg.get("resblock", "1"), |
| 32 | output_sampling_rate=( |
| 33 | output_sampling_rate if output_sampling_rate is not None else cfg.get("output_sampling_rate", 24000) |
| 34 | ), |
| 35 | activation=cfg.get("activation", "snake"), |
| 36 | use_tanh_at_final=cfg.get("use_tanh_at_final", True), |
| 37 | apply_final_activation=apply_final_activation, |
| 38 | use_bias_at_final=cfg.get("use_bias_at_final", True), |
| 39 | ) |
| 40 | |
| 41 | |
| 42 | class VocoderConfigurator(ModelConfigurator[Vocoder]): |
| 43 | """Configurator that auto-detects the checkpoint format. |
| 44 | Returns a plain Vocoder for pre-ltx-2.3 checkpoints (flat config) or a |
| 45 | VocoderWithBWE for ltx-2.3+ checkpoints (nested "vocoder" + "bwe" config). |
| 46 | """ |
| 47 | |
| 48 | @classmethod |
| 49 | def from_config(cls: type[Vocoder], config: dict) -> Vocoder | VocoderWithBWE: |
| 50 | cfg = config.get("vocoder", {}) |
| 51 | |
| 52 | if "bwe" not in cfg: |
| 53 | check_config_value(cfg, "resblock", "1") |
| 54 | check_config_value(cfg, "stereo", True) |
| 55 | return _vocoder_from_config(cfg) |
| 56 | |
| 57 | vocoder_cfg = cfg.get("vocoder", {}) |
| 58 | bwe_cfg = cfg["bwe"] |
| 59 | |
| 60 | check_config_value(vocoder_cfg, "resblock", "AMP1") |
| 61 | check_config_value(vocoder_cfg, "stereo", True) |
| 62 | check_config_value(vocoder_cfg, "activation", "snakebeta") |
| 63 | check_config_value(bwe_cfg, "resblock", "AMP1") |
| 64 | check_config_value(bwe_cfg, "stereo", True) |
| 65 | check_config_value(bwe_cfg, "activation", "snakebeta") |
| 66 | |
| 67 | vocoder = _vocoder_from_config( |
| 68 | vocoder_cfg, |
| 69 | output_sampling_rate=bwe_cfg["input_sampling_rate"], |
| 70 | ) |
| 71 | bwe_generator = _vocoder_from_config( |
| 72 | bwe_cfg, |
| 73 | apply_final_activation=False, |
| 74 | output_sampling_rate=bwe_cfg["output_sampling_rate"], |
| 75 | ) |
| 76 | mel_stft = MelSTFT( |
| 77 | filter_length=bwe_cfg["n_fft"], |
| 78 | hop_length=bwe_cfg["hop_length"], |
| 79 | win_length=bwe_cfg["n_fft"], |
| 80 | n_mel_channels=bwe_cfg["num_mels"], |
| 81 | ) |
| 82 | return VocoderWithBWE( |
| 83 | vocoder=vocoder, |
| 84 | bwe_generator=bwe_generator, |
| 85 | mel_stft=mel_stft, |
| 86 | input_sampling_rate=bwe_cfg["input_sampling_rate"], |
| 87 | output_sampling_rate=bwe_cfg["output_sampling_rate"], |
| 88 | hop_length=bwe_cfg["hop_length"], |
| 89 | ) |
| 90 | |
| 91 | |
| 92 | def _strip_vocoder_prefix(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]: |
| 93 | """Strip the leading 'vocoder.' prefix exactly once. |
| 94 | Uses removeprefix instead of str.replace so that BWE keys like |
| 95 | 'vocoder.vocoder.conv_pre' become 'vocoder.conv_pre' (not 'conv_pre'). |
| 96 | Works identically for legacy keys like 'vocoder.conv_pre' → 'conv_pre'. |
| 97 | """ |
| 98 | return [KeyValueOperationResult(key.removeprefix("vocoder."), value)] |
| 99 | |
| 100 | |
| 101 | VOCODER_COMFY_KEYS_FILTER = ( |
| 102 | SDOps("VOCODER_COMFY_KEYS_FILTER") |
| 103 | .with_matching(prefix="vocoder.") |
| 104 | .with_kv_operation(operation=_strip_vocoder_prefix, key_prefix="vocoder.") |
| 105 | ) |
| 106 | |
| 107 | |
| 108 | class AudioDecoderConfigurator(ModelConfigurator[AudioDecoder]): |
| 109 | @classmethod |
| 110 | def from_config(cls: type[AudioDecoder], config: dict) -> AudioDecoder: |
| 111 | audio_vae_cfg = config.get("audio_vae", {}) |
| 112 | model_cfg = audio_vae_cfg.get("model", {}) |
| 113 | model_params = model_cfg.get("params", {}) |
| 114 | ddconfig = model_params.get("ddconfig", {}) |
| 115 | preprocessing_cfg = audio_vae_cfg.get("preprocessing", {}) |
| 116 | stft_cfg = preprocessing_cfg.get("stft", {}) |
| 117 | mel_cfg = preprocessing_cfg.get("mel", {}) |
| 118 | variables_cfg = audio_vae_cfg.get("variables", {}) |
| 119 | |
| 120 | sample_rate = model_params.get("sampling_rate", 16000) |
| 121 | mel_hop_length = stft_cfg.get("hop_length", 160) |
| 122 | is_causal = stft_cfg.get("causal", True) |
| 123 | mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") or variables_cfg.get("mel_bins") |
| 124 | |
| 125 | return AudioDecoder( |
| 126 | ch=ddconfig.get("ch", 128), |
| 127 | out_ch=ddconfig.get("out_ch", 2), |
| 128 | ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))), |
| 129 | num_res_blocks=ddconfig.get("num_res_blocks", 2), |
| 130 | attn_resolutions=ddconfig.get("attn_resolutions", {8, 16, 32}), |
| 131 | resolution=ddconfig.get("resolution", 256), |
| 132 | z_channels=ddconfig.get("z_channels", 8), |
| 133 | norm_type=NormType(ddconfig.get("norm_type", "pixel")), |
| 134 | causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")), |
| 135 | dropout=ddconfig.get("dropout", 0.0), |
| 136 | mid_block_add_attention=ddconfig.get("mid_block_add_attention", True), |
| 137 | sample_rate=sample_rate, |
| 138 | mel_hop_length=mel_hop_length, |
| 139 | is_causal=is_causal, |
| 140 | mel_bins=mel_bins, |
| 141 | ) |
| 142 | |
| 143 | |
| 144 | class AudioEncoderConfigurator(ModelConfigurator[AudioEncoder]): |
| 145 | @classmethod |
| 146 | def from_config(cls: type[AudioEncoder], config: dict) -> AudioEncoder: |
| 147 | audio_vae_cfg = config.get("audio_vae", {}) |
| 148 | model_cfg = audio_vae_cfg.get("model", {}) |
| 149 | model_params = model_cfg.get("params", {}) |
| 150 | ddconfig = model_params.get("ddconfig", {}) |
| 151 | preprocessing_cfg = audio_vae_cfg.get("preprocessing", {}) |
| 152 | stft_cfg = preprocessing_cfg.get("stft", {}) |
| 153 | mel_cfg = preprocessing_cfg.get("mel", {}) |
| 154 | variables_cfg = audio_vae_cfg.get("variables", {}) |
| 155 | |
| 156 | sample_rate = model_params.get("sampling_rate", 16000) |
| 157 | mel_hop_length = stft_cfg.get("hop_length", 160) |
| 158 | n_fft = stft_cfg.get("filter_length", 1024) |
| 159 | is_causal = stft_cfg.get("causal", True) |
| 160 | mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") or variables_cfg.get("mel_bins") |
| 161 | |
| 162 | return AudioEncoder( |
| 163 | ch=ddconfig.get("ch", 128), |
| 164 | ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))), |
| 165 | num_res_blocks=ddconfig.get("num_res_blocks", 2), |
| 166 | attn_resolutions=ddconfig.get("attn_resolutions", {8, 16, 32}), |
| 167 | resolution=ddconfig.get("resolution", 256), |
| 168 | z_channels=ddconfig.get("z_channels", 8), |
| 169 | double_z=ddconfig.get("double_z", True), |
| 170 | dropout=ddconfig.get("dropout", 0.0), |
| 171 | resamp_with_conv=ddconfig.get("resamp_with_conv", True), |
| 172 | in_channels=ddconfig.get("in_channels", 2), |
| 173 | attn_type=AttentionType(ddconfig.get("attn_type", "vanilla")), |
| 174 | mid_block_add_attention=ddconfig.get("mid_block_add_attention", True), |
| 175 | norm_type=NormType(ddconfig.get("norm_type", "pixel")), |
| 176 | causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")), |
| 177 | sample_rate=sample_rate, |
| 178 | mel_hop_length=mel_hop_length, |
| 179 | n_fft=n_fft, |
| 180 | is_causal=is_causal, |
| 181 | mel_bins=mel_bins, |
| 182 | ) |
| 183 | |
| 184 | |
| 185 | AUDIO_VAE_DECODER_COMFY_KEYS_FILTER = ( |
| 186 | SDOps("AUDIO_VAE_DECODER_COMFY_KEYS_FILTER") |
| 187 | .with_matching(prefix="audio_vae.decoder.") |
| 188 | .with_matching(prefix="audio_vae.per_channel_statistics.") |
| 189 | .with_replacement("audio_vae.decoder.", "") |
| 190 | .with_replacement("audio_vae.per_channel_statistics.", "per_channel_statistics.") |
| 191 | ) |
| 192 | |
| 193 | |
| 194 | AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER = ( |
| 195 | SDOps("AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER") |
| 196 | .with_matching(prefix="audio_vae.encoder.") |
| 197 | .with_matching(prefix="audio_vae.per_channel_statistics.") |
| 198 | .with_replacement("audio_vae.encoder.", "") |
| 199 | .with_replacement("audio_vae.per_channel_statistics.", "per_channel_statistics.") |
| 200 | ) |
| 201 |