# Copyright 2026 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import argparse import json import re import safetensors import torch import torch.nn as nn from transformers import ( AutoConfig, AutoFeatureExtractor, Xcodec2Config, Xcodec2FeatureExtractor, Xcodec2Model, logging, ) logging.set_verbosity_info() logger = logging.get_logger("transformers.models.xcodec2") # Original key names: https://huggingface.co/HKUSTAudio/xcodec2/blob/main/modeling_xcodec2.py#L12 # Mappings are applied sequentially per key; order matters. # fmt: off STATE_DICT_MAPPING = { # ── Semantic encoder (semantic_model -> semantic_encoder) ── r"^semantic_model\.": r"semantic_encoder.", # ── Acoustic encoder (CodecEnc -> acoustic_encoder) ── r"^CodecEnc\.conv_blocks\.0\.": r"acoustic_encoder.conv1.", r"^CodecEnc\.conv_final_block\.0\.": r"acoustic_encoder.snake1.", r"^CodecEnc\.conv_final_block\.1\.": r"acoustic_encoder.conv2.", r"^CodecEnc\.conv_blocks\.(\d+)\.": r"acoustic_encoder.block.SHIFT_DOWN.\1.", r"^CodecEnc\.": r"acoustic_encoder.", # -- EncoderBlock: block.{0,1,2} -> res_unit{1,2,3}, block.3 -> snake1, block.4 -> conv1 # NOTE: These must come BEFORE ResidualUnit mappings so keys are transformed first r"(\.block\.\d+)\.block\.0\.": r"\1.res_unit1.", r"(\.block\.\d+)\.block\.1\.": r"\1.res_unit2.", r"(\.block\.\d+)\.block\.2\.": r"\1.res_unit3.", r"(\.block\.\d+)\.block\.3\.": r"\1.snake1.", r"(\.block\.\d+)\.block\.4\.": r"\1.conv1.", # -- ResidualUnit: block.{0,1,2,3} -> {snake1,conv1,snake2,conv2} r"(\.res_unit\d+)\.block\.0\.": r"\1.snake1.", r"(\.res_unit\d+)\.block\.1\.": r"\1.conv1.", r"(\.res_unit\d+)\.block\.2\.": r"\1.snake2.", r"(\.res_unit\d+)\.block\.3\.": r"\1.conv2.", # -- ResidualUnit nested structure: res_unit{1,2,3,4} -> {snake1,conv1,snake2,conv2} r"(\.res_unit\d+)\.res_unit1\.": r"\1.snake1.", r"(\.res_unit\d+)\.res_unit2\.": r"\1.conv1.", r"(\.res_unit\d+)\.res_unit3\.": r"\1.snake2.", r"(\.res_unit\d+)\.res_unit4\.": r"\1.conv2.", # -- DownSample1d: lowpass.filter -> filter (LowPassFilter1d inlined) r"\.downsample\.lowpass\.filter": r".downsample.filter", # ── Quantizer: lives on main model, not inside decoder ── r"^generator\.quantizer\.layers\.0\.": r"quantizer.quantizer.", r"^generator\.quantizer\.": r"quantizer.", # ── Decoder (generator -> acoustic_decoder) ── # -- Handle backbone components: now directly in decoder (no backbone. prefix) r"^generator\.backbone\.prior_net\.": r"acoustic_decoder.prior_net.", r"^generator\.backbone\.post_net\.": r"acoustic_decoder.post_net.", r"^generator\.backbone\.final_layer_norm\.": r"acoustic_decoder.norm.", r"^generator\.backbone\.": r"acoustic_decoder.", # -- General generator mapping r"^generator\.": r"acoustic_decoder.", # -- ISTFT head: out -> linear r"acoustic_decoder\.head\.out\.": r"acoustic_decoder.head.linear.", # -- Transformer layers: transformers -> layers r"\.transformers\.": r".layers.", r"\.att\.c_proj\.": r".self_attn.o_proj.", r"\.att_norm\.": r".input_layernorm.", r"\.ffn_norm\.": r".post_attention_layernorm.", # ── Semantic adapter (SemanticEncoder_module -> semantic_adapter) ── r"^SemanticEncoder_module\.": r"semantic_adapter.", r"semantic_adapter\.initial_conv\.": r"semantic_adapter.conv1.", r"semantic_adapter\.final_conv\.": r"semantic_adapter.conv4.", r"semantic_adapter\.residual_blocks\.0\.": r"semantic_adapter.act1.", r"semantic_adapter\.residual_blocks\.1\.": r"semantic_adapter.conv2.", r"semantic_adapter\.residual_blocks\.2\.": r"semantic_adapter.act2.", r"semantic_adapter\.residual_blocks\.3\.": r"semantic_adapter.conv3.", # ── Linear layers (fc_prior -> fc_encoder, fc_post_a -> acoustic_decoder.fc) ── r"^fc_prior\.": r"fc_encoder.", r"^fc_post_a\.": r"acoustic_decoder.fc.", } # fmt: on def permute_for_rope(input_tensor, n_heads, dim1, dim2): input_tensor = input_tensor.reshape(dim1, dim2) input_tensor = input_tensor.view(n_heads, dim1 // n_heads // 2, 2, dim2) input_tensor = input_tensor.transpose(1, 2).reshape(dim1, dim2) return input_tensor def map_old_key_to_new(old_key): new_key = old_key for pattern, replacement in STATE_DICT_MAPPING.items(): match = re.search(pattern, new_key) if match: # Handle index shift: conv_blocks index N (N>=1) -> encoder_blocks index N-1 if "SHIFT_DOWN" in replacement and match.groups(): layer_idx = int(match.group(1)) replacement = replacement.replace(r"SHIFT_DOWN.\1", str(layer_idx - 1)) new_key = re.sub(pattern, replacement, new_key) return new_key def _is_unused_semantic_layer(key, num_layers): """Check if a key belongs to a semantic encoder layer beyond the ones we keep.""" match = re.search(r"semantic_encoder\.encoder\.layers\.(\d+)\.", key) return match is not None and int(match.group(1)) >= num_layers def convert_state_dict(state_dict, hf_model): new_state_dict = {} for old_key, tensor in state_dict.items(): # Special case: c_attn is a fused QKV projection that must be split into 3 tensors if ".att.c_attn." in old_key: w_q, w_k, w_v = torch.chunk(tensor, chunks=3, dim=0) w_q = permute_for_rope(w_q, hf_model.config.num_attention_heads, w_q.shape[0], w_q.shape[1]) w_k = permute_for_rope(w_k, hf_model.config.num_attention_heads, w_k.shape[0], w_k.shape[1]) base_key = map_old_key_to_new(old_key.replace(".att.c_attn.", ".self_attn.PROJ.")) new_state_dict[base_key.replace("PROJ", "q_proj")] = w_q new_state_dict[base_key.replace("PROJ", "k_proj")] = w_k new_state_dict[base_key.replace("PROJ", "v_proj")] = w_v continue new_key = map_old_key_to_new(old_key) new_state_dict[new_key] = tensor if old_key != new_key: logger.debug(f"Converted: {old_key} -> {new_key}") state_dict = new_state_dict # Filter out non-persistent buffers (recomputed in _init_weights) and # unused semantic encoder layers (original has 24 layers, we only use 16) num_semantic_layers = hf_model.config.semantic_model_config.num_hidden_layers state_dict = { k: v for k, v in state_dict.items() if not ( k.endswith("upsample.filter") or k.endswith("downsample.filter") or k.endswith(".window") or k.endswith(".levels") or k.endswith(".basis") or k.endswith(".codebook") or _is_unused_semantic_layer(k, num_semantic_layers) ) } extra_keys = set(state_dict.keys()) - set(hf_model.state_dict().keys()) missing_keys = set(hf_model.state_dict().keys()) - set(state_dict.keys()) if len(extra_keys) != 0: raise ValueError(f"{len(extra_keys)} extra keys found: {extra_keys}") if len(missing_keys) != 0: raise ValueError(f"{len(missing_keys)} missing keys found: {missing_keys}") hf_model.load_state_dict(state_dict, strict=True) n_params = sum(p.numel() for p in hf_model.parameters()) logger.info(f"model loaded: {round(n_params / 1e6, 1)}M params") del state_dict return hf_model def apply_weight_norm(model): weight_norm = nn.utils.weight_norm weight_norm(model.acoustic_encoder.conv1) for encoder_block in model.acoustic_encoder.block: weight_norm(encoder_block.res_unit1.conv1) weight_norm(encoder_block.res_unit1.conv2) weight_norm(encoder_block.res_unit2.conv1) weight_norm(encoder_block.res_unit2.conv2) weight_norm(encoder_block.res_unit3.conv1) weight_norm(encoder_block.res_unit3.conv2) weight_norm(encoder_block.conv1) weight_norm(model.acoustic_encoder.conv2) def remove_weight_norm(model): param_remove = nn.utils.remove_weight_norm param_remove(model.acoustic_encoder.conv1) for encoder_block in model.acoustic_encoder.block: param_remove(encoder_block.res_unit1.conv1) param_remove(encoder_block.res_unit1.conv2) param_remove(encoder_block.res_unit2.conv1) param_remove(encoder_block.res_unit2.conv2) param_remove(encoder_block.res_unit3.conv1) param_remove(encoder_block.res_unit3.conv2) param_remove(encoder_block.conv1) param_remove(model.acoustic_encoder.conv2) @torch.no_grad() def convert_checkpoint( checkpoint_path, pytorch_dump_folder_path=None, config_path=None, repo_id=None, ): if config_path is not None: with open(config_path, "r", encoding="utf-8") as f: model_config = json.load(f) hidden_size = model_config["codec_decoder_hidden_size"] else: hidden_size = 1024 # default to https://huggingface.co/HKUSTAudio/xcodec2/blob/main/config.json config = Xcodec2Config( encoder_hidden_size=48, # https://huggingface.co/HKUSTAudio/xcodec2/blob/main/vq/codec_encoder.py#L21, like in DAC hidden_size=hidden_size, # use hardcoded semantic model: https://huggingface.co/HKUSTAudio/xcodec2/blob/main/modeling_xcodec2.py#L19 semantic_model_config=AutoConfig.from_pretrained("facebook/w2v-bert-2.0", num_hidden_layers=16), ) if not torch.cuda.is_available(): raise ValueError("Run this script on a machine with a GPU for weight norm layers to be correctly copied.") torch_device = "cuda" model = Xcodec2Model(config).to(torch_device).eval() original_checkpoint = safetensors.torch.load_file(checkpoint_path) if "best_state" in original_checkpoint: # we might have a training state saved, in which case discard the yaml results and just retain the weights original_checkpoint = original_checkpoint["best_state"] apply_weight_norm(model) model = convert_state_dict(original_checkpoint, model) remove_weight_norm(model) dac_fe = AutoFeatureExtractor.from_pretrained("descript/dac_16khz") semantic_fe = AutoFeatureExtractor.from_pretrained("facebook/w2v-bert-2.0") if semantic_fe.sampling_rate != dac_fe.sampling_rate: raise ValueError( f"Sampling rates for DAC and semantic feature extractor must match, got {dac_fe.sampling_rate} and {semantic_fe.sampling_rate}." ) feature_extractor = Xcodec2FeatureExtractor( feature_size=semantic_fe.feature_size, sampling_rate=semantic_fe.sampling_rate, padding_value=semantic_fe.padding_value, hop_length=dac_fe.hop_length, ) if pytorch_dump_folder_path is not None: model.save_pretrained(pytorch_dump_folder_path) feature_extractor.save_pretrained(pytorch_dump_folder_path) if repo_id: print("Pushing to the hub...") feature_extractor.push_to_hub(repo_id) model.push_to_hub(repo_id) """ While there is training code on GitHub: https://github.com/zhenye234/X-Codec-2.0 The modeling code on the Hub is used by the xcodec2 (pip) package: https://huggingface.co/HKUSTAudio/xcodec2/blob/main/modeling_xcodec2.py Moreover given their example usage: https://huggingface.co/HKUSTAudio/xcodec2/blob/main/modeling_xcodec2.py It seems best to use the config and weights from the Hub. Conversion example usage: ``` # download model and config wget https://huggingface.co/HKUSTAudio/xcodec2/resolve/main/config.json -P /raid/eric/xcodec2_original wget https://huggingface.co/HKUSTAudio/xcodec2/resolve/main/model.safetensors -P /raid/eric/xcodec2_original # run conversion python src/transformers/models/xcodec2/convert_xcodec2_checkpoint.py \ --checkpoint_path /raid/eric/xcodec2_original/model.safetensors \ --config_path /raid/eric/xcodec2_original/config.json \ --push_to_hub bezzam/xcodec2 ``` """ if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--checkpoint_path", required=True, default=None, type=str, help="Path to original checkpoint") parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert") parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.") parser.add_argument( "--push_to_hub", default=None, type=str, help="Where to upload the converted model on the 🤗 hub." ) args = parser.parse_args() convert_checkpoint( args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.push_to_hub, )