Transformers documentation

Pocket TTS

Transformers

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.2).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

该模型于 2026-10-20 贡献至 Hugging Face Transformers。

Pocket TTS

概述

Pocket TTS 是 Kyutai 推出的 1 亿参数文本转语音模型，属于一种 Continuous Audio Language Model，可在 CPU 上以快于实时的速度运行。该因果 Transformer 通过单步流头（one-step flow head）每帧（12.5 Hz）输出一个连续音频潜在变量，并由语音结束头（end-of-speech head）决定何时停止。

处理器包含 Pocket TTS 音频分词器，这是一种基于 Mimi 的编解码器，用于对语音提示进行编码，并将生成的潜在变量解码为 24 kHz 音频。预定义的声音以 KV 缓存的形式提供（voices/）。

发布了两个 flow heads：Lagrangian 自蒸馏（flow_type="lsd"，默认）和 drifting（flow_type="drifting"，english_drifting_26-09）。支持微调，请参阅 Training。

该模型由 Eustache Le Bihan 贡献。原始代码位于 kyutai-labs/pocket-tts。

使用技巧

使用 apply_chat_template() 准备输入（每段文本对应一条 user 消息，并设置 add_generation_prompt=True）。分词器的 normalizer 会按照检查点语言的相关设置，以与原始方式相同的方式处理文本。
generate 会返回 audio_latents 和 audio_latents_mask；请使用 PocketTTSProcessor.decode() 对其进行解码。
max_new_tokens 的默认值基于文本长度设定的预算，当语音结束（end-of-speech）的 logit 超过 eos_threshold 后不久，生成即停止。
do_sample=False 在温度为 0 时与原始结果一致。启用采样时，temperature（默认值 0.3）表示流起始噪声的方差。
使用 PocketTTSProcessor.split_text() 拆分长文本，每个分块对应一段对话。

用法

语音

语音可以通过三种方式指定。

1. 克隆录音：对话中的一个 audio 项。

from transformers import AutoModelForTextToWaveform, AutoProcessor

model_id = "eustlb/pocket-tts"
processor = AutoProcessor.from_pretrained(model_id, device_map="auto")
model = AutoModelForTextToWaveform.from_pretrained(model_id, device_map="auto")

url = "https://huggingface.co/datasets/hf-internal-testing/dummy-audio-samples/resolve/main/en-Alice_woman.wav"
conversation = [
    {
        "role": "user",
        "content": [
            {"type": "audio", "url": url},
            {"type": "text", "text": "Hello world. I am Kyutai's Pocket TTS."},
        ],
    }
]
inputs = processor.apply_chat_template(
    conversation, tokenize=True, return_dict=True, add_generation_prompt=True
).to(model.device)

outputs = model.generate(**inputs)
audio = processor.decode(outputs.audio_latents, outputs.audio_latents_mask)
processor.save_audio(audio, "cloned.wav")

2. 预计算潜在变量：一个空的 audio 项，潜在变量通过处理器传入。

from transformers import AutoModelForTextToWaveform, AutoProcessor
from transformers.audio_utils import load_audio

model_id = "eustlb/pocket-tts"
processor = AutoProcessor.from_pretrained(model_id, device_map="auto")
model = AutoModelForTextToWaveform.from_pretrained(model_id, device_map="auto")

# e.g. a recording encoded once by the processor's audio tokenizer,
# to reuse without re-encoding it
url = "https://huggingface.co/datasets/hf-internal-testing/dummy-audio-samples/resolve/main/en-Alice_woman.wav"
features = processor.feature_extractor(
    load_audio(url, sampling_rate=24000), sampling_rate=24000
).to(processor.audio_tokenizer.device)
# (num_frames, latent_dim)
voice = processor.audio_tokenizer.encode(**features).latents[0]

conversation = [
    {
        "role": "user",
        "content": [
            {"type": "audio"},
            {"type": "text", "text": "Hello world. I am Kyutai's Pocket TTS."},
        ],
    }
]
inputs = processor.apply_chat_template(
    conversation,
    tokenize=True,
    return_dict=True,
    add_generation_prompt=True,
    processor_kwargs={"audio_prompt_latents": [voice]},
).to(model.device)

outputs = model.generate(**inputs)
audio = processor.decode(outputs.audio_latents, outputs.audio_latents_mask)
processor.save_audio(audio, "latents.wav")

3. 缓存的上下文：语音提示的 KV 缓存，通过 past_key_values 传入。随附的语音（voices/）使用此形式。

from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoModelForTextToWaveform, AutoProcessor, DynamicCache

model_id = "eustlb/pocket-tts"
processor = AutoProcessor.from_pretrained(model_id, device_map="auto")
model = AutoModelForTextToWaveform.from_pretrained(model_id, device_map="auto")

voice = load_file(hf_hub_download(model_id, "voices/alba.safetensors"))
layers = []
for i, layer in enumerate(model.model.layers):
    # each layer's cache on that layer's device
    device = layer.self_attn.k_proj.weight.device
    layers.append(
        (voice[f"layers.{i}.keys"].to(device), voice[f"layers.{i}.values"].to(device))
    )
voice_cache = DynamicCache(layers, config=model.config)

conversation = [
    {
        "role": "user",
        "content": [{"type": "text", "text": "Hello world. I am Kyutai's Pocket TTS."}],
    }
]
inputs = processor.apply_chat_template(
    conversation, tokenize=True, return_dict=True, add_generation_prompt=True
).to(model.device)

outputs = model.generate(**inputs, past_key_values=voice_cache)
audio = processor.decode(outputs.audio_latents, outputs.audio_latents_mask)
processor.save_audio(audio, "cached.wav")

批处理与长文本

批量处理多个文本和语音，并将长文本拆分为多个片段：

import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoModelForTextToWaveform, AutoProcessor, DynamicCache

model_id = "eustlb/pocket-tts"
processor = AutoProcessor.from_pretrained(model_id, device_map="auto")
model = AutoModelForTextToWaveform.from_pretrained(model_id, device_map="auto")


def load_voices(names):
    """The shipped voices, one per conversation, as one cache, left-padded to the
    longest voice, and the mask of their positions.

    `generate` extends the cache: load it for every call.
    """
    voices = [
        load_file(hf_hub_download(model_id, f"voices/{name}.safetensors"))
        for name in names
    ]
    lengths = [voice["layers.0.keys"].shape[2] for voice in voices]
    length = max(lengths)
    layers = []
    for i, layer in enumerate(model.model.layers):
        device = layer.self_attn.k_proj.weight.device
        keys, values = (
            torch.cat(
                [
                    torch.nn.functional.pad(voice[f"layers.{i}.{kind}"], (0, 0, length - n, 0))
                    for voice, n in zip(voices, lengths)
                ]
            ).to(device)
            for kind in ("keys", "values")
        )
        layers.append((keys, values))
    mask = torch.tensor([[0] * (length - n) + [1] * n for n in lengths])
    return DynamicCache(layers, config=model.config), mask


text = (
    "A first sentence that is long enough to be worth splitting. "
    "Then a second one, and a third one for good measure."
)
chunks = processor.split_text(text)
conversations = [
    [{"role": "user", "content": [{"type": "text", "text": chunk}]}]
    for chunk in chunks + ["Hello world"]
]
inputs = processor.apply_chat_template(
    conversations, tokenize=True, return_dict=True, add_generation_prompt=True
).to(model.device)

voice_cache, voice_mask = load_voices(["alba"] * len(chunks) + ["jean"])
# the mask covers the cached voices too
inputs["attention_mask"] = torch.cat(
    [voice_mask.to(inputs["attention_mask"]), inputs["attention_mask"]], dim=-1
)
outputs = model.generate(**inputs, past_key_values=voice_cache)
audio = processor.decode(outputs.audio_latents, outputs.audio_latents_mask)
processor.save_audio(
    [torch.cat(audio[: len(chunks)]), audio[-1]], ["long.wav", "hello.wav"]
)

漂移检查点

用法相同，但使用检查点自身的语音：

from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoModelForTextToWaveform, AutoProcessor, DynamicCache

model_id = "eustlb/pocket-tts-english-drifting"
processor = AutoProcessor.from_pretrained(model_id, device_map="auto")
model = AutoModelForTextToWaveform.from_pretrained(model_id, device_map="auto")

voice = load_file(hf_hub_download(model_id, "voices/alba.safetensors"))
layers = []
for i, layer in enumerate(model.model.layers):
    device = layer.self_attn.k_proj.weight.device
    layers.append(
        (voice[f"layers.{i}.keys"].to(device), voice[f"layers.{i}.values"].to(device))
    )

conversation = [
    {
        "role": "user",
        "content": [{"type": "text", "text": "Hello world. I am Kyutai's Pocket TTS."}],
    }
]
inputs = processor.apply_chat_template(
    conversation, tokenize=True, return_dict=True, add_generation_prompt=True
).to(model.device)

outputs = model.generate(
    **inputs, past_key_values=DynamicCache(layers, config=model.config)
)
audio = processor.decode(outputs.audio_latents, outputs.audio_latents_mask)
processor.save_audio(audio, "drifting.wav")

训练

将目标音频放在 assistant 轮次中。处理器将其潜在表示作为 labels 返回，模型返回原始损失（LSD 流损失加上 config.eos_loss_weight 倍的语音结束损失）。语音提示通常是同一句话的开头部分。

import torch
from transformers import AutoModelForTextToWaveform, AutoProcessor

model_id = "eustlb/pocket-tts"
processor = AutoProcessor.from_pretrained(model_id, device_map="auto")
model = AutoModelForTextToWaveform.from_pretrained(model_id, device_map="auto")
model.train()

conversation = [
    {
        "role": "user",
        "content": [
            {"type": "audio", "audio": voice_prompt},
            {"type": "text", "text": target_text},
        ],
    },
    {"role": "assistant", "content": [{"type": "audio", "audio": target_audio}]},
]
inputs = processor.apply_chat_template(
    conversation, tokenize=True, return_dict=True
).to(model.device)
loss = model(**inputs).loss
loss.backward()

以录音或潜在表示的形式提供语音，而不是缓存。
对于无分类器引导，请删除 user 轮次中的 audio 或 text 项（原始实现以 0.2 的概率分别删除每一项）。
预计算的潜在变量应放入 processor_kwargs={"audio_prompt_latents": [...], "labels": [...]} 中，并使用空的 {"type": "audio"} 项作为占位符。
原始配方：AdamW，learning_rate=2e-5（新语言为 2e-4），adam_beta2=0.95，weight_decay=0.1，max_grad_norm=1.0，1000 步预热后保持恒定，权重 EMA（衰减系数 0.999）。

Transformers

Pocket TTS

概述

使用技巧

用法

语音

批处理与长文本

漂移检查点

训练

PocketTTSConfig

class transformers.PocketTTSConfig

PocketTTSFlowHeadConfig

类 transformers.PocketTTSFlowHeadConfig

PocketTTSFeatureExtractor

class transformers.PocketTTSFeatureExtractor

__call__

PocketTTSProcessor

class transformers.PocketTTSProcessor

__call__

split_text

decode

save_audio

PocketTTSModel

class transformers.PocketTTSModel

forward

PocketTTSForConditionalGeneration

class transformers.PocketTTSForConditionalGeneration

forward

generate

sample_latent

denormalize_latents

PocketTTSGenerateOutput

class transformers.models.pocket_tts.generation_pocket_tts.PocketTTSGenerateOutput

call

call