speechbrain.lobes.models.discrete.wavtokenizer 模块

该模块支持集成预训练的WavTokenizer。

请注意，您需要 pip 安装 git+https://github.com/Tomiinek/WavTokenizer 才能使用此模块。

仓库: https://github.com/jishengpeng/WavTokenizer/ 论文: https://arxiv.org/abs/2408.16532

Authors

  • Pooneh Mousavi 2024

摘要

类：

WavTokenizer

该模块支持集成预训练的WavTokenizer模型，这是一种用于音频语言建模的具有单一码本的离散编解码模型。

参考

class speechbrain.lobes.models.discrete.wavtokenizer.WavTokenizer(source, save_path=None, config='wavtokenizer_smalldata_frame40_3s_nq1_code4096_dim512_kmeans200_attn.yaml', checkpoint='WavTokenizer_small_600_24k_4096.ckpt', sample_rate=24000, freeze=True)[source]

基础：Module

该模块支持集成预训练的WavTokenizer模型，这是一种用于音频语言建模的具有单码本的离散编解码模型。

Source paper:

https://arxiv.org/abs/2408.16532

你需要 pip 安装 git+https://github.com/Tomiinek/WavTokenizer 来使用这个模块。

该代码改编自官方的WavTokenizer仓库： https://github.com/jishengpeng/WavTokenizer/

Parameters:

  • source (str) – 一个 HuggingFace 仓库标识符或路径

  • save_path (str) – 预训练模型将被保存的位置

  • config (str) – HF配置文件的名称。

  • 检查点 (str) – HF检查点文件的名称。

  • sample_rate (int (默认值: 24000)) – 音频采样率

  • freeze (bool) – 模型是否会被冻结（例如，如果用作训练另一个模型的一部分，则不可训练）

Example

>>> model_hub = "novateur/WavTokenizer"
>>> save_path = "savedir"
>>> config="wavtokenizer_smalldata_frame40_3s_nq1_code4096_dim512_kmeans200_attn.yaml"
>>> checkpoint="WavTokenizer_small_600_24k_4096.ckpt"
>>> model = WavTokenizer(model_hub, save_path,config=config,checkpoint=checkpoint)
>>> audio = torch.randn(4, 48000)
>>> length = torch.tensor([1.0, .5, .75, 1.0])
>>> tokens, embs= model.encode(audio)
>>> tokens.shape
torch.Size([4, 1, 80])
>>> embs.shape
torch.Size([4, 80, 512])
>>> rec = model.decode(tokens)
>>> rec.shape
torch.Size([4, 48000])
forward(inputs)[source]

将输入音频编码为令牌和嵌入，并从令牌解码音频

Parameters:

inputs (torch.Tensor) – 一个 (Batch x Samples) 的音频张量

Returns:

  • tokens (torch.Tensor) – 一个 (Batch x Tokens x Heads) 的音频令牌张量

  • emb (torch.Tensor) – 来自模型量化器的原始向量嵌入

  • audio (torch.Tensor) – 重构的音频

encode(inputs)[source]

将输入音频编码为标记和嵌入

Parameters:

inputs (torch.Tensor) – 一个 (Batch x Samples) 或 (Batch x Channel x Samples) 的音频张量

Returns:

  • tokens (torch.Tensor) – 一个 (Batch x NQ x Length) 的音频令牌张量

  • emb (torch.Tensor) – 来自模型量化器的原始向量嵌入

decode(tokens)[source]

从标记解码音频

Parameters:

tokens (torch.Tensor) – 一个 (Batch x NQ x Length) 的音频令牌张量

Returns:

audio – 重建的音频

Return type:

torch.Tensor