XCodec2

https://huggingface.co/HKUSTAudio/xcodec2
Admin
XCodec2 is a speech tokenizer that converts audio to high-quality discrete codes with single vector...
Tokens:4,469
Snippets:45
Trust Score:10
Update:6 months ago
Show doc for...
Tokens:
Raw
### Install XCodec2 Environment

Source: https://huggingface.co/HKUSTAudio/xcodec2/blob/main/README.md

Commands to create a conda environment and install the xcodec2 package.

```bash
conda create -n xcodec2 python=3.9
conda activate xcodec2
pip install xcodec2  (Use `xcodec2==0.1.5` for codec inference and llasa fine-tuning. I’ve removed unnecessary dependencies, and it works fine in my testing. However,  I’m not sure if other problems may arise. If you prefer more stability, I recommend using `xcodec2==0.1.3` which accurately aligns during my codec training.)
```

--------------------------------

### Python Code for Audio Generation Setup

Source: https://huggingface.co/HKUSTAudio/xcodec2/discussions/13

This Python code sets up the environment for audio generation using XCodec2. It includes necessary imports, path configurations, and debugging flags. The code demonstrates how to load model configurations and initialize models with empty weights, addressing potential weight initialization issues.

```python
import omegaconf

import typing
import torch

import soundfile as sf
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, PreTrainedModel
from transformers import AutoFeatureExtractor, Wav2Vec2BertModel

from accelerate import init_empty_weights

from collections import OrderedDict

import collections

from configuration_bigcodec import BigCodecConfig
from xcodec2.modeling_xcodec2 import XCodec2Model

from safetensors.torch import load_file, save_file

ROOT='/aifs4su/yiakwy/shared/audio'

PROJECT_ROOT='/home/yiakwy/workspace/Github/sglang/3rdparty/X-Codec-2.0'

model_path = f"{ROOT}/xcodec2"

Debug=True

def find_meta_tensors(module, prefix=""):
    metas = []
    for name, param in module.named_parameters(recurse=True):
        if param.is_meta:
            metas.append(prefix + name)
    for name, buf in module.named_buffers(recurse=True):
        if buf.is_meta:
            metas.append(prefix + name)
    return metas


if Debug:
    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)

    with init_empty_weights():
        model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)

    # torch >= 2.6
    torch.serialization.add_safe_globals([list])
    torch.serialization.add_safe_globals([dict])
    torch.serialization.add_safe_globals([int])

    torch.serialization.add_safe_globals([typing.Any])
    torch.serialization.add_safe_globals([collections.defaultdict])

    torch.serialization.add_safe_globals([omegaconf.base.ContainerMetadata])
    torch.serialization.add_safe_globals([omegaconf.listconfig.ListConfig])
    torch.serialization.add_safe_globals([omegaconf.dictconfig.DictConfig])
    torch.serialization.add_safe_globals([omegaconf.nodes.AnyNode])
    torch.serialization.add_safe_globals([omegaconf.base.Metadata])

    ckpt = f"{model_path}/ckpt/epoch=4-step=1400000.ckpt"
    state_dict = torch.load(ckpt, map_location="cpu", weights_only=True)
    state_dict = state_dict['state_dict']

    filtered_state_dict_codec = OrderedDict()
```

--------------------------------

### XCodec2 Configuration File

Source: https://huggingface.co/HKUSTAudio/xcodec2/discussions/13

Example JSON configuration for the XCodec2 model environment.

```json
{
  "Wav2Vec2BertModelROOT": "/aifs4su/yiakwy/shared/audio/",
  "architectures": [
    "XCodec2Model"
  ],
  "auto_map": {
    "AutoConfig": "configuration_bigcodec.BigCodecConfig",
    "AutoModelForCausalLM": "modeling_xcodec2.XCodec2Model"
  },
  "codec_decoder_hidden_size": 1024,
  "codec_encoder_hidden_size": 1024,
  "model_type": "bigcodec",
  "semantic_hidden_size": 1024,
  "torch_dtype": "float32",
  "transformers_version": "4.51.3",
  "use_vocos": true
}
```

--------------------------------

### Install XCodec2

Source: https://huggingface.co/HKUSTAudio/xcodec2/blob/main/README.md?code=true

Install XCodec2 using conda and pip. Use version 0.1.5 for codec inference and LLaSA fine-tuning, or 0.1.3 for stable codec training.

```bash
conda create -n xcodec2 python=3.9
conda activate xcodec2
pip install xcodec2 (Use `xcodec2==0.1.5` for codec inference and llasa fine-tuning. I’ve removed unnecessary dependencies, and it works fine in my testing. However, I’m not sure if other problems may arise. If you prefer more stability, I recommend using `xcodec2==0.1.3` which accurately aligns during my codec training.)
```

--------------------------------

### Python Module Initialization

Source: https://huggingface.co/HKUSTAudio/xcodec2/tree/main/vq/alias_free_torch

This file initializes the xcodec2 package. It is part of the core library structure.

```python
Safe
199 Bytes
Initial commit over 1 year ago
```

--------------------------------

### Run Inference with XCodec2

Source: https://huggingface.co/HKUSTAudio/xcodec2/blob/main/README.md

Python script to load the model, encode an audio file, and reconstruct the waveform.

```python
import torch
import soundfile as sf
from transformers import AutoConfig

 
from xcodec2.modeling_xcodec2 import XCodec2Model
 
model_path = "HKUSTAudio/xcodec2"  
 
model = XCodec2Model.from_pretrained(model_path)
model.eval().cuda()   

 
wav, sr = sf.read("test.wav")   
wav_tensor = torch.from_numpy(wav).float().unsqueeze(0)  # Shape: (1, T)

 
with torch.no_grad():
   # Only 16khz speech
   # Only supports single input. For batch inference, please refer to the link below.
    vq_code = model.encode_code(input_waveform=wav_tensor)
    print("Code:", vq_code )  

    recon_wav = model.decode_code(vq_code).cpu()       # Shape: (1, 1, T')

 
sf.write("reconstructed.wav", recon_wav[0, 0, :].numpy(), sr)
print("Done! Check reconstructed.wav")
```

--------------------------------

### Load XCodec2 Model Configuration

Source: https://huggingface.co/HKUSTAudio/xcodec2/discussions/13

Loads the configuration for the XCodec2 model from a specified path. Ensure the model path is correct and the configuration files are present.

```python
config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
```

--------------------------------

### Initialize XCodec2 Model with Empty Weights

Source: https://huggingface.co/HKUSTAudio/xcodec2/discussions/13

Initializes the XCodec2 model with empty weights using `init_empty_weights` from the `accelerate` library. This is useful for loading large models without immediately consuming significant memory. Note that some weights might be newly initialized, requiring model training.

```python
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
```

--------------------------------

### Use XCodec2 for Speech Encoding and Decoding

Source: https://huggingface.co/HKUSTAudio/xcodec2/blob/main/README.md?code=true

Load the XCodec2 model and use it to encode a WAV file into discrete codes and then decode these codes back into reconstructed audio.

```python
import torch
import soundfile as sf
from transformers import AutoConfig


from xcodec2.modeling_xcodec2 import XCodec2Model

model_path = "HKUSTAudio/xcodec2" 

model = XCodec2Model.from_pretrained(model_path)
model.eval().cuda() 


wav, sr = sf.read("test.wav") 
wav_tensor = torch.from_numpy(wav).float().unsqueeze(0) # Shape: (1, T)  


with torch.no_grad():
 # Only 16khz speech
 # Only supports single input. For batch inference, please refer to the link below.
 vq_code = model.encode_code(input_waveform=wav_tensor)
 print("Code:", vq_code ) 

 recon_wav = model.decode_code(vq_code).cpu() # Shape: (1, 1, T')  


sf.write("reconstructed.wav", recon_wav[0, 0, :].numpy(), sr)
print("Done! Check reconstructed.wav")
```

--------------------------------

### Encode and Decode Audio with XCodec2

Source: https://huggingface.co/HKUSTAudio/xcodec2/resolve/main/README.md?download=true

Load the XCodec2 model, encode a WAV file into discrete codes, and then decode these codes back into an audio waveform. Ensure the input audio is 16kHz.

```python
import torch
import soundfile as sf
from transformers import AutoConfig

from xcodec2.modeling_xcodec2 import XCodec2Model

model_path = "HKUSTAudio/xcodec2"  

model = XCodec2Model.from_pretrained(model_path)
model.eval().cuda()   

wav, sr = sf.read("test.wav")   
wav_tensor = torch.from_numpy(wav).float().unsqueeze(0)  # Shape: (1, T)

with torch.no_grad():
   # Only 16khz speech
   # Only supports single input. For batch inference, please refer to the link below.
    vq_code = model.encode_code(input_waveform=wav_tensor)
    print("Code:", vq_code )

    recon_wav = model.decode_code(vq_code).cpu()       # Shape: (1, 1, T')

sf.write("reconstructed.wav", recon_wav[0, 0, :].numpy(), sr)
print("Done! Check reconstructed.wav")
```

--------------------------------

### Python Activation Module

Source: https://huggingface.co/HKUSTAudio/xcodec2/tree/main/vq/alias_free_torch

This file contains activation functions used within the xcodec2 model. It is essential for the neural network's forward pass.

```python
Safe
854 Bytes
Initial commit over 1 year ago
```

--------------------------------

### Python Resample Module

Source: https://huggingface.co/HKUSTAudio/xcodec2/tree/main/vq/alias_free_torch

This file handles audio resampling functionalities for the xcodec2 model. It is crucial for managing different audio sample rates.

```python
Safe
1.86 kB
Initial commit over 1 year ago
```

--------------------------------

### xcodec2 Model Configuration

Source: https://huggingface.co/HKUSTAudio/xcodec2/blame/main/config.json

JSON configuration file defining the architecture and hidden size parameters for the xcodec2 model.

```json
{
    "model_type": "xcodec2",
    "semantic_hidden_size": 1024,
    "codec_encoder_hidden_size": 1024,
    "codec_decoder_hidden_size": 1024,
    "use_vocos": true,
    "architectures": [
      "XCodec2Model"
    ]
  }
```

--------------------------------

### Configure Model Metadata

Source: https://huggingface.co/HKUSTAudio/xcodec2/blob/main/README.md

YAML configuration for model metadata including license and pipeline tags.

```yaml
license: cc-by-nc-4.0
tags:
  - audio-to-audio
pipeline_tag: audio-to-audio
```

--------------------------------

### Python Filter Module

Source: https://huggingface.co/HKUSTAudio/xcodec2/tree/main/vq/alias_free_torch

This file implements filtering operations, likely for audio signal processing within the xcodec2 model.

```python
Safe
3.46 kB
Initial commit over 1 year ago
```

--------------------------------

### Perform Audio Reconstruction Inference

Source: https://huggingface.co/HKUSTAudio/xcodec2/discussions/13

Reads an audio file, encodes it into VQ codes, and reconstructs the waveform.

```python
    wav, sr = sf.read(f"{ROOT}/xcodec2/test.flac")
    wav_tensor = torch.from_numpy(wav).float().unsqueeze(0)  # Shape: (1, T)


    with torch.no_grad():
        vq_code = model.encode_code(input_waveform=wav_tensor)
        print("Code:", vq_code )  

        recon_wav = model.decode_code(vq_code).cpu()       # Shape: (1, 1, T')

    
    sf.write("reconstructed.wav", recon_wav[0, 0, :].numpy(), sr)
    print("Done! Check reconstructed.wav")
```

--------------------------------

### Load and Filter Checkpoint State Dictionary

Source: https://huggingface.co/HKUSTAudio/xcodec2/discussions/13

Loads the state dictionary from a PyTorch checkpoint file and filters it. This process may involve renaming keys to match the model's expected structure, especially when migrating between different library versions.

```python
ckpt = f"{model_path}/ckpt/epoch=4-step=1400000.ckpt"
state_dict = torch.load(ckpt, map_location="cpu", weights_only=True)
state_dict = state_dict['state_dict']

filtered_state_dict_codec = OrderedDict()
```

--------------------------------

### Filter and Load Model State Dictionaries

Source: https://huggingface.co/HKUSTAudio/xcodec2/discussions/13

Iterates through a state dictionary to filter and rename keys for specific model modules before loading them into the model components.

```python
    filtered_state_dict_semantic_encoder = OrderedDict()
    filtered_state_dict_gen = OrderedDict()
    filtered_state_dict_fc_post_a = OrderedDict()
    filtered_state_dict_fc_prior = OrderedDict()
    for key, value in state_dict.items():
        if key.startswith('CodecEnc.'):
            print(f"adding key = {key} ...")
            new_key = key[len('CodecEnc.'):]

            if "beta" in new_key:
                old_key = new_key
                new_key = new_key.replace("beta", "bias")
                print(f"updating key from {old_key} to {new_key}")

            filtered_state_dict_codec[new_key] = value
        elif key.startswith('generator.'):
            new_key = key[len('generator.'):]
            filtered_state_dict_gen[new_key] = value
        elif key.startswith('fc_post_a.'):
            new_key = key[len('fc_post_a.'):]
            filtered_state_dict_fc_post_a[new_key] = value
        elif key.startswith('SemanticEncoder_module.'):
            new_key = key[len('SemanticEncoder_module.'):]
            filtered_state_dict_semantic_encoder[new_key] = value
        elif key.startswith('fc_prior.'):
            new_key = key[len('fc_prior.'):]
            filtered_state_dict_fc_prior[new_key] = value

    model.semantic_model = Wav2Vec2BertModel.from_pretrained(f"{config.Wav2Vec2BertModelROOT}/w2v-bert-2.0", output_hidden_states=True)
    model.semantic_model.eval().cuda()

    model.SemanticEncoder_module.load_state_dict(filtered_state_dict_semantic_encoder, strict=False, assign=True)
    model.SemanticEncoder_module.eval().cuda()

    model.CodecEnc.load_state_dict(filtered_state_dict_codec, strict=False, assign=True)
    metas = find_meta_tensors(model.CodecEnc)
    model.CodecEnc.eval().cuda()

    model.generator.load_state_dict(filtered_state_dict_gen, strict=False, assign=True)
    model.generator.eval().cuda()

    model.fc_prior.load_state_dict(filtered_state_dict_fc_prior, strict=False, assign=True)
    model.fc_prior.eval().cuda()

    model.fc_post_a.load_state_dict(filtered_state_dict_fc_post_a, strict=False, assign=True)
    model.fc_post_a.eval().cuda()

    model.feature_extractor = AutoFeatureExtractor.from_pretrained(f"{config.Wav2Vec2BertModelROOT}/w2v-bert-2.0")

    output_dir='/aifs4su/yiakwy/shared/audio/test_cvt_hf_xcodec2'

    model.eval().cuda()
```

--------------------------------

### Register Safe Globals for Serialization

Source: https://huggingface.co/HKUSTAudio/xcodec2/discussions/13

Registers necessary global types for PyTorch's serialization to handle custom objects like lists, dictionaries, and omegaconf configurations. This is crucial when loading checkpoints that contain these types.

```python
torch.serialization.add_safe_globals([list])
torch.serialization.add_safe_globals([dict])
torch.serialization.add_safe_globals([int])
torch.serialization.add_safe_globals([typing.Any])
torch.serialization.add_safe_globals([collections.defaultdict])
torch.serialization.add_safe_globals([omegaconf.base.ContainerMetadata])
torch.serialization.add_safe_globals([omegaconf.listconfig.ListConfig])
torch.serialization.add_safe_globals([omegaconf.dictconfig.DictConfig])
torch.serialization.add_safe_globals([omegaconf.nodes.AnyNode])
torch.serialization.add_safe_globals([omegaconf.base.Metadata])
```
Tokens:
Raw