Implementation:Facebookresearch Audiocraft Export encodec
| Knowledge Sources | |
|---|---|
| Domains | |
| Last Updated | 2026-02-13 00:00 GMT |
Overview
Concrete implementation of the EnCodec model export utility within Audiocraft. The export_encodec() function extracts the best model weights and configuration from a full training checkpoint and saves them in a compact format suitable for inference and distribution.
Description
The function loads a full training checkpoint from disk, extracts the best_state['model'] weights dictionary and the xp.cfg configuration, serializes the config to a YAML string, tags the package with the current library version and an exported flag, creates the output directory if needed, and saves the result using torch.save().
The exported file is significantly smaller than the training checkpoint because it omits optimizer state, EMA buffers, scheduler state, and other training metadata. The resulting file can be loaded via the standard Audiocraft pretrained model API.
Usage
Import and call directly after training completes:
from audiocraft.utils.export import export_encodec
Typically invoked from a script or notebook after a Dora training run finishes.
Code Reference
Source Location
- Repository:
facebookresearch/audiocraft - File:
audiocraft/utils/export.py(lines 20--33)
Signature
def export_encodec(
checkpoint_path: Union[Path, str],
out_file: Union[Path, str],
) -> Union[Path, str]:
"""Export only the best state from the given EnCodec checkpoint.
This should be used if you trained your own EnCodec model.
Args:
checkpoint_path: Path to the full training checkpoint file.
out_file: Destination path for the exported lightweight checkpoint.
Returns:
The out_file path.
"""
pkg = torch.load(checkpoint_path, 'cpu')
new_pkg = {
'best_state': pkg['best_state']['model'],
'xp.cfg': OmegaConf.to_yaml(pkg['xp.cfg']),
'version': __version__,
'exported': True,
}
Path(out_file).parent.mkdir(exist_ok=True, parents=True)
torch.save(new_pkg, out_file)
return out_file
Import
from audiocraft.utils.export import export_encodec
Dependencies
torch-- fortorch.load()andtorch.save()checkpoint I/Oomegaconf-- forOmegaConf.to_yaml()serialization of the Hydra configaudiocraft.__version__-- library version string embedded in the export
I/O Contract
Inputs
| Name | Type | Description |
|---|---|---|
checkpoint_path |
Union[Path, str] |
Path to the full training checkpoint produced by CompressionSolver (via Dora experiment management). The checkpoint must contain keys best_state (with nested model dict) and xp.cfg (an OmegaConf DictConfig).
|
out_file |
Union[Path, str] |
Destination file path for the exported checkpoint. Parent directories are created automatically if they do not exist. |
Outputs
| Name | Type | Description |
|---|---|---|
| Return value | Union[Path, str] |
The out_file path, passed through for convenience.
|
| Exported file | File on disk | A torch.save() file containing: best_state (model state dict), xp.cfg (YAML string), version (str), exported (True).
|
Exported Checkpoint Structure
| Key | Type | Description |
|---|---|---|
best_state |
dict |
The model's state_dict() containing all encoder, decoder, and quantizer weights. Flattened from the training checkpoint's best_state['model'].
|
xp.cfg |
str |
The full Hydra configuration serialized as a YAML string via OmegaConf.to_yaml(). Contains model architecture parameters, training hyperparameters, dataset configuration, and all other settings.
|
version |
str |
The audiocraft.__version__ string at export time, enabling version compatibility tracking.
|
exported |
bool |
Always True. Used by the loading code to distinguish exported checkpoints from training checkpoints.
|
Usage Examples
Example 1: Exporting a Trained EnCodec Model
Exporting a training checkpoint to a lightweight inference checkpoint.
from audiocraft.utils.export import export_encodec
# Export from a Dora training checkpoint
export_encodec(
checkpoint_path='/checkpoints/xps/my_encodec_run/checkpoint.th',
out_file='/exports/my_encodec_model.th',
)
Example 2: Loading the Exported Model
After export, the model can be loaded via the pretrained model API.
from audiocraft import models
# Load the exported checkpoint
model = models.CompressionModel.get_pretrained('/exports/my_encodec_model.th')
model.eval()
# Use for inference
import torch
audio = torch.randn(1, 1, 32000) # 1 second at 32kHz
with torch.no_grad():
compressed = model.encode(audio)
reconstructed = model.decode(compressed)
Example 3: Exporting for Use as MusicGen Tokenizer
The exported compression model can be packaged alongside a MusicGen language model.
from audiocraft.utils.export import export_encodec, export_pretrained_compression_model
# First export the EnCodec model
export_encodec(
checkpoint_path='/checkpoints/xps/my_encodec_run/checkpoint.th',
out_file='/exports/my_encodec.th',
)
# Then package it for use with a MusicGen model
export_pretrained_compression_model(
pretrained_encodec='/exports/my_encodec.th',
out_file='/exports/musicgen_compression.th',
)