Views
No views yet
unsloth and SNAC for speech generation:1from unsloth import FastLanguageModel
2import torch
3from snac import SNAC
4
5model, tokenizer = FastLanguageModel.from_pretrained(
6 model_name = "Vyvo/VyvoTTS-LFM2-Itto",
7 max_seq_length= 8192,
8 dtype = None,
9 load_in_4bit = False,
10)
11snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
12tokeniser_length = 64400
13start_of_text = 1
14end_of_text = 7
15
16start_of_speech = tokeniser_length + 1
17end_of_speech = tokeniser_length + 2
18start_of_human = tokeniser_length + 3
19end_of_human = tokeniser_length + 4
20pad_token = tokeniser_length + 7
21
22audio_tokens_start = tokeniser_length + 10
23prompts = ["Hey there my name is Elise, and I'm a speech generation model that can sound like a person."]
24chosen_voice = None
25
26FastLanguageModel.for_inference(model)
27snac_model.to("cpu")
28prompts_ = [(f"{chosen_voice}: " + p) if chosen_voice else p for p in prompts]
29
30all_input_ids = []
31for prompt in prompts_:
32 input_ids = tokenizer(prompt, return_tensors="pt").input_ids
33 all_input_ids.append(input_ids)
34
35start_token = torch.tensor([[start_of_human]], dtype=torch.int64)
36end_tokens = torch.tensor([[end_of_text, end_of_human]], dtype=torch.int64)
37
38all_modified_input_ids = []
39for input_ids in all_input_ids:
40 modified_input_ids = torch.cat([start_token, input_ids, end_tokens], dim=1)
41 all_modified_input_ids.append(modified_input_ids)
42
43all_padded_tensors, all_attention_masks = [], []
44max_length = max([m.shape[1] for m in all_modified_input_ids])
45for m in all_modified_input_ids:
46 padding = max_length - m.shape[1]
47 padded_tensor = torch.cat([torch.full((1, padding), pad_token, dtype=torch.int64), m], dim=1)
48 attention_mask = torch.cat([torch.zeros((1, padding), dtype=torch.int64), torch.ones((1, m.shape[1]), dtype=torch.int64)], dim=1)
49 all_padded_tensors.append(padded_tensor)
50 all_attention_masks.append(attention_mask)
51
52input_ids = torch.cat(all_padded_tensors, dim=0).to("cuda")
53attention_mask = torch.cat(all_attention_masks, dim=0).to("cuda")
54
55generated_ids = model.generate(
56 input_ids=input_ids,
57 attention_mask=attention_mask,
58 max_new_tokens=1200,
59 do_sample=True,
60 temperature=0.6,
61 top_p=0.95,
62 repetition_penalty=1.1,
63 num_return_sequences=1,
64 eos_token_id=end_of_speech,
65 use_cache=True
66)
67
68token_to_find = start_of_speech
69token_to_remove = end_of_speech
70token_indices = (generated_ids == token_to_find).nonzero(as_tuple=True)
71
72if len(token_indices[1]) > 0:
73 last_occurrence_idx = token_indices[1][-1].item()
74 cropped_tensor = generated_ids[:, last_occurrence_idx+1:]
75else:
76 cropped_tensor = generated_ids
77
78processed_rows = []
79for row in cropped_tensor:
80 masked_row = row[row != token_to_remove]
81 processed_rows.append(masked_row)
82
83code_lists = []
84for row in processed_rows:
85 row_length = row.size(0)
86 new_length = (row_length // 7) * 7
87 trimmed_row = row[:new_length]
88 trimmed_row = [t - audio_tokens_start for t in trimmed_row]
89 code_lists.append(trimmed_row)
90
91def redistribute_codes(code_list):
92 layer_1, layer_2, layer_3 = [], [], []
93 for i in range((len(code_list)+1)//7):
94 layer_1.append(code_list[7*i])
95 layer_2.append(code_list[7*i+1]-4096)
96 layer_3.append(code_list[7*i+2]-(2*4096))
97 layer_3.append(code_list[7*i+3]-(3*4096))
98 layer_2.append(code_list[7*i+4]-(4*4096))
99 layer_3.append(code_list[7*i+5]-(5*4096))
100 layer_3.append(code_list[7*i+6]-(6*4096))
101 codes = [
102 torch.tensor(layer_1).unsqueeze(0),
103 torch.tensor(layer_2).unsqueeze(0),
104 torch.tensor(layer_3).unsqueeze(0)
105 ]
106 audio_hat = snac_model.decode(codes)
107 return audio_hat
108
109my_samples = []
110for code_list in code_lists:
111 samples = redistribute_codes(code_list)
112 my_samples.append(samples)
113
114from IPython.display import display, Audio
115if len(prompts) != len(my_samples):
116 raise Exception("Number of prompts and samples do not match")
117else:
118 for i in range(len(my_samples)):
119 print(prompts[i])
120 samples = my_samples[i]
121 display(Audio(samples.detach().squeeze().to("cpu").numpy(), rate=24000))
122
123del my_samples, samples1@misc{VyvoTTS-LFM2-350M,
2 title={VyvoTTS-LFM2-350M},
3 author={Vyvo},
4 year={2025},
5 howpublished={\url{https://huggingface.co/Vyvo/VyvoTTS-LFM2-350M}}
6}