# -*- 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!
Join Discord if you need help + ⭐ Star us on Github
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 ) """ ### 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( "maxbsoft/mrdragonfox-elise", revision="2cc657c3f94a83df18fcd968b7531ca1a19c7f88", 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) """ ### 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...") """ ### 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...") """ ### Inference Let's run the model! You can change the prompts """ prompts = [ "Hey there my name is Elise, and I'm a speech generation model that can sound like a person.", "I missed you so much! It's been way too long.", "This is absolutely amazing I can't believe it worked!", ] # Orpheus supports emotion tags: , , , , # , , , , , 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)!
Join Discord if you need help + ⭐️ Star us on Github ⭐️
"""