from unsloth import FastLanguageModelfrom snac import SNACimport soundfile as sfimport numpy as npfrom loguru import loggerfrom df.enhance import init_df, enhance, save_audioimport torchimport librosaimport torchaudioimport os #Name of the modelMODEL_NAME = 'snorbyte/snorTTS-Indic-v0'MAX_SEQ_LENGTH = 4096HUGGINGFACE_TOKEN = "" # ! Add your hugging face token # Load the model and tokenizer.model, tokenizer = FastLanguageModel.from_pretrained( model_name=MODEL_NAME, # load_in_4bit=True, max_seq_length=MAX_SEQ_LENGTH, token=HUGGINGFACE_TOKEN,)logger.success(f"Loaded model: {MODEL_NAME}") # Load the end of speech token for the tokenizer.tokeniser_length = 128256end_of_speech_id = tokeniser_length + 2pad_token_id = tokeniser_length + 7audio_start_id = tokeniser_length + 10 pad_token = tokenizer.decode([pad_token_id])logger.success("Load special tokens for the tokenizer.") # Wrap Model for InferenceFastLanguageModel.for_inference(model)logger.success(f"{MODEL_NAME} is ready for inference.") # Set the padding token and padding side.tokenizer.pad_token = pad_tokentokenizer.padding_side = "left"logger.success("Set padding token and padding side for the tokenizer.") # Load the SNAC model for audio decoding.snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")logger.success("Loaded SNAC model for audio decoding.") # Load DeepFilter for optional post processingdf_model, df_state, _ = init_df() # Function to generate audiodef generate_audio( row, model, tokenizer, user=False, temperature=0.4, top_p=0.9, repetition_penalty=1.05): try: if user: prompt = row["eval_text_user"] else: prompt = row["eval_text_no_user"] inputs = tokenizer(prompt, add_special_tokens=False, return_tensors="pt") max_tokens = MAX_SEQ_LENGTH - inputs.input_ids.shape[1] output = model.generate( input_ids=inputs.input_ids.to("cuda"), attention_mask=inputs.attention_mask.to("cuda"), max_new_tokens=max_tokens, temperature=temperature, top_p=top_p, repetition_penalty=repetition_penalty, eos_token_id=end_of_speech_id, ) audio_ids = [] for id in output[0]: if id >= audio_start_id: audio_ids.append(id.item()) clean_audio_ids = [] for i in range((len(audio_ids) + 1) // 7): for j in range(7): clean_audio_ids += [audio_ids[7 * i + j] - audio_start_id] codes = [[], [], []] for i in range((len(clean_audio_ids) + 1) // 7): codes[0].append(clean_audio_ids[7 * i]) codes[1].append(clean_audio_ids[7 * i + 1] - 4096) codes[2].append(clean_audio_ids[7 * i + 2] - (2 * 4096)) codes[2].append(clean_audio_ids[7 * i + 3] - (3 * 4096)) codes[1].append(clean_audio_ids[7 * i + 4] - (4 * 4096)) codes[2].append(clean_audio_ids[7 * i + 5] - (5 * 4096)) codes[2].append(clean_audio_ids[7 * i + 6] - (6 * 4096)) codes = [ torch.tensor(codes[0]).unsqueeze(0), torch.tensor(codes[1]).unsqueeze(0), torch.tensor(codes[2]).unsqueeze(0), ] try: audio = snac_model.decode(codes) except Exception as e: logger.error(f"Error decoding audio: {e}") return None return audio.detach().squeeze().to("cpu").numpy() except Exception as e: logger.error(f"Error decoding audio: {e}") return None # Run inference.# * Please refer to the training script to create prompt from SNAC tokens.row = { "eval_text_user": f"<custom_token_3><|begin_of_text|>kannada142: ಅಯ್ಯಯ್ಯೋ... Whitefield ಗೆ reach ಆಗೋಕೆ almost 10 hours ಆಯ್ತು you know... traffic was so terrible today <|eot_id|><custom_token_4><custom_token_5><custom_token_1>" } eval_sample = generate_audio(row, model, tokenizer, True)if eval_sample is None: logger.error("Failed to generate audio for evaluation sample.")else: logger.success("Audio Generated. Post Processing Started") ## post-processing settings filename = "eval.wav" speed = 1.05 #add speed up according to speaker denoise = False #denoise if you want output = eval_sample.astype(np.float32) #speed up if abs(speed - 1.0) > 1e-4: output_t = torch.from_numpy(output).unsqueeze(0) output_speed, _ = torchaudio.sox_effects.apply_effects_tensor(output_t, 24_000, effects=[["tempo", f"{speed}"]]) output = output_speed.squeeze(0).cpu().numpy() #denoise if denoise: resampled_48k = librosa.resample(output, orig_sr=24_000, target_sr=48_000) resampled_48k = torch.from_numpy(resampled_48k).unsqueeze(0) output_48k = enhance(df_model, df_state, resampled_48k) output_48k = output_48k.squeeze(0).cpu().numpy() output = librosa.resample(output_48k, orig_sr=48_000, target_sr=24_000) logger.success("Saving Final Output...") #save sf.write(filename, output, 24_000) logger.success(f"Generated and saved evaluation sample audio as {filename}.")