{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# ๐Ÿ’ง LFM3-Audio - Teach the model a new voice with Transformers\n", "\n", "This tutorial demonstrates how to fine-tune our LFM3 speech model [`LiquidAI/LFM3-Audio-2B`](https://huggingface.co/LiquidAI/LFM3-Audio-2B)\n", "so it speaks with a **new, consistent voice**.\n", "\n", "LFM3-Audio generates speech itself: for every 80 ms frame of audio it predicts 12 codebook tokens of the Liquid codec, alongside the text tokens.\n", "Fine-tuning therefore trains **two losses at once**: the usual next-token loss on text, and an **audio loss** on the speech frames.\n", "The model computes both for you; you only need to build the right batches.\n", "\n", "We use [LJSpeech](https://huggingface.co/datasets/MikhailT/lj-speech), ~24 hours of a single speaker reading non-fiction books,\n", "and teach the model to use this voice when asked for *\"Linda's voice\"*, a voice it has never heard of.\n", "\n", "## ๐ŸŽฏ What You'll Find:\n", "- **Data preparation** - turn (text, recording) pairs into training conversations with speech targets\n", "- **LoRA + SFT** - parameter-efficient fine-tuning of the language model and the audio decoder with the ๐Ÿค— `Trainer`\n", "- **Evaluation** - listen before/after, check the speech is still intelligible, and measure how close the voice is to the target speaker\n", "\n", "## ๐Ÿ“‹ Prerequisites:\n", "- **GPU Runtime**: Select GPU in `Runtime` โ†’ `Change runtime type` (an A100 is recommended; ~15 GB of GPU memory is used with the default settings)\n", "- **Hugging Face Account**: For accessing models and datasets" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# ๐Ÿ“ฆ Installation & Setup\n", "\n", "First, let's install all the required packages:" ] }, { "cell_type": "code", "metadata": {}, "source": [ "!pip install -qqq \"transformers==5.20.2\" \"torchao==0.18.0\" accelerate datasets peft librosa soundfile speechbrain --progress-bar off" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Let's now verify the packages are installed correctly" ] }, { "cell_type": "code", "metadata": {}, "source": [ "import os\n", "import torch\n", "import transformers\n", "import peft\n", "os.environ[\"WANDB_DISABLED\"] = \"true\"\n", "\n", "print(f\"๐Ÿ“ฆ PyTorch version: {torch.__version__}\")\n", "print(f\"๐Ÿค— Transformers version: {transformers.__version__}\")\n", "print(f\"๐Ÿงฉ PEFT version: {peft.__version__}\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Loading the model from Transformers ๐Ÿค—\n", "\n", "The **Liquid codec** converts between 24 kHz waveforms and the discrete speech codes the model reads and writes.\n", "It lives in the `audio_tokenizer/` folder of the model repository; the processor loads it as `processor.audio_tokenizer`, and we move it to the GPU." ] }, { "cell_type": "code", "metadata": {}, "source": [ "from transformers import AutoModelForMultimodalLM, AutoProcessor\n", "\n", "model_id = \"LiquidAI/LFM3-Audio-2B\"\n", "device = \"cuda\"\n", "\n", "print(\"๐Ÿง  Loading model...\")\n", "model = AutoModelForMultimodalLM.from_pretrained(\n", " model_id, dtype=torch.bfloat16, device_map=device, trust_remote_code=True\n", ")\n", "\n", "print(\"๐Ÿ“š Loading processor and audio codec...\")\n", "processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)\n", "codec = processor.audio_tokenizer.to(device)\n", "tokenizer = processor.tokenizer\n", "\n", "print(\"\\nโœ… Model loaded successfully!\")\n", "print(f\"๐Ÿ”ข Parameters: {model.num_parameters():,}\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## ๐Ÿ› ๏ธ Helpers for speech in and out\n", "\n", "`tts()` reads a sentence aloud with a given system prompt, and `asr()` transcribes a waveform.\n", "We will use the model's own ASR to check that the fine-tuned voice stays intelligible." ] }, { "cell_type": "code", "metadata": {}, "source": [ "import numpy as np\n", "import soundfile as sf\n", "from IPython.display import Audio, display\n", "\n", "SR = 24_000 # the codec works at 24 kHz\n", "\n", "\n", "@torch.inference_mode()\n", "def tts(model, text, system_prompt, seed=0):\n", " torch.manual_seed(seed)\n", " conversation = [\n", " {\"role\": \"system\", \"content\": system_prompt},\n", " {\"role\": \"user\", \"content\": text},\n", " {\"role\": \"assistant\", \"content\": \"<|audio_start|>\"},\n", " ]\n", " inputs = processor.apply_chat_template(\n", " conversation, continue_final_message=True, tokenize=True, return_dict=True, return_tensors=\"pt\"\n", " ).to(device)\n", " out = model.generate(\n", " **inputs,\n", " max_new_tokens=600,\n", " audio_temperature=1.0,\n", " audio_top_k=128,\n", " audio_repetition_penalty=1.1,\n", " audio_repetition_window=4096,\n", " )\n", " codes = out.audio_codes[:1, :, : out.audio_lengths[0]]\n", " return codec.decode(codes).audio_values[0].float().cpu().numpy().squeeze()\n", "\n", "\n", "@torch.inference_mode()\n", "def asr(model, waveform):\n", " sf.write(\"_asr.wav\", waveform, SR)\n", " inputs = processor.apply_chat_template(\n", " [\n", " {\"role\": \"system\", \"content\": \"Perform ASR in English.\"},\n", " {\"role\": \"user\", \"content\": [{\"type\": \"audio\", \"audio\": \"_asr.wav\"}]},\n", " ],\n", " add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors=\"pt\",\n", " ).to(device)\n", " out = model.generate(**inputs, max_new_tokens=256)\n", " return processor.decode(out.sequences[0, inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True)\n", "\n", "\n", "def play(waveform, sr=SR):\n", " display(Audio(waveform, rate=sr))" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# ๐Ÿ“š Load the dataset\n", "\n", "We use one of the 8 shards of [MikhailT/lj-speech](https://huggingface.co/datasets/MikhailT/lj-speech) (~1,600 clips, ~3 hours).\n", "We keep the audio as raw bytes (`decode=False`) and decode it with `soundfile` ourselves." ] }, { "cell_type": "code", "metadata": {}, "source": [ "import io\n", "from datasets import Audio as AudioFeature, load_dataset\n", "\n", "N_TRAIN, N_EVAL = 1200, 20\n", "\n", "# verification_mode=\"no_checks\": we load one shard, which doesn't match the \"full\" split in the dataset card\n", "raw_ds = load_dataset(\"MikhailT/lj-speech\", data_files=\"data/full-00000-of-00008-*.parquet\", split=\"train\", verification_mode=\"no_checks\")\n", "raw_ds = raw_ds.cast_column(\"audio\", AudioFeature(decode=False)).shuffle(seed=0).select(range(N_TRAIN + N_EVAL))\n", "\n", "print(\"โœ… Dataset loaded:\", raw_ds)\n", "sample = raw_ds[0]\n", "waveform, sr = sf.read(io.BytesIO(sample[\"audio\"][\"bytes\"]))\n", "print(f\"\\n๐Ÿ“ Single Sample: {sample['normalized_text']}\")\n", "play(waveform, sr)" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## ๐ŸŽง Before fine-tuning\n", "\n", "We will ask for **\"Linda's voice\"**. The model knows how to read in a female voice, but has no idea who Linda is,\n", "so let's hear what it does before training. Keep these clips to compare with after training." ] }, { "cell_type": "code", "metadata": {}, "source": [ "SYSTEM_PROMPT = \"Perform TTS in American English using Linda's voice.\"\n", "\n", "test_sentences = [\n", " \"The museum opens at nine o'clock and closes shortly after sunset.\",\n", " \"Please remember to bring your umbrella tomorrow.\",\n", " \"Our team shipped the new feature ahead of schedule.\",\n", " \"The river was calm, and the boats drifted slowly past the old bridge.\",\n", " \"Can you tell me where the nearest train station is?\",\n", "]\n", "\n", "before = [tts(model, s, SYSTEM_PROMPT) for s in test_sentences]\n", "for sentence, wav in zip(test_sentences[:2], before[:2]):\n", " print(\"๐Ÿ—ฃ๏ธ\", sentence)\n", " play(wav)" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# ๐Ÿ”ข Encode the recordings into speech codes\n", "\n", "The model learns to predict the **codec codes** of each recording, so we encode every clip once, up front, with\n", "`processor.encode_audio()` (24 kHz mono in, a `(num_frames, 12)` array of codes out, one frame per 80 ms).\n", "This takes a couple of minutes." ] }, { "cell_type": "code", "metadata": {}, "source": [ "import librosa\n", "\n", "\n", "def encode(example):\n", " wav, sr = sf.read(io.BytesIO(example[\"audio\"][\"bytes\"]), dtype=\"float32\")\n", " wav = librosa.resample(wav, orig_sr=sr, target_sr=SR)\n", " example[\"codes\"] = processor.encode_audio(wav).tolist()\n", " return example\n", "\n", "\n", "encoded_ds = raw_ds.map(encode, remove_columns=[\"audio\"], desc=\"Encoding speech\")\n", "train_dataset = encoded_ds.select(range(N_TRAIN))\n", "eval_dataset = encoded_ds.select(range(N_TRAIN, N_TRAIN + N_EVAL))\n", "\n", "print(\"โœ… Speech encoded:\")\n", "print(f\" ๐Ÿ“š Train samples: {len(train_dataset)}\")\n", "print(f\" ๐Ÿงช Eval samples: {len(eval_dataset)}\")\n", "print(f\" ๐ŸŽต Codes of the first clip: {np.array(train_dataset[0]['codes']).shape}\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Collate function\n", "\n", "Each training example is a TTS conversation whose assistant turn is the speech itself:\n", "\n", "```\n", "<|im_start|>system\n", "Perform TTS in American English using Linda's voice.<|im_end|>\n", "<|im_start|>user\n", "{text}<|im_end|>\n", "<|im_start|>assistant\n", "<|audio_start|><|audio_out|> ร— (num_frames + 1)<|im_end|>\n", "```\n", "\n", "Each `<|audio_out|>` placeholder stands for one speech frame; its embedding is replaced by the frame's codes (`audio_out_codes`).\n", "The extra last frame is an **end-of-speech** frame, after which the model switches back to text and closes the turn.\n", "\n", "The model takes two kinds of labels:\n", "- `labels` (text): `-100` everywhere except `<|im_end|>` and the newline after it, so the model learns to end the turn after speaking\n", "- `audio_labels` (speech): the codes of every frame, one row per `<|audio_out|>` placeholder, in order across the batch\n", "\n", "`<|audio_start|>` is not trained: at inference we prefill it, as in the `tts()` helper above." ] }, { "cell_type": "code", "metadata": {}, "source": [ "AUDIO_OUT = processor.audio_out_token\n", "EOS_CODE = processor.audio_eos_code\n", "\n", "\n", "def create_collate_fn(system_prompt):\n", " def collate_fn(examples):\n", " input_ids, labels, codes = [], [], []\n", " for example in examples:\n", " prompt = processor.apply_chat_template(\n", " [{\"role\": \"system\", \"content\": system_prompt}, {\"role\": \"user\", \"content\": example[\"normalized_text\"]}],\n", " add_generation_prompt=True, tokenize=False,\n", " )\n", " frames = torch.tensor(example[\"codes\"])\n", " frames = torch.cat([frames, torch.full((1, frames.shape[1]), EOS_CODE)]) # end-of-speech frame\n", " prompt_ids = tokenizer(prompt, add_special_tokens=False)[\"input_ids\"]\n", " answer_ids = tokenizer(\"<|audio_start|>\" + AUDIO_OUT * len(frames) + \"<|im_end|>\\n\", add_special_tokens=False)[\"input_ids\"]\n", " answer_labels = [t if t != processor.audio_out_token_id else -100 for t in answer_ids]\n", " answer_labels[0] = -100 # <|audio_start|> is prefilled at inference\n", " input_ids.append(prompt_ids + answer_ids)\n", " labels.append([-100] * len(prompt_ids) + answer_labels)\n", " codes.append(frames)\n", "\n", " max_len = max(len(ids) for ids in input_ids)\n", " pad = lambda seq, value: seq + [value] * (max_len - len(seq))\n", " batch = {\n", " \"input_ids\": torch.tensor([pad(ids, tokenizer.pad_token_id) for ids in input_ids]),\n", " \"attention_mask\": torch.tensor([pad([1] * len(ids), 0) for ids in input_ids]),\n", " \"labels\": torch.tensor([pad(lab, -100) for lab in labels]),\n", " \"audio_out_codes\": torch.cat(codes),\n", " }\n", " batch[\"audio_labels\"] = batch[\"audio_out_codes\"].clone()\n", " return batch\n", " return collate_fn\n", "\n", "\n", "collate_fn = create_collate_fn(SYSTEM_PROMPT)\n", "\n", "batch = collate_fn([train_dataset[0], train_dataset[1]])\n", "print({k: tuple(v.shape) for k, v in batch.items()})\n", "with torch.no_grad():\n", " out = model(**{k: v.to(device) for k, v in batch.items()})\n", "# the only supervised text tokens are <|im_end|> and the newline after it, which the model already predicts well: almost all the loss is on speech\n", "print(f\"๐Ÿ“‰ Loss before training: {out.loss.item():.3f} (text {out.text_loss.item():.3f}, audio {out.audio_loss.item():.3f})\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Wrap the model with PEFT\n", "\n", "We add LoRA adapters to the attention projections (`q_proj`, `k_proj`, `v_proj`, `o_proj`), gated short convolutions (`in_proj`, `out_proj`) and MLP (`w1`, `w2`, `w3`) of the **language model** (which decides what to say and plans the speech)\n", "and the **depth decoder** (the small transformer that turns each hidden state into the 12 codes of a speech frame).\n", "The audio encoder, used only for speech *input*, stays frozen. LFM3's attention layers also have a per-head output gate, `gate_proj`; adding it to the targets is a reasonable choice. Feel free to modify the `r` (rank) value:\n", "- Higher โ†’ better approximation of full fine-tuning\n", "- Lower โ†’ needs even less compute resources" ] }, { "cell_type": "code", "metadata": {}, "source": [ "from peft import LoraConfig, get_peft_model\n", "\n", "peft_config = LoraConfig(\n", " r=16,\n", " lora_alpha=32,\n", " lora_dropout=0.05,\n", " bias=\"none\",\n", " # optional: add \"|gate_proj\" to the group for LFM3's per-head attention gate\n", " target_modules=r\"model\\.(language_model|depth_decoder)\\..*\\.(q_proj|k_proj|v_proj|o_proj|in_proj|out_proj|w1|w2|w3)\",\n", ")\n", "\n", "model = get_peft_model(model, peft_config)\n", "model.print_trainable_parameters()" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# ๐Ÿš€ Launch Training\n", "\n", "We are now ready to launch an SFT run with the ๐Ÿค— `Trainer`. Our collate function builds the batches, so we keep all dataset columns\n", "(`remove_unused_columns=False`). Two epochs over 1,200 clips take a few minutes on an A100.\n", "\n", "> **Why not TRL's `SFTTrainer`?** It computes its own next-token loss on the text logits, which ignores the speech frames.\n", "> The plain `Trainer` uses the loss returned by the model, which combines the text loss and the audio loss." ] }, { "cell_type": "code", "metadata": {}, "source": [ "from transformers import Trainer, TrainingArguments\n", "\n", "training_args = TrainingArguments(\n", " output_dir=\"lfm3-audio-linda\",\n", " num_train_epochs=2,\n", " per_device_train_batch_size=8,\n", " per_device_eval_batch_size=8,\n", " learning_rate=2e-4,\n", " lr_scheduler_type=\"cosine\",\n", " warmup_steps=10,\n", " logging_steps=10,\n", " eval_strategy=\"steps\",\n", " eval_steps=50,\n", " save_strategy=\"no\",\n", " bf16=True,\n", " remove_unused_columns=False,\n", " report_to=\"none\",\n", ")\n", "\n", "print(\"๐Ÿ—๏ธ Creating trainer...\")\n", "trainer = Trainer(\n", " model=model,\n", " args=training_args,\n", " train_dataset=train_dataset,\n", " eval_dataset=eval_dataset,\n", " data_collator=collate_fn,\n", ")\n", "\n", "print(\"\\n๐Ÿš€ Starting SFT training...\")\n", "trainer.train()\n", "print(\"๐ŸŽ‰ SFT training completed!\")\n", "print(f\"๐Ÿ’พ Peak GPU memory: {torch.cuda.max_memory_allocated() / 1e9:.1f} GB\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# ๐ŸŽง After fine-tuning\n", "\n", "Same sentences, same prompt, same seeds. Listen and compare with the clips above." ] }, { "cell_type": "code", "metadata": {}, "source": [ "model.eval()\n", "after = [tts(model, s, SYSTEM_PROMPT) for s in test_sentences]\n", "for sentence, wav_before, wav_after in zip(test_sentences, before, after):\n", " print(\"๐Ÿ—ฃ๏ธ\", sentence)\n", " print(\" before:\"); play(wav_before)\n", " print(\" after:\"); play(wav_after)" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## โœ… Is it still intelligible?\n", "\n", "A new voice is only useful if the words are still right. We transcribe the generated speech with the model's own ASR." ] }, { "cell_type": "code", "metadata": {}, "source": [ "for sentence, wav in zip(test_sentences, after):\n", " print(f\"๐Ÿ“ {sentence}\\n๐ŸŽ™๏ธ {asr(model, wav)}\\n\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## ๐Ÿ“ How close is the voice to the target speaker?\n", "\n", "We embed each clip with a speaker-verification model ([ECAPA-TDNN](https://huggingface.co/speechbrain/spkrec-ecapa-voxceleb))\n", "and compute the cosine similarity to the average embedding of held-out LJSpeech recordings. Higher means closer to the target voice;\n", "other recordings of the real speaker typically score around 0.84." ] }, { "cell_type": "code", "metadata": {}, "source": [ "from speechbrain.inference.speaker import EncoderClassifier\n", "\n", "speaker_encoder = EncoderClassifier.from_hparams(\"speechbrain/spkrec-ecapa-voxceleb\", savedir=\"ecapa\", run_opts={\"device\": device})\n", "\n", "\n", "def speaker_embedding(wav, sr):\n", " wav = torch.tensor(librosa.resample(wav, orig_sr=sr, target_sr=16_000), dtype=torch.float32)[None]\n", " return torch.nn.functional.normalize(speaker_encoder.encode_batch(wav).squeeze().cpu(), dim=0)\n", "\n", "\n", "reference_ds = load_dataset(\"MikhailT/lj-speech\", data_files=\"data/full-00001-of-00008-*.parquet\", split=\"train\", verification_mode=\"no_checks\")\n", "reference_ds = reference_ds.cast_column(\"audio\", AudioFeature(decode=False)).select(range(30))\n", "reference = torch.stack([speaker_embedding(*sf.read(io.BytesIO(r[\"audio\"][\"bytes\"]), dtype=\"float32\")) for r in reference_ds])\n", "reference = torch.nn.functional.normalize(reference.mean(0), dim=0)\n", "\n", "similarity = lambda clips: np.mean([float(speaker_embedding(w, SR) @ reference) for w in clips])\n", "print(f\"๐ŸŽฏ Speaker similarity to LJSpeech before: {similarity(before):.2f} after: {similarity(after):.2f}\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# ๐Ÿ’พ Save the model\n", "\n", "Save the LoRA adapter (a few tens of MB), or merge it into the base weights to obtain a regular checkpoint that loads exactly like the original model." ] }, { "cell_type": "code", "metadata": {}, "source": [ "model.save_pretrained(\"./lfm3-audio-linda-lora\")\n", "print(\"๐Ÿ’พ Adapter saved to: ./lfm3-audio-linda-lora\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": {}, "source": [ "print(\"๐Ÿ”„ Merging LoRA weights...\")\n", "merged = model.merge_and_unload()\n", "merged.save_pretrained(\"./lfm3-audio-linda\")\n", "processor.save_pretrained(\"./lfm3-audio-linda\")\n", "print(\"๐Ÿ’พ Model saved to: ./lfm3-audio-linda\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# ๐Ÿš€ Next steps\n", "\n", "- **Your own voice**: replace LJSpeech with recordings of the voice you want; a few hundred clips of a single speaker with accurate transcripts are enough to start. Keep the clips under ~20 seconds.\n", "- **Several custom voices**: give each speaker its own name in the system prompt and mix them in the same training set.\n", "- **Other tasks**: the same collate function pattern works for any task with speech output (e.g. speech-to-speech translation); for text-only outputs, drop `audio_out_codes`/`audio_labels` and put the audio in the user turn instead." ] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "A100", "provenance": [], "machine_shape": "hm" }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }