🚀 CSM-1B-HF

CSM-1B-HF是一个基于Hugging Face实现的语音模型，它能将文本转化为语音，为语音合成领域提供了新的解决方案。

🚀 快速开始

CSM-HF是 Sesame的对话语音模型（CSM） 在Hugging Face上的实现。它完全重写了 Sesame提供的PyTorch代码，并且从推理到训练都与Hugging Face的 transformers 库完全兼容。

✨ 主要特性

• 创建了 CSMModel 类。
• 用HF transformers的 LllamaModel 替换了骨干网络和解码器的TorchTune模型。
• 添加了一个处理器类，用于为模型准备输入。
• 增加了标签支持和 解码器训练摊销。
• 为模型类添加了 generate_frame 和 generate 方法，用于生成音频。
• 完全支持HuggingFace的 Trainer。

💻 使用示例

基础用法

你可以使用该模型从文本输入生成音频。以下是一个语音克隆的示例：

import torch
from modeling_csm import CSMModel
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer
from tokenizers.processors import TemplateProcessing
from moshi.models import loaders
from processor import CSMProcessor
import torchaudio

device = 'cuda'

def load_llama3_tokenizer():
    """
    https://github.com/huggingface/transformers/issues/22794#issuecomment-2092623992
    """
    tokenizer_name = "meta-llama/Llama-3.2-1B"
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
    bos = tokenizer.bos_token
    eos = tokenizer.eos_token
    tokenizer._tokenizer.post_processor = TemplateProcessing(
        single=f"{bos}:0 $A:0 {eos}:0",
        pair=f"{bos}:0 $A:0 {eos}:0 {bos}:1 $B:1 {eos}:1",
        special_tokens=[(f"{bos}", tokenizer.bos_token_id), (f"{eos}", tokenizer.eos_token_id)],
    )

    return tokenizer

text_tokenizer = load_llama3_tokenizer()

mimi_weight = hf_hub_download(loaders.DEFAULT_REPO, loaders.MIMI_NAME)
audio_tokenizer = loaders.get_mimi(mimi_weight, device=device)
audio_tokenizer.set_num_codebooks(32)

processor = CSMProcessor(text_tokenizer, audio_tokenizer)


def load_audio(path, target_sr):
    audio, sr = torchaudio.load(path)
    audio = audio.squeeze(0)
    if sr != target_sr:
        audio = torchaudio.functional.resample(audio, orig_freq=sr, new_freq=target_sr)
    return audio


model = CSMModel.from_pretrained("thomasgauthier/csm-1b-hf", torch_dtype=torch.bfloat16)
model.to('cuda')


inputs = processor(
    messages=[
        {
        "role": "speaker_0",
        "content": [
            {"type": "text", "text": "<AUDIO_CLIP_TRANSCRIPT>"},
            {"type": "audio"} # This placeholder is required for audio tokenization (it maps to the first element in the `audios` list passed to the processor)
        ]
    },
            {
        "role": "speaker_0",
        "content": [
            {"type": "text", "text": "Hello, this is voice cloning speaking"},
            # does not include audio as the model will generate it
        ]
    }
        ], 
    audios=[load_audio('AUDIO_CLIP_FOR_VOICE_CLONING.wav', audio_tokenizer.sample_rate)],
    return_tensors="pt"
)

import torch

with torch.inference_mode():
    # Generate up to 50 new frames
    gen_frames = model.generate(
        input_ids=inputs['input_ids'].cuda(),
        attention_mask=inputs['attention_mask'].cuda(),
        max_new_frames=50,
        topk=50,
        temperature=1.0,
        use_cache=True,
        stop_on_all_zeros=True,

    )

decoded_audio = audio_tokenizer.decode(gen_frames.permute(0, 2, 1)).squeeze(0).squeeze(0)

audio_array = (decoded_audio * 32768).to(torch.int16).cpu().numpy()

# Audio can be played with the following code:
# from IPython.display import Audio
# Audio(audio_array, rate=audio_tokenizer.sample_rate)

📚 详细文档

架构

模型架构在 ARCHITECTURE.md 中进行了讨论（由O1编写）。

训练

数据格式

CSM-HF期望训练数据采用JSONL格式，其中每行是一个包含对话的JSON对象。每个对话由以下部分组成：

• messages：消息对象数组，每个对象包含：
  • role：说话者标识符（例如，"speaker_0"，"speaker_1"）
  • content：内容对象数组，可以是：
    • 文本：{"type": "text", "text": "消息文本"}
    • 音频：{"type": "audio", "url": "音频文件路径.wav"}
• training_mask：布尔数组，指示哪些消息应用于训练（true）或作为上下文（false）

示例数据格式：

{
  "messages": [
    {
      "role": "speaker_0",
      "content": [
        {"type": "text", "text": "We have a chance for a new life here."},
        {"type": "audio", "url": "clips/example_audio.wav"}
      ]
    },
    {
      "role": "speaker_1",
      "content": [
        {"type": "text", "text": "Uncle?"},
        {"type": "audio", "url": "clips/response_audio.wav"}
      ]
    }
  ],
  "training_mask": [false, true]
}

训练过程

该模型采用两阶段自回归架构：

1. 骨干网络（帧间处理）：
   • 处理整个帧序列
   • 每个帧代表所有码本的组合嵌入
   • 处理话语之间的长距离依赖关系
2. 解码器（帧内处理）：
   • 一次处理一个帧
   • 按顺序生成32个码本（1个语义码本 + 31个声学码本）
   • 每个码本被视为序列中的一个标记

训练利用计算摊销技术：

• 第零个（语义）码本在所有帧上进行训练
• 其余码本（1 - 31）仅在 amortization_ratio 的帧上进行训练
• 这在保持质量的同时显著减少了内存使用

要训练模型，请运行以下命令：

python train.py \
  --train_file path/to/training_data.jsonl \
  --output_dir ./output \
  --num_train_epochs 3 \
  --per_device_train_batch_size 1 \
  --gradient_accumulation_steps 8 \
  --learning_rate 5e-6

📄 许可证

本项目采用 apache-2.0 许可证。

🔜 待办事项

• [x] 两阶段自回归架构实现
• [x] 多码本音频分词
• [x] 计算摊销以实现高效训练
• [x] 准备包含交错文本/音频的数据集
• [x] 自定义训练循环，分别计算骨干网络/解码器损失
• [x] 正确处理解码器摊销的epoch重复
• [x] 内存优化技术（混合精度、梯度累积）
• [ ] 支持LoRA以实现高效微调
• [ ] 使用 torch.compile 实现更快的推理
• [ ] 通过提示调整/前缀优化实现语音克隆
• [ ] 支持DPO
• [ ] 支持RL（GRPO、RLOO等）

🙏 致谢

特别感谢：

• Sesame Labs 提供原始架构设计和实现
• Hugging Face 提供Transformers库和训练基础设施
• Claude 和 ChatGPT 在文档和代码开发方面提供的帮助

本项目基于开源社区的研究和工具构建。感谢开源社区的协作精神，使这样的项目成为可能。