Files
tts/notebook/vyvotts-lfm2-train.ipynb

584 KiB
Raw Permalink Blame History

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.

You will learn how to do data prep, how to train, how to run the model, & how to save it

News

Unsloth now supports Text-to-Speech (TTS) models. Read our guide here.

Read our Gemma 3N Guide and check out our new Dynamic 2.0 quants which outperforms other quantization methods!

Visit our docs for all our model uploads and notebooks.

Installation

In [ ]:
%%capture
import os
if "COLAB_" not in "".join(os.environ.keys()):
    !pip install unsloth
else:
    # Do this only in Colab notebooks! Otherwise use pip install unsloth
    !pip install --no-deps bitsandbytes accelerate xformers==0.0.29.post3 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 snac

Unsloth

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

Thank you to Etherl for creating this notebook!

In [ ]:
from unsloth import FastLanguageModel
import torch

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "Vyvo/VyvoTTS-LFM2-PT,
    max_seq_length= 8192, # Choose any for long context!
    dtype = None, # Select None for auto detection
    load_in_4bit = False, # Select True for 4bit which reduces memory usage
    load_in_8bit = True, # Select True for 8bit which reduces memory usage
    full_finetuning = True
    # token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
)
==((====))==  Unsloth 2025.3.19: Fast Llama patching. Transformers: 4.50.0.
   \\   /|    Tesla T4. Num GPUs = 1. Max memory: 14.741 GB. Platform: Linux.
O^O/ \_/ \    Torch: 2.6.0+cu124. CUDA: 7.5. CUDA Toolkit: 12.4. Triton: 3.2.0
\        /    Bfloat16 = FALSE. FA [Xformers = 0.0.29.post3. FA2 = False]
 "-____-"     Free license: http://github.com/unslothai/unsloth
Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!

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

In [ ]:
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 = 3407,
    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.

In [ ]:
from datasets import load_dataset
dataset = load_dataset("MrDragonFox/Elise", split = "train")
In [ ]:
tokeniser_length = 64400
start_of_text = 1
end_of_text = 7

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
In [ ]:
#@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()+audio_tokens_start)
    all_codes.append(codes[1][0][2*i].item()+audio_tokens_start+4096)
    all_codes.append(codes[2][0][4*i].item()+audio_tokens_start+(2*4096))
    all_codes.append(codes[2][0][(4*i)+1].item()+audio_tokens_start+(3*4096))
    all_codes.append(codes[1][0][(2*i)+1].item()+audio_tokens_start+(4*4096))
    all_codes.append(codes[2][0][(4*i)+2].item()+audio_tokens_start+(5*4096))
    all_codes.append(codes[2][0][(4*i)+3].item()+audio_tokens_start+(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"])
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)
Map:   0%|          | 0/1195 [00:00<?, ? examples/s]
Filter:   0%|          | 0/1195 [00:00<?, ? examples/s]
Filter:   0%|          | 0/1195 [00:00<?, ? examples/s]
Map:   0%|          | 0/1195 [00:00<?, ? examples/s]
*** 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']}"

Map:   0%|          | 0/1195 [00:00<?, ? examples/s]

Train the model

Now let's use Huggingface Trainer! More docs here: Transformers docs. 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).

In [ ]:
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.
        max_steps = 300,
        learning_rate = 2e-4,
        logging_steps = 1,
        optim = "adamw_8bit",
        weight_decay = 0.01,
        lr_scheduler_type = "linear",
        seed = 3407,
        output_dir = "outputs",
        report_to = "none", # Use this for WandB etc
    ),
)
In [ ]:
# @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.")
GPU = Tesla T4. Max memory = 14.741 GB.
5.713 GB of memory reserved.
In [ ]:
trainer_stats = trainer.train()
==((====))==  Unsloth - 2x faster free finetuning | Num GPUs used = 1
   \\   /|    Num examples = 1,195 | Num Epochs = 1 | Total steps = 298
O^O/ \_/ \    Batch size per device = 1 | Gradient accumulation steps = 4
\        /    Data Parallel GPUs = 1 | Total batch size (1 x 4 x 1) = 4
 "-____-"     Trainable parameters = 97,255,424/3,000,000,000 (3.24% trained)
[298/298 18:25, Epoch 0/1]
Step Training Loss
1 5.113800
2 4.870600
3 4.988100
4 4.852700
5 5.009200
6 4.836700
7 4.804600
8 4.941900
9 4.988000
10 4.666500
11 4.952900
12 4.774700
13 4.728000
14 4.658200
15 4.647300
16 4.809000
17 4.885800
18 4.775500
19 4.859700
20 4.533900
21 4.821600
22 4.759600
23 4.844500
24 4.372000
25 4.513600
26 4.670800
27 4.614600
28 4.669200
29 4.533900
30 4.502700
31 4.977600
32 4.601000
33 4.522600
34 4.557400
35 4.307500
36 4.645900
37 4.572300
38 4.645600
39 4.772800
40 4.543400
41 4.594200
42 4.548800
43 4.598800
44 4.458900
45 4.618300
46 4.632200
47 4.649500
48 4.599900
49 4.477100
50 4.606300
51 4.484800
52 4.536300
53 4.523300
54 4.141800
55 4.780900
56 4.400400
57 4.542100
58 4.492900
59 4.542500
60 4.578300
61 4.715400
62 4.843400
63 4.396200
64 4.419400
65 4.525100
66 4.639600
67 4.532200
68 4.650900
69 4.680500
70 4.602700
71 4.648100
72 4.605000
73 4.592900
74 4.576600
75 4.662500
76 4.342400
77 4.542600
78 4.494200
79 4.563300
80 4.535600
81 4.553600
82 4.411700
83 4.431900
84 4.280100
85 4.416700
86 4.338000
87 4.587000
88 4.395700
89 4.390400
90 4.487500
91 4.532900
92 4.418700
93 4.740300
94 4.589200
95 4.613500
96 4.452400
97 4.535600
98 4.538300
99 4.432100
100 4.395900
101 4.468100
102 4.540800
103 4.499400
104 4.437100
105 4.391600
106 4.289700
107 4.376700
108 4.474300
109 4.708600
110 4.627000
111 4.610700
112 4.433900
113 4.413500
114 4.429300
115 4.553900
116 4.429100
117 4.675700
118 4.634600
119 4.388800
120 4.530900
121 4.438900
122 4.509700
123 4.470200
124 4.594200
125 4.313500
126 4.557300
127 4.309400
128 4.480300
129 4.426600
130 4.434600
131 4.414400
132 4.425600
133 4.581100
134 4.601200
135 4.452500
136 4.534100
137 4.484200
138 4.325700
139 4.503200
140 4.335100
141 4.387700
142 4.489700
143 4.381000
144 4.465700
145 4.483200
146 4.489400
147 4.519100
148 4.252900
149 4.384400
150 4.679200
151 4.448900
152 4.409900
153 4.659700
154 4.422400
155 4.532200
156 4.400100
157 4.371400
158 4.636900
159 4.477200
160 4.559400
161 4.385800
162 4.253300
163 4.334500
164 4.688000
165 4.456600
166 4.289500
167 4.474500
168 4.051200
169 4.454400
170 4.584100
171 4.505800
172 4.562000
173 4.492300
174 4.441000
175 4.511700
176 4.533600
177 4.405700
178 4.570700
179 4.729400
180 4.335000
181 4.397600
182 4.405200
183 4.362600
184 4.209000
185 4.401600
186 4.412000
187 4.305400
188 4.653600
189 4.484300
190 4.584500
191 4.372600
192 4.517400
193 4.398100
194 4.561600
195 4.544700
196 4.382200
197 4.309600
198 4.248300
199 4.466300
200 4.320800
201 4.551600
202 4.588900
203 4.533600
204 4.355500
205 4.417700
206 4.363200
207 4.453800
208 4.504600
209 4.298700
210 4.462900
211 4.541100
212 4.291200
213 4.663200
214 4.561500
215 4.618500
216 4.662000
217 4.355400
218 4.300400
219 4.352100
220 4.505100
221 4.445300
222 4.462000
223 4.486900
224 4.427700
225 4.528700
226 4.468100
227 4.582500
228 4.483600
229 4.317800
230 4.516200
231 4.276400
232 4.302900
233 4.590200
234 4.379800
235 4.482000
236 4.435500
237 4.404200
238 4.349000
239 4.456400
240 4.512600
241 4.587600
242 4.411700
243 4.472200
244 4.587900
245 4.475000
246 4.353400
247 4.577700
248 4.489700
249 4.452200
250 4.472400
251 4.617300
252 4.456900
253 4.738200
254 4.488000
255 4.476700
256 4.359100
257 4.513100
258 4.549900
259 4.632100
260 4.497100
261 4.330600
262 4.615500
263 4.272100
264 4.363700
265 4.373800
266 4.476800
267 4.191800
268 4.662300
269 4.343100
270 4.544400
271 4.401500
272 4.190000
273 4.405000
274 4.489900
275 4.571100
276 4.515900
277 4.634800
278 4.361800
279 4.403200
280 4.456400
281 4.603500
282 4.252100
283 4.485200
284 4.568500
285 4.444000
286 4.345700
287 4.234200
288 4.370200
289 4.528300
290 4.368400
291 4.622700
292 4.246400
293 4.399400
294 4.466000
295 4.334400
296 4.387500
297 4.348200
298 4.498000

In [ ]:
# @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} %.")

Inference

Let's run the model! You can change the prompts

In [ ]:
prompts = [
    "Hey there my name is Elise, <giggles> and I'm a speech generation model that can sound like a person.",
]

chosen_voice = None # None for single-speaker
In [ ]:
#@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([[ start_of_human]], dtype=torch.int64) # Start of human
end_tokens = torch.tensor([[end_of_text, end_of_human]], 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), 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)

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=end_of_speech,
     use_cache = True
  )
token_to_find = start_of_speech
token_to_remove = end_of_speech

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 - audio_tokens_start for t in trimmed_row]
    code_lists.append(trimmed_row)


def redistribute_codes(code_list):
  layer_1 = []
  layer_2 = []
  layer_3 = []
  for i in range((len(code_list)+1)//7):
    layer_1.append(code_list[7*i])
    layer_2.append(code_list[7*i+1]-4096)
    layer_3.append(code_list[7*i+2]-(2*4096))
    layer_3.append(code_list[7*i+3]-(3*4096))
    layer_2.append(code_list[7*i+4]-(4*4096))
    layer_3.append(code_list[7*i+5]-(5*4096))
    layer_3.append(code_list[7*i+6]-(6*4096))
  codes = [torch.tensor(layer_1).unsqueeze(0),
         torch.tensor(layer_2).unsqueeze(0),
         torch.tensor(layer_3).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)
from IPython.display import display, Audio
if len(prompts) != len(my_samples):
  raise Exception("Number of prompts and samples do not match")
else:
  for i in range(len(my_samples)):
    print(prompts[i])
    samples = my_samples[i]
    display(Audio(samples.detach().squeeze().to("cpu").numpy(), rate=24000))
# Clean up to save RAM
del my_samples,samples
Hey there my name is Elise, <giggles> and I'm a speech generation model that can sound like a person.
Your browser does not support the audio element.

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!

In [ ]:
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
('lora_model/tokenizer_config.json',
 'lora_model/special_tokens_map.json',
 'lora_model/tokenizer.json')

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.

In [ ]:
# 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 = "")
Unsloth: You have 1 CPUs. Using `safe_serialization` is 10x slower.
We shall switch to Pytorch saving, which might take 3 minutes and not 30 minutes.
To force `safe_serialization`, set it to `None` instead.
Unsloth: Kaggle/Colab has limited disk space. We need to delete the downloaded
model which will save 4-16GB of disk space, allowing you to save on Kaggle/Colab.
Unsloth: Will remove a cached repo with size 15.1G
Unsloth: Merging 4bit and LoRA weights to 16bit...
Unsloth: Will use up to 3.99 out of 12.67 RAM for saving.
Unsloth: Saving model... This might take 5 minutes ...
100%|██████████| 28/28 [00:01<00:00, 27.83it/s]
Unsloth: Saving tokenizer... Done.
Unsloth: Saving model/pytorch_model-00001-of-00002.bin...
Unsloth: Saving model/pytorch_model-00002-of-00002.bin...
Done.

And we're done! If you have any questions on Unsloth, we have a Discord 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
  2. Saving finetunes to Ollama. Free notebook
  3. Llama 3.2 Vision finetuning - Radiography use case. Free Colab
  4. See notebooks for DPO, ORPO, Continued pretraining, conversational finetuning and more on our documentation!

Join Discord if you need help + ⭐️ Star us on Github ⭐️