orpheus

第7章 · 模型后训练 · 配套项目 chapter7/orpheus

项目说明

Orpheus TTS - Text-to-Speech Fine-tuning with Unsloth

This project demonstrates how to fine-tune the Orpheus 3B text-to-speech model using Unsloth for efficient training and inference.

Overview

Orpheus is a text-to-speech (TTS) model that converts text into natural-sounding speech. This implementation uses: - Unsloth for efficient LoRA fine-tuning (30% less VRAM, 2x larger batch sizes) - SNAC (Stochastic Neural Audio Codec) for audio tokenization at 24kHz - Hugging Face Transformers for model training and inference

Features

• ✅ Fine-tune Orpheus 3B model with minimal GPU memory
• ✅ Support for single-speaker and multi-speaker TTS
• ✅ Expressive speech with emotion tags (laugh, sigh, gasp, etc.)
• ✅ Audio generation with customizable parameters
• ✅ Automatic WAV file export of generated speech
• ✅ LoRA adapter training for efficient fine-tuning

Installation

Prerequisites

• Python 3.8+
• CUDA-compatible GPU (recommended: 16GB+ VRAM)
• CUDA toolkit installed

Install Dependencies

IMPORTANT: The datasets package version must be between 3.4.1 and 4.0.0 for compatibility.

pip install sentencepiece protobuf "datasets>=3.4.1,<4.0.0" "huggingface_hub>=0.34.0" hf_transfer
pip install -r requirements.txt

For Colab or specific environments, you may need:

pip install --no-deps bitsandbytes accelerate xformers peft trl triton cut_cross_entropy unsloth_zoo
pip install --no-deps unsloth

Project Structure

orpheus/
├── orpheus_sft_unsloth.py    # Full training + inference script
├── inference.py               # Standalone inference script
├── requirements.txt           # Python dependencies
├── README.md                  # This file
├── lora_model/               # Saved LoRA adapters (after training)
└── generated_audio/          # Generated audio outputs

Usage

Training

Fine-tune the model on your dataset:

python orpheus_sft_unsloth.py

The training script will: 1. Load the pre-trained Orpheus 3B model 2. Apply LoRA adapters for efficient fine-tuning 3. Train on the MrDragonFox/Elise dataset (or your custom dataset) 4. Save the fine-tuned LoRA adapters to lora_model/

Training Parameters: - Batch size: 1 per device - Gradient accumulation: 4 steps - Learning rate: 2e-4 - Training steps: 60 (configurable) - LoRA rank: 64

Inference

Generate speech from text using the fine-tuned model:

python inference.py

Or use it as a module:

from inference import OrpheusInference

# Initialize the model
tts = OrpheusInference(
    model_path="unsloth/orpheus-3b-0.1-ft",
    lora_path="lora_model"  # Optional: load fine-tuned adapters
)

# Generate speech with emotion tags
prompts = [
    "Hey there my name is Elise, <giggles> and I'm a speech generation model.",
    "I missed you <laugh> so much! It's been way too long.",
    "This is absolutely amazing <gasp> I can't believe it worked!"
]

audio_files = tts.generate(
    prompts=prompts,
    output_dir="generated_audio",
    temperature=0.6,
    top_p=0.95,
    max_new_tokens=1200
)

print(f"Generated {len(audio_files)} audio files")

Emotion Tags (Expressive Speech)

Orpheus supports special emotion/expression tags to create more expressive and natural-sounding speech:

Supported Tags: - Laughter: <laugh>, <giggles>, <chuckle> - Emotions: <sigh>, <gasp> - Physical sounds: <yawn>, <cough>, <sniffle>, <groan>

Usage:

prompts = [
    "Hey there <giggles> welcome to my channel!",
    "I missed you <laugh> so much!",
    "That's so beautiful <sigh> it brings back memories.",
    "I'm so tired <yawn> after working all day.",
    "This is incredible <gasp> I can't believe my eyes!"
]

audio_files = tts.generate(prompts=prompts)

How it works: - Tags are enclosed in angle brackets: <tag> - During training, the model learns to associate these tags with audio patterns - The Elise dataset contains 336 occurrences of "laughs", 156 of "sighs", etc. - If your custom dataset lacks these tags, you can manually annotate transcripts where the audio contains those expressions

Multi-Speaker Support

For multi-speaker models, specify the voice name:

tts.generate(
    prompts=["This is a test <laugh> with emotion."],
    voice="speaker_name"  # Specify the speaker
)

Dataset Format

The training script expects datasets with the following structure:

Single-speaker: - text: The text to be spoken - audio: Audio file with array and sampling_rate fields

Multi-speaker: - source: Speaker identifier - text: The text to be spoken - audio: Audio file with array and sampling_rate fields

Example dataset: MrDragonFox/Elise

Model Architecture

Audio Tokenization (SNAC)

• Sample rate: 24kHz
• Multi-layer hierarchical codec (3 layers)
• 7 tokens per frame (1 + 2 + 4 from the three layers)
• Duplicate frame removal for efficiency

Special Tokens

• Start of human: 128259
• End of human: 128260
• Start of AI: 128261
• End of AI: 128262
• Start of speech: 128257
• End of speech: 128258
• Pad token: 128263

Output

Generated audio files are saved as WAV files in the generated_audio/ directory: - Format: WAV (PCM) - Sample rate: 24kHz - Naming: output_0.wav, output_1.wav, etc.

Memory Usage

Typical memory requirements: - Training: ~12-16GB VRAM (with LoRA and 4-bit quantization) - Inference: ~8-10GB VRAM - CPU RAM: ~16GB recommended

Troubleshooting

Common Issues

1. Dataset version error:

Ensure datasets>=3.4.1,<4.0.0 is installed
pip install "datasets>=3.4.1,<4.0.0" --force-reinstall

2. CUDA out of memory: - Reduce max_new_tokens during inference - Use load_in_4bit=True when loading the model - Reduce batch size or enable gradient checkpointing

3. Multi-GPU issues: - Set CUDA_VISIBLE_DEVICES=0 to use only one GPU - per_device_train_batch_size >1 may cause errors on multi-GPU setups

Performance Tips

1. Faster Inference: Use FastLanguageModel.for_inference(model) before generation
2. Memory Optimization: Enable 4-bit quantization with load_in_4bit=True
3. Better Quality: Adjust temperature (0.4-0.8) and top_p (0.9-0.95) parameters
4. Longer Audio: Increase max_new_tokens (each ~7 tokens = 1 audio frame)

Resources

• Unsloth Documentation
• Unsloth TTS Guide
• SNAC Model
• Orpheus Model
• Discord Support

License

This project uses models and libraries with their respective licenses: - Unsloth: Apache 2.0 - Transformers: Apache 2.0 - Orpheus model: Check model card on Hugging Face

Citation

If you use this code in your research, please cite:

@misc{orpheus-tts-unsloth,
  title={Orpheus TTS Fine-tuning with Unsloth},
  author={Unsloth AI Team},
  year={2024},
  url={https://github.com/unslothai/unsloth}
}

Contributing

Contributions are welcome! Please feel free to submit issues or pull requests.

Acknowledgments

• Thanks to Etherl for creating the original notebook
• Unsloth AI team for the efficient training framework
• Hugging Face for hosting models and datasets

源代码

inference.py

#!/usr/bin/env python3
"""
Orpheus TTS Inference Script

This script provides a standalone inference interface for the Orpheus text-to-speech model.
It supports both single-speaker and multi-speaker TTS generation.

Usage:
    python inference.py

Or import as a module:
    from inference import OrpheusInference
    tts = OrpheusInference()
    audio_files = tts.generate(prompts=["Hello world"])
"""

import os
import torch
import torchaudio.transforms as T
from unsloth import FastLanguageModel
from snac import SNAC
import soundfile as sf
from typing import List, Optional


class OrpheusInference:
    """
    Orpheus TTS Inference Engine

    Supports expressive speech generation with emotion tags like:
    <laugh>, <giggles>, <chuckle>, <sigh>, <cough>, <sniffle>, 
    <groan>, <yawn>, <gasp>, etc.

    Example usage:
        tts = OrpheusInference()
        tts.generate(prompts=["I missed you <laugh> so much!"])
    """

    def __init__(
        self,
        model_path: str = "unsloth/orpheus-3b-0.1-ft",
        lora_path: Optional[str] = None,
        max_seq_length: int = 2048,
        load_in_4bit: bool = False,
        device: str = "cuda"
    ):
        """
        Initialize the Orpheus TTS inference engine.

        Args:
            model_path: HuggingFace model path or local path
            lora_path: Optional path to LoRA adapters
            max_seq_length: Maximum sequence length for the model
            load_in_4bit: Whether to use 4-bit quantization
            device: Device to run inference on ('cuda' or 'cpu')
        """
        self.device = device
        self.sample_rate = 24000  # SNAC model uses 24kHz

        print(f"Loading model from {model_path}...")
        self.model, self.tokenizer = FastLanguageModel.from_pretrained(
            model_name=model_path,
            max_seq_length=max_seq_length,
            dtype=None,
            load_in_4bit=load_in_4bit,
        )

        # Load LoRA adapters if provided
        if lora_path:
            print(f"Loading LoRA adapters from {lora_path}...")
            from peft import PeftModel
            self.model = PeftModel.from_pretrained(self.model, lora_path)

        # Enable fast inference
        FastLanguageModel.for_inference(self.model)

        # Load SNAC audio codec
        print("Loading SNAC audio codec...")
        self.snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
        self.snac_model = self.snac_model.to("cpu")  # Keep on CPU to save VRAM

        # Special tokens
        self.start_of_human = 128259
        self.end_of_text = 128009
        self.end_of_human = 128260
        self.start_of_ai = 128261
        self.start_of_speech = 128257
        self.end_of_speech = 128258
        self.pad_token = 128263

        print("Model ready for inference!")

    def _prepare_inputs(
        self,
        prompts: List[str],
        voice: Optional[str] = None
    ) -> tuple:
        """
        Prepare input tensors for the model.

        Args:
            prompts: List of text prompts to convert to speech
            voice: Optional voice/speaker name for multi-speaker models

        Returns:
            Tuple of (input_ids, attention_mask)
        """
        # Add voice prefix if specified
        prompts_ = [(f"{voice}: " + p) if voice else p for p in prompts]

        # Tokenize all prompts
        all_input_ids = []
        for prompt in prompts_:
            input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
            all_input_ids.append(input_ids)

        # Add special tokens: SOH SOT Text EOT EOH
        start_token = torch.tensor([[self.start_of_human]], dtype=torch.int64)
        end_tokens = torch.tensor([[self.end_of_text, self.end_of_human]], dtype=torch.int64)

        all_modified_input_ids = []
        for input_ids in all_input_ids:
            modified_input_ids = torch.cat([start_token, input_ids, end_tokens], dim=1)
            all_modified_input_ids.append(modified_input_ids)

        # Pad all sequences to the same length
        max_length = max([ids.shape[1] for ids in all_modified_input_ids])
        all_padded_tensors = []
        all_attention_masks = []

        for modified_input_ids in all_modified_input_ids:
            padding = max_length - modified_input_ids.shape[1]
            padded_tensor = torch.cat(
                [torch.full((1, padding), self.pad_token, dtype=torch.int64), modified_input_ids],
                dim=1
            )
            attention_mask = torch.cat(
                [torch.zeros((1, padding), dtype=torch.int64),
                 torch.ones((1, modified_input_ids.shape[1]), dtype=torch.int64)],
                dim=1
            )
            all_padded_tensors.append(padded_tensor)
            all_attention_masks.append(attention_mask)

        input_ids = torch.cat(all_padded_tensors, dim=0).to(self.device)
        attention_mask = torch.cat(all_attention_masks, dim=0).to(self.device)

        return input_ids, attention_mask

    def _decode_audio(self, generated_ids: torch.Tensor) -> List[torch.Tensor]:
        """
        Decode generated token IDs into audio waveforms.

        Args:
            generated_ids: Tensor of generated token IDs

        Returns:
            List of audio waveform tensors
        """
        # Find start of speech token
        token_to_find = self.start_of_speech
        token_to_remove = self.end_of_speech

        token_indices = (generated_ids == token_to_find).nonzero(as_tuple=True)

        # Crop to speech tokens only
        if len(token_indices[1]) > 0:
            last_occurrence_idx = token_indices[1][-1].item()
            cropped_tensor = generated_ids[:, last_occurrence_idx+1:]
        else:
            cropped_tensor = generated_ids

        # Remove end of speech tokens
        processed_rows = []
        for row in cropped_tensor:
            masked_row = row[row != token_to_remove]
            processed_rows.append(masked_row)

        # Prepare code lists for SNAC decoder
        code_lists = []
        for row in processed_rows:
            row_length = row.size(0)
            new_length = (row_length // 7) * 7  # Each frame has 7 tokens
            trimmed_row = row[:new_length]
            trimmed_row = [t.item() - 128266 for t in trimmed_row]  # Offset for audio tokens
            code_lists.append(trimmed_row)

        # Decode using SNAC
        audio_samples = []
        for code_list in code_lists:
            audio = self._redistribute_codes(code_list)
            audio_samples.append(audio)

        return audio_samples

    def _redistribute_codes(self, code_list: List[int]) -> torch.Tensor:
        """
        Redistribute flattened codes back into SNAC's 3-layer format.

        Terminates early if invalid codes are detected (out of range 0-4095)
        to prevent machine noise at the end of audio.

        Args:
            code_list: Flattened list of audio codes

        Returns:
            Audio waveform tensor
        """
        layer_1 = []
        layer_2 = []
        layer_3 = []

        # SNAC codebook size is 4096 per layer (valid range: 0-4095)
        max_code_value = 4095

        for i in range((len(code_list) + 1) // 7):
            # Extract codes with offsets
            c0 = code_list[7*i]
            c1 = code_list[7*i+1] - 4096
            c2 = code_list[7*i+2] - (2*4096)
            c3 = code_list[7*i+3] - (3*4096)
            c4 = code_list[7*i+4] - (4*4096)
            c5 = code_list[7*i+5] - (5*4096)
            c6 = code_list[7*i+6] - (6*4096)

            # Check if any code is out of valid range
            # If so, terminate audio generation to avoid machine noise
            if (c0 < 0 or c0 > max_code_value or
                c1 < 0 or c1 > max_code_value or
                c2 < 0 or c2 > max_code_value or
                c3 < 0 or c3 > max_code_value or
                c4 < 0 or c4 > max_code_value or
                c5 < 0 or c5 > max_code_value or
                c6 < 0 or c6 > max_code_value):
                print(f"Invalid audio code detected at frame {i}, terminating audio generation")
                break

            layer_1.append(c0)
            layer_2.append(c1)
            layer_3.append(c2)
            layer_3.append(c3)
            layer_2.append(c4)
            layer_3.append(c5)
            layer_3.append(c6)

        # Return empty/silent audio if no valid codes were found
        if not layer_1:
            print("Warning: No valid audio codes found, returning silence")
            return torch.zeros(1, 1, 1000)  # Small silent audio

        codes = [
            torch.tensor(layer_1, dtype=torch.long).unsqueeze(0),
            torch.tensor(layer_2, dtype=torch.long).unsqueeze(0),
            torch.tensor(layer_3, dtype=torch.long).unsqueeze(0)
        ]

        audio_hat = self.snac_model.decode(codes)
        return audio_hat

    def generate(
        self,
        prompts: List[str],
        output_dir: str = "generated_audio",
        voice: Optional[str] = None,
        max_new_tokens: int = 1200,
        temperature: float = 0.6,
        top_p: float = 0.95,
        repetition_penalty: float = 1.1,
        do_sample: bool = True
    ) -> List[str]:
        """
        Generate speech from text prompts.

        Orpheus supports emotion/expression tags in your prompts:
        - <laugh>, <giggles>, <chuckle> - Laughter variations
        - <sigh>, <gasp> - Emotional expressions
        - <yawn>, <cough>, <sniffle>, <groan> - Physical sounds

        These tags are treated as special tokens that trigger corresponding
        audio patterns learned during training. The Elise dataset contains
        hundreds of examples with these tags.

        Example prompts:
            "Hey there <giggles> welcome to my channel!"
            "I missed you <laugh> so much!"
            "That's so beautiful <sigh> it brings back memories."

        Args:
            prompts: List of text prompts to convert to speech (can include tags)
            output_dir: Directory to save generated audio files
            voice: Optional voice/speaker name for multi-speaker models
            max_new_tokens: Maximum number of tokens to generate
            temperature: Sampling temperature (higher = more random)
            top_p: Nucleus sampling threshold
            repetition_penalty: Penalty for repeating tokens
            do_sample: Whether to use sampling (vs greedy decoding)

        Returns:
            List of paths to generated audio files
        """
        print(f"Generating speech for {len(prompts)} prompt(s)...")

        # Prepare inputs
        input_ids, attention_mask = self._prepare_inputs(prompts, voice)

        # Generate tokens
        print("Generating tokens...")
        with torch.inference_mode():
            generated_ids = self.model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_new_tokens=max_new_tokens,
                do_sample=do_sample,
                temperature=temperature,
                top_p=top_p,
                repetition_penalty=repetition_penalty,
                num_return_sequences=1,
                eos_token_id=self.end_of_speech,
                use_cache=True
            )

        # Decode to audio
        print("Decoding audio...")
        audio_samples = self._decode_audio(generated_ids)

        # Save to files
        os.makedirs(output_dir, exist_ok=True)
        output_paths = []

        for idx, audio_sample in enumerate(audio_samples):
            # Convert tensor to numpy (detach first to avoid gradient tracking)
            audio_numpy = audio_sample.squeeze().detach().cpu().numpy()

            # Save as WAV file
            output_path = os.path.join(output_dir, f"output_{idx}.wav")
            sf.write(output_path, audio_numpy, self.sample_rate)
            output_paths.append(output_path)
            print(f"✓ Saved: {output_path}")

        return output_paths


def main():
    """Main function for standalone execution."""

    # Example prompts with emotion tags
    # Orpheus supports special tags like <laugh>, <giggles>, <chuckle>, <sigh>, 
    # <cough>, <sniffle>, <groan>, <yawn>, <gasp>, etc.
    # These tags are enclosed in angle brackets and will be treated as special tokens
    # that the model learned during training to generate corresponding audio patterns.
    prompts = [
        "Hey there my name is Elise, <giggles> and I'm a speech generation model that can sound like a person.",
        "I missed you <laugh> so much! It's been way too long.",
        "This is absolutely amazing <gasp> I can't believe it worked!",
        "I'm so tired <yawn> after working all day on this project.",
        "That's really touching <sigh> it reminds me of home.",
    ]

    # Initialize inference engine
    tts = OrpheusInference(
        model_path="unsloth/orpheus-3b-0.1-ft",
        lora_path="lora_model" if os.path.exists("lora_model") else None,
        load_in_4bit=False
    )

    # Generate speech
    output_files = tts.generate(
        prompts=prompts,
        output_dir="generated_audio",
        temperature=0.6,
        top_p=0.95,
        max_new_tokens=1200
    )

    print(f"\n✅ Successfully generated {len(output_files)} audio file(s)!")
    print(f"📁 Output directory: generated_audio/")

    # For multi-speaker models, you can specify a voice:
    # output_files = tts.generate(prompts=prompts, voice="speaker_name")


if __name__ == "__main__":
    main()

orpheus_sft_unsloth.py

# -*- coding: utf-8 -*-
"""Orpheus_(3B)-TTS.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Orpheus_(3B)-TTS.ipynb

To run this, press "*Runtime*" and press "*Run all*" on a **free** Tesla T4 Google Colab instance!
<div class="align-center">
<a href="https://unsloth.ai/"><img src="https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png" width="115"></a>
<a href="https://discord.gg/unsloth"><img src="https://github.com/unslothai/unsloth/raw/main/images/Discord button.png" width="145"></a>
<a href="https://docs.unsloth.ai/"><img src="https://github.com/unslothai/unsloth/blob/main/images/documentation%20green%20button.png?raw=true" width="125"></a></a> Join Discord if you need help + ⭐ <i>Star us on <a href="https://github.com/unslothai/unsloth">Github</a> </i> ⭐
</div>

To install Unsloth on your own computer, follow the installation instructions on our Github page [here](https://docs.unsloth.ai/get-started/installing-+-updating).

You will learn how to do [data prep](#Data), how to [train](#Train), how to [run the model](#Inference), & [how to save it](#Save)

### News

Unsloth's [Docker image](https://hub.docker.com/r/unsloth/unsloth) is here! Start training with no setup & environment issues. [Read our Guide](https://docs.unsloth.ai/new/how-to-train-llms-with-unsloth-and-docker).

[gpt-oss RL](https://docs.unsloth.ai/new/gpt-oss-reinforcement-learning) is now supported with the fastest inference & lowest VRAM. Try our [new notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/gpt-oss-(20B)-GRPO.ipynb) which creates kernels!

Introducing [Vision](https://docs.unsloth.ai/new/vision-reinforcement-learning-vlm-rl) and [Standby](https://docs.unsloth.ai/basics/memory-efficient-rl) for RL! Train Qwen, Gemma etc. VLMs with GSPO - even faster with less VRAM.

Unsloth now supports Text-to-Speech (TTS) models. Read our [guide here](https://docs.unsloth.ai/basics/text-to-speech-tts-fine-tuning).

Visit our docs for all our [model uploads](https://docs.unsloth.ai/get-started/all-our-models) and [notebooks](https://docs.unsloth.ai/get-started/unsloth-notebooks).

### Installation
"""

# Commented out IPython magic to ensure Python compatibility.
# %%capture
# import os, re
# if "COLAB_" not in "".join(os.environ.keys()):
#     !pip install unsloth
# else:
#     # Do this only in Colab notebooks! Otherwise use pip install unsloth
#     import torch; v = re.match(r"[0-9\.]{3,}", str(torch.__version__)).group(0)
#     xformers = "xformers==" + ("0.0.32.post2" if v == "2.8.0" else "0.0.29.post3")
#     !pip install --no-deps bitsandbytes accelerate {xformers} peft trl triton cut_cross_entropy unsloth_zoo
#     !pip install sentencepiece protobuf "datasets>=3.4.1,<4.0.0" "huggingface_hub>=0.34.0" hf_transfer
#     !pip install --no-deps unsloth
# !pip install transformers==4.55.4
# !pip install --no-deps trl==0.22.2
# !pip install snac
# !pip install soundfile librosa

"""### Unsloth

`FastModel` supports loading nearly any model now! This includes Vision and Text models!

Thank you to [Etherl](https://huggingface.co/Etherll) for creating this notebook!
"""

from unsloth import FastLanguageModel
import torch

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/orpheus-3b-0.1-ft",
    max_seq_length= 2048, # Choose any for long context!
    dtype = None, # Select None for auto detection
    load_in_4bit = False, # Select True for 4bit which reduces memory usage
)

"""We now add LoRA adapters so we only need to update 1 to 10% of all parameters!"""

model = FastLanguageModel.get_peft_model(
    model,
    r = 64, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj",],
    lora_alpha = 64,
    lora_dropout = 0, # Supports any, but = 0 is optimized
    bias = "none",    # Supports any, but = "none" is optimized
    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
    use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
    random_state = 42,
    use_rslora = False,  # We support rank stabilized LoRA
    loftq_config = None, # And LoftQ
)

"""<a name="Data"></a>
### Data Prep  

We will use the `MrDragonFox/Elise`, which is designed for training TTS models. Ensure that your dataset follows the required format: **text, audio** for single-speaker models or **source, text, audio** for multi-speaker models. You can modify this section to accommodate your own dataset, but maintaining the correct structure is essential for optimal training.
"""

from datasets import load_dataset
dataset = load_dataset("MrDragonFox/Elise", split = "train")

#@title Tokenization Function

import locale
import torchaudio.transforms as T
import os
import torch
from snac import SNAC
locale.getpreferredencoding = lambda: "UTF-8"
ds_sample_rate = dataset[0]["audio"]["sampling_rate"]

snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
snac_model = snac_model.to("cuda")
def tokenise_audio(waveform):
  waveform = torch.from_numpy(waveform).unsqueeze(0)
  waveform = waveform.to(dtype=torch.float32)
  resample_transform = T.Resample(orig_freq=ds_sample_rate, new_freq=24000)
  waveform = resample_transform(waveform)

  waveform = waveform.unsqueeze(0).to("cuda")

  #generate the codes from snac
  with torch.inference_mode():
    codes = snac_model.encode(waveform)

  all_codes = []
  for i in range(codes[0].shape[1]):
    all_codes.append(codes[0][0][i].item()+128266)
    all_codes.append(codes[1][0][2*i].item()+128266+4096)
    all_codes.append(codes[2][0][4*i].item()+128266+(2*4096))
    all_codes.append(codes[2][0][(4*i)+1].item()+128266+(3*4096))
    all_codes.append(codes[1][0][(2*i)+1].item()+128266+(4*4096))
    all_codes.append(codes[2][0][(4*i)+2].item()+128266+(5*4096))
    all_codes.append(codes[2][0][(4*i)+3].item()+128266+(6*4096))


  return all_codes

def add_codes(example):
    # Always initialize codes_list to None
    codes_list = None

    try:
        answer_audio = example.get("audio")
        # If there's a valid audio array, tokenise it
        if answer_audio and "array" in answer_audio:
            audio_array = answer_audio["array"]
            codes_list = tokenise_audio(audio_array)
    except Exception as e:
        print(f"Skipping row due to error: {e}")
        # Keep codes_list as None if we fail
    example["codes_list"] = codes_list

    return example

dataset = dataset.map(add_codes, remove_columns=["audio"])

tokeniser_length = 128256
start_of_text = 128000
end_of_text = 128009

start_of_speech = tokeniser_length + 1
end_of_speech = tokeniser_length + 2

start_of_human = tokeniser_length + 3
end_of_human = tokeniser_length + 4

start_of_ai = tokeniser_length + 5
end_of_ai =  tokeniser_length + 6
pad_token = tokeniser_length + 7

audio_tokens_start = tokeniser_length + 10

dataset = dataset.filter(lambda x: x["codes_list"] is not None)
dataset = dataset.filter(lambda x: len(x["codes_list"]) > 0)

def remove_duplicate_frames(example):
    vals = example["codes_list"]
    if len(vals) % 7 != 0:
        raise ValueError("Input list length must be divisible by 7")

    result = vals[:7]

    removed_frames = 0

    for i in range(7, len(vals), 7):
        current_first = vals[i]
        previous_first = result[-7]

        if current_first != previous_first:
            result.extend(vals[i:i+7])
        else:
            removed_frames += 1

    example["codes_list"] = result

    return example

dataset = dataset.map(remove_duplicate_frames)

tok_info = '''*** HERE you can modify the text prompt
If you are training a multi-speaker model (e.g., canopylabs/orpheus-3b-0.1-ft),
ensure that the dataset includes a "source" field and format the input accordingly:
- Single-speaker: f"{example['text']}"
- Multi-speaker: f"{example['source']}: {example['text']}"
'''
print(tok_info)

def create_input_ids(example):
    # Determine whether to include the source field
    text_prompt = f"{example['source']}: {example['text']}" if "source" in example else example["text"]

    text_ids = tokenizer.encode(text_prompt, add_special_tokens=True)
    text_ids.append(end_of_text)

    example["text_tokens"] = text_ids
    input_ids = (
        [start_of_human]
        + example["text_tokens"]
        + [end_of_human]
        + [start_of_ai]
        + [start_of_speech]
        + example["codes_list"]
        + [end_of_speech]
        + [end_of_ai]
    )
    example["input_ids"] = input_ids
    example["labels"] = input_ids
    example["attention_mask"] = [1] * len(input_ids)

    return example


dataset = dataset.map(create_input_ids, remove_columns=["text", "codes_list"])
columns_to_keep = ["input_ids", "labels", "attention_mask"]
columns_to_remove = [col for col in dataset.column_names if col not in columns_to_keep]

dataset = dataset.remove_columns(columns_to_remove)

"""<a name="Train"></a>
### Train the model
Now let's use Huggingface  `Trainer`! More docs here: [Transformers docs](https://huggingface.co/docs/transformers/main_classes/trainer). We do 60 steps to speed things up, but you can set `num_train_epochs=1` for a full run, and turn off `max_steps=None`.

**Note:** Using a per_device_train_batch_size >1 may lead to errors if multi-GPU setup to avoid issues, ensure CUDA_VISIBLE_DEVICES is set to a single GPU (e.g., CUDA_VISIBLE_DEVICES=0).
"""

from transformers import TrainingArguments,Trainer,DataCollatorForSeq2Seq
trainer = Trainer(
    model = model,
    train_dataset = dataset,
    args = TrainingArguments(
        per_device_train_batch_size = 1,
        gradient_accumulation_steps = 4,
        warmup_steps = 5,
        num_train_epochs = 1, # Set this for 1 full training run.
        learning_rate = 2e-4,
        logging_steps = 1,
        optim = "adamw_8bit",
        weight_decay = 0.01,
        lr_scheduler_type = "linear",
        seed = 42,
        output_dir = "outputs",
        report_to = "none", # Use this for WandB etc
    ),
)

# @title Show current memory stats
gpu_stats = torch.cuda.get_device_properties(0)
start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.")
print(f"{start_gpu_memory} GB of memory reserved.")

trainer_stats = trainer.train()

# @title Show final memory and time stats
used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
used_memory_for_lora = round(used_memory - start_gpu_memory, 3)
used_percentage = round(used_memory / max_memory * 100, 3)
lora_percentage = round(used_memory_for_lora / max_memory * 100, 3)
print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.")
print(
    f"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training."
)
print(f"Peak reserved memory = {used_memory} GB.")
print(f"Peak reserved memory for training = {used_memory_for_lora} GB.")
print(f"Peak reserved memory % of max memory = {used_percentage} %.")
print(f"Peak reserved memory for training % of max memory = {lora_percentage} %.")


print("Saving model...")
"""<a name="Save"></a>
### Saving, loading finetuned models
To save the final model as LoRA adapters, either use Huggingface's `push_to_hub` for an online save or `save_pretrained` for a local save.

**[NOTE]** This ONLY saves the LoRA adapters, and not the full model. To save to 16bit or GGUF, scroll down!
"""

model.save_pretrained("lora_model")  # Local saving
tokenizer.save_pretrained("lora_model")
# model.push_to_hub("your_name/lora_model", token = "...") # Online saving
# tokenizer.push_to_hub("your_name/lora_model", token = "...") # Online saving

"""### Saving to float16

We also support saving to `float16` directly. Select `merged_16bit` for float16 or `merged_4bit` for int4. We also allow `lora` adapters as a fallback. Use `push_to_hub_merged` to upload to your Hugging Face account! You can go to https://huggingface.co/settings/tokens for your personal tokens.
"""

# Merge to 16bit
if False: model.save_pretrained_merged("model", tokenizer, save_method = "merged_16bit",)
if False: model.push_to_hub_merged("hf/model", tokenizer, save_method = "merged_16bit", token = "")

# Merge to 4bit
if False: model.save_pretrained_merged("model", tokenizer, save_method = "merged_4bit",)
if False: model.push_to_hub_merged("hf/model", tokenizer, save_method = "merged_4bit", token = "")

# Just LoRA adapters
if False:
    model.save_pretrained("model")
    tokenizer.save_pretrained("model")
if False:
    model.push_to_hub("hf/model", token = "")
    tokenizer.push_to_hub("hf/model", token = "")

print("Inference...")
"""<a name="Inference"></a>
### Inference
Let's run the model! You can change the prompts


"""

prompts = [
    "Hey there my name is Elise, <giggles> and I'm a speech generation model that can sound like a person.",
    "I missed you <laugh> so much! It's been way too long.",
    "This is absolutely amazing <gasp> I can't believe it worked!",
]

# Orpheus supports emotion tags: <laugh>, <giggles>, <chuckle>, <sigh>,
# <gasp>, <yawn>, <cough>, <sniffle>, <groan>, etc.
chosen_voice = None # None for single-speaker

#@title Run Inference


FastLanguageModel.for_inference(model) # Enable native 2x faster inference

# Moving snac_model cuda to cpu
snac_model.to("cpu")

prompts_ = [(f"{chosen_voice}: " + p) if chosen_voice else p for p in prompts]

all_input_ids = []

for prompt in prompts_:
  input_ids = tokenizer(prompt, return_tensors="pt").input_ids
  all_input_ids.append(input_ids)

start_token = torch.tensor([[ 128259]], dtype=torch.int64) # Start of human
end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64) # End of text, End of human

all_modified_input_ids = []
for input_ids in all_input_ids:
  modified_input_ids = torch.cat([start_token, input_ids, end_tokens], dim=1) # SOH SOT Text EOT EOH
  all_modified_input_ids.append(modified_input_ids)

all_padded_tensors = []
all_attention_masks = []
max_length = max([modified_input_ids.shape[1] for modified_input_ids in all_modified_input_ids])
for modified_input_ids in all_modified_input_ids:
  padding = max_length - modified_input_ids.shape[1]
  padded_tensor = torch.cat([torch.full((1, padding), 128263, dtype=torch.int64), modified_input_ids], dim=1)
  attention_mask = torch.cat([torch.zeros((1, padding), dtype=torch.int64), torch.ones((1, modified_input_ids.shape[1]), dtype=torch.int64)], dim=1)
  all_padded_tensors.append(padded_tensor)
  all_attention_masks.append(attention_mask)

all_padded_tensors = torch.cat(all_padded_tensors, dim=0)
all_attention_masks = torch.cat(all_attention_masks, dim=0)

input_ids = all_padded_tensors.to("cuda")
attention_mask = all_attention_masks.to("cuda")
generated_ids = model.generate(
      input_ids=input_ids,
      attention_mask=attention_mask,
      max_new_tokens=1200,
      do_sample=True,
      temperature=0.6,
      top_p=0.95,
      repetition_penalty=1.1,
      num_return_sequences=1,
      eos_token_id=128258,
     use_cache = True
  )
token_to_find = 128257
token_to_remove = 128258

token_indices = (generated_ids == token_to_find).nonzero(as_tuple=True)

if len(token_indices[1]) > 0:
    last_occurrence_idx = token_indices[1][-1].item()
    cropped_tensor = generated_ids[:, last_occurrence_idx+1:]
else:
    cropped_tensor = generated_ids

mask = cropped_tensor != token_to_remove

processed_rows = []

for row in cropped_tensor:
    masked_row = row[row != token_to_remove]
    processed_rows.append(masked_row)

code_lists = []

for row in processed_rows:
    row_length = row.size(0)
    new_length = (row_length // 7) * 7
    trimmed_row = row[:new_length]
    trimmed_row = [t - 128266 for t in trimmed_row]
    code_lists.append(trimmed_row)


def redistribute_codes(code_list):
  layer_1 = []
  layer_2 = []
  layer_3 = []
  # SNAC codebook size is 4096 per layer (valid range: 0-4095)
  max_code_value = 4095
  for i in range((len(code_list)+1)//7):
    # Extract codes with offsets
    c0 = code_list[7*i]
    c1 = code_list[7*i+1]-4096
    c2 = code_list[7*i+2]-(2*4096)
    c3 = code_list[7*i+3]-(3*4096)
    c4 = code_list[7*i+4]-(4*4096)
    c5 = code_list[7*i+5]-(5*4096)
    c6 = code_list[7*i+6]-(6*4096)

    # Check if any code is out of valid range
    # If so, terminate audio generation to avoid machine noise
    if (c0 < 0 or c0 > max_code_value or
        c1 < 0 or c1 > max_code_value or
        c2 < 0 or c2 > max_code_value or
        c3 < 0 or c3 > max_code_value or
        c4 < 0 or c4 > max_code_value or
        c5 < 0 or c5 > max_code_value or
        c6 < 0 or c6 > max_code_value):
      print(f"Invalid audio code detected at frame {i}, terminating audio generation")
      break

    layer_1.append(c0)
    layer_2.append(c1)
    layer_3.append(c2)
    layer_3.append(c3)
    layer_2.append(c4)
    layer_3.append(c5)
    layer_3.append(c6)

  # Return empty/silent audio if no valid codes were found
  if not layer_1:
    print("Warning: No valid audio codes found, returning silence")
    return torch.zeros(1, 1, 1000)  # Small silent audio

  codes = [torch.tensor(layer_1, dtype=torch.long).unsqueeze(0),
         torch.tensor(layer_2, dtype=torch.long).unsqueeze(0),
         torch.tensor(layer_3, dtype=torch.long).unsqueeze(0)]

  # codes = [c.to("cuda") for c in codes]
  audio_hat = snac_model.decode(codes)
  return audio_hat

my_samples = []
for code_list in code_lists:
  samples = redistribute_codes(code_list)
  my_samples.append(samples)

# Save generated audio samples to WAV files
import soundfile as sf
import os

output_dir = "generated_audio"
os.makedirs(output_dir, exist_ok=True)

for idx, audio_sample in enumerate(my_samples):
  # Convert tensor to numpy array and squeeze to remove batch dimension
  # Detach from computation graph to avoid gradient tracking error
  audio_numpy = audio_sample.squeeze().detach().cpu().numpy()
  # Save to WAV file with 24kHz sample rate (matching SNAC model)
  output_path = os.path.join(output_dir, f"output_{idx}.wav")
  sf.write(output_path, audio_numpy, 24000)
  print(f"Saved audio to {output_path}")

# Clean up to save RAM
del my_samples,samples

"""And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!

Some other links:
1. Train your own reasoning model - Llama GRPO notebook [Free Colab](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.1_(8B)-GRPO.ipynb)
2. Saving finetunes to Ollama. [Free notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3_(8B)-Ollama.ipynb)
3. Llama 3.2 Vision finetuning - Radiography use case. [Free Colab](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.2_(11B)-Vision.ipynb)
6. See notebooks for DPO, ORPO, Continued pretraining, conversational finetuning and more on our [documentation](https://docs.unsloth.ai/get-started/unsloth-notebooks)!

<div class="align-center">
  <a href="https://unsloth.ai"><img src="https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png" width="115"></a>
  <a href="https://discord.gg/unsloth"><img src="https://github.com/unslothai/unsloth/raw/main/images/Discord.png" width="145"></a>
  <a href="https://docs.unsloth.ai/"><img src="https://github.com/unslothai/unsloth/blob/main/images/documentation%20green%20button.png?raw=true" width="125"></a>

  Join Discord if you need help + ⭐️ <i>Star us on <a href="https://github.com/unslothai/unsloth">Github</a> </i> ⭐️
</div>

"""