Views
No views yet
1from unsloth import FastLanguageModel
2import torch
3from snac import SNAC
4from peft import PeftModel
5import soundfile as sf
6
7base_model, tokenizer = FastLanguageModel.from_pretrained(
8 model_name="unsloth/orpheus-3b-0.1-ft",
9 max_seq_length=4096,
10 dtype=None,
11 load_in_4bit=False,
12)
13
14model = PeftModel.from_pretrained(
15 base_model,
16 "lhaatveit/orpheus-3b-0.1-ft-no-lora"
17)
18
19FastLanguageModel.for_inference(model)
20
21snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
22snac_model.to("cpu")
23
24prompts = [
25 "Hei! Jeg er en språkmodell som kan snakke norsk. Jeg har trent på flere timer med norsk tale.",
26 "Jeg bygger på modellen Orpheus fra Canopy AI.",
27 "Man skal ikke selge skinnet før bjørnen er skutt.",
28 "Ibsens ripsbærbusker og andre buskvekster."
29]
30
31chosen_voice = "zoe"
32
33prompts_ = [(f"{chosen_voice}: " + p) if chosen_voice else p for p in prompts]
34
35all_input_ids = []
36
37for prompt in prompts_:
38 input_ids = tokenizer(prompt, return_tensors = "pt").input_ids
39 all_input_ids.append(input_ids)
40
41start_of_human_token = 128259
42end_of_text_token = 128009
43end_of_human_token = 128260
44pad_token = 128263
45start_of_speech_token = 128257
46end_of_speech_token = 128258
47
48
49start_token = torch.tensor([[start_of_human_token]], dtype = torch.int64)
50end_tokens = torch.tensor([[end_of_text_token, end_of_human_token]], dtype = torch.int64)
51
52all_modified_input_ids = []
53
54for input_ids in all_input_ids:
55 modified_input_ids = torch.cat([start_token, input_ids, end_tokens], dim = 1)
56 all_modified_input_ids.append(modified_input_ids)
57
58all_padded_tensors = []
59all_attention_masks = []
60max_length = max([modified_input_ids.shape[1] for modified_input_ids in all_modified_input_ids])
61
62for modified_input_ids in all_modified_input_ids:
63 padding = max_length - modified_input_ids.shape[1]
64 padded_tensor = torch.cat([torch.full((1, padding), pad_token, dtype = torch.int64), modified_input_ids], dim = 1)
65 attention_mask = torch.cat([torch.zeros((1, padding), dtype = torch.int64), torch.ones((1, modified_input_ids.shape[1]), dtype = torch.int64)], dim = 1)
66 all_padded_tensors.append(padded_tensor)
67 all_attention_masks.append(attention_mask)
68
69all_padded_tensors = torch.cat(all_padded_tensors, dim = 0)
70all_attention_masks = torch.cat(all_attention_masks, dim = 0)
71
72input_ids = all_padded_tensors.to("cuda")
73attention_mask = all_attention_masks.to("cuda")
74
75generated_ids = model.generate(
76 input_ids = input_ids,
77 attention_mask = attention_mask,
78 max_new_tokens = 3000,
79 do_sample = True,
80 temperature = 0.6,
81 top_p = 0.95,
82 repetition_penalty = 1.1,
83 num_return_sequences = 1,
84 eos_token_id = 128258,
85 use_cache = True,
86 )
87
88token_to_find = start_of_speech_token
89token_to_remove = end_of_speech_token
90
91processed_rows = []
92code_lists = []
93
94for row in generated_ids:
95
96 sos = (row == token_to_find).nonzero()[0]
97 eos_tokens = (row == token_to_remove).nonzero()
98 eos = (row == token_to_remove).nonzero()[0]
99 snac_tokens = row[(sos[-1] + 1):eos[0]]
100
101 invalid_offset = ((snac_tokens < 128266) + (snac_tokens >= (128266 + 7 * 4096))).nonzero()
102 invalid = len(invalid_offset)
103
104 if invalid:
105 print ("row has illegal SNAC tokens")
106 new_offset = invalid_offset[0]
107 snac_tokens = snac_tokens[0:new_offset]
108
109 modified_len = len(snac_tokens) // 7 * 7
110
111 snac_codes = snac_tokens[:modified_len]
112 snac_codes = [t - 128266 for t in snac_codes]
113
114 code_lists.append(snac_codes)
115
116def snac_decode(code_list):
117 layer_1 = []
118 layer_2 = []
119 layer_3 = []
120 for i in range((len(code_list)+1)//7):
121 layer_1.append(code_list[7*i])
122 layer_2.append(code_list[7*i+1]-4096)
123 layer_3.append(code_list[7*i+2]-(2*4096))
124 layer_3.append(code_list[7*i+3]-(3*4096))
125 layer_2.append(code_list[7*i+4]-(4*4096))
126 layer_3.append(code_list[7*i+5]-(5*4096))
127 layer_3.append(code_list[7*i+6]-(6*4096))
128 codes = [torch.tensor(layer_1).unsqueeze(0),
129 torch.tensor(layer_2).unsqueeze(0),
130 torch.tensor(layer_3).unsqueeze(0)]
131 audio_hat = snac_model.decode(codes)
132 return audio_hat
133
134my_samples = []
135
136for code_list in code_lists:
137 samples = snac_decode(code_list)
138 my_samples.append(samples)
139
140for i in range(len(my_samples)):
141 print(prompts[i])
142 samples = my_samples[i]
143 audio = samples.detach().squeeze().to("cpu").numpy()
144 sf.write(f"sample_{i}.wav", audio, 24000)