Module audiocraft.utils.export_legacy

Legacy functions used at the time of the first release, kept for referencd.

Functions

def export_encodec(checkpoint_path: str | pathlib.Path, out_file: str | pathlib.Path)

Expand source code

def export_encodec(checkpoint_path: tp.Union[Path, str], out_file: tp.Union[Path, str]):
    pkg = torch.load(checkpoint_path, 'cpu')
    new_pkg = {
        'best_state': pkg['ema']['state']['model'],
        'xp.cfg': OmegaConf.to_yaml(pkg['xp.cfg']),
        # The following params were NOT exported for the first release of MusicGen.
        'version': __version__,
        'exported': True,
    }
    Path(out_file).parent.mkdir(exist_ok=True, parents=True)
    torch.save(new_pkg, out_file)
    return out_file

def export_lm(checkpoint_path: str | pathlib.Path, out_file: str | pathlib.Path)

Expand source code

def export_lm(checkpoint_path: tp.Union[Path, str], out_file: tp.Union[Path, str]):
    pkg = torch.load(checkpoint_path, 'cpu')
    if pkg['fsdp_best_state']:
        best_state = pkg['fsdp_best_state']['model']
    else:
        best_state = pkg['best_state']['model']
    new_pkg = {
        'best_state': best_state,
        'xp.cfg': OmegaConf.to_yaml(_clean_lm_cfg(pkg['xp.cfg'])),
        # The following params were NOT exported for the first release of MusicGen.
        'version': __version__,
        'exported': True,
    }
    Path(out_file).parent.mkdir(exist_ok=True, parents=True)
    torch.save(new_pkg, out_file)
    return out_file