Add inference code

inference.py

@@ -0,0 +1,189 @@
+import librosa
+
+import torchaudio
+import torch
+
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+)
+
+from speechtokenizer import SpeechTokenizer
+from WavTokenizer.decoder.pretrained import WavTokenizer
+from audiotools import AudioSignal
+
+
+def resample(audio_data: torch.Tensor, sample_rate: int):
+    print("Inout sample rate:", sample_rate)
+    if sample_rate == 24000:
+      audio_data24k = audio_data
+      audio_data16k = torch.tensor(
+          librosa.resample(
+              audio_data.cpu().detach().numpy(), orig_sr=sample_rate, target_sr=16000
+          )
+      )
+    elif sample_rate == 16000:
+      audio_data16k = audio_data
+      audio_data24k = torch.tensor(
+          librosa.resample(
+              audio_data.cpu().detach().numpy(), orig_sr=sample_rate, target_sr=24000
+          )
+      )
+    else:
+      print("Resampling everything")
+      audio_data16k = torch.tensor(
+          librosa.resample(
+              audio_data.cpu().detach().numpy(), orig_sr=sample_rate, target_sr=16000
+          )
+      )
+      audio_data24k = torch.tensor(
+          librosa.resample(
+              audio_data.cpu().detach().numpy(), orig_sr=sample_rate, target_sr=24000
+          )
+      )
+
+    return (audio_data16k.view(1, -1).float().to(device), 
+           audio_data24k.view(1, -1).float().to(device))
+
+
+def decode_tts(tokens, quantizer, n_codebooks, n_original_tokens, start_audio_token_id, end_audio_token_id):
+    # find start and end indices of audio tokens
+    start = torch.nonzero(tokens == start_audio_token_id)
+    end = torch.nonzero(tokens == end_audio_token_id)
+
+    start = start[0, -1] + 1 if len(start) else 0
+    end = end[0, -1] if len(end) else tokens.shape[-1]
+
+    # subtract length of original vocabulary -> tokens in range [0, 1024)
+    audio_tokens = tokens[start:end] % n_original_tokens
+    reminder = audio_tokens.shape[-1] % n_codebooks
+
+    if reminder:
+        # pad if last frame is incomplete
+        pad_tokens = torch.zeros(n_codebooks - reminder, device="cuda")
+        audio_tokens = torch.cat([audio_tokens, pad_tokens], dim=0)
+
+    transposed = audio_tokens.view(-1, n_codebooks).t()
+    codes = transposed.view(n_codebooks, 1, -1).to(device)
+
+    audio = quantizer.decode(codes).squeeze(0)
+
+    del tokens
+    del audio_tokens
+    torch.cuda.empty_cache()
+
+    return AudioSignal(audio.detach().cpu().numpy(), quantizer.sample_rate)
+
+
+def infer_text_to_audio(text, model, tokenizer, quantizer, max_seq_length=1024, top_k=20):
+    text_tokenized = tokenizer(text, return_tensors="pt")
+    text_input_tokens = text_tokenized["input_ids"].to(device)
+
+    soa = tokenizer(start_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+    eoa = tokenizer(end_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+
+    text_tokens = torch.cat([text_input_tokens, soa], dim=1)
+    attention_mask = torch.ones(text_tokens.size(), device=device)
+
+    output_audio_tokens = model.generate(
+        text_tokens,
+        attention_mask=attention_mask,
+        max_new_tokens=max_seq_length,
+        top_k=top_k,
+        do_sample=True,
+        temperature=0.1,
+        repetition_penalty=1.1,
+        length_penalty=1.2,
+        no_repeat_ngram_size=3,
+    )
+
+    audio_signal = decode_tts(output_audio_tokens[0], quantizer, 3, len(tokenizer), soa, eoa)
+
+    return audio_signal
+
+
+def infer_audio_to_text(audio_path, model, tokenizer, quantizer_speech, quantizer_wav, max_seq_length=1024, top_k=20):
+    audio_data, sample_rate = torchaudio.load(audio_path)
+
+    audio_16k, audio_24k = resample(audio_data, sample_rate)
+    bandwidth_id = torch.tensor([0])
+
+    codes_semantics = quantizer_speech.encode(audio_16k.reshape(1, 1, -1))
+    raw_semantic_tokens = codes_semantics + len(tokenizer)
+    raw_semantic_tokens = raw_semantic_tokens[:1].view(1, -1)
+
+    _, codes = quantizer_wav.encode_infer(audio_24k, bandwidth_id=bandwidth_id)
+    raw_acoustic_tokens = codes + len(tokenizer) + 1024
+    raw_acoustic_tokens = raw_acoustic_tokens.view(1, -1)
+
+    audio_tokens = torch.cat([raw_semantic_tokens, raw_acoustic_tokens], dim=1)
+
+    soa = tokenizer(start_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+    eoa = tokenizer(end_audio_token, return_tensors="pt")["input_ids"][:, -1:].to(device)
+    audio_tokens = torch.cat([soa, audio_tokens, eoa], dim=1)
+    
+    # text_tokens = tokenizer("is said with", return_tensors="pt")["input_ids"].to(device)
+    tokens = torch.cat([audio_tokens], dim=1)
+
+    attention_mask = torch.ones(tokens.size(), device=device)
+
+    output_text_tokens = model.generate(
+        tokens,
+        attention_mask=attention_mask,
+        max_new_tokens=max_seq_length,
+        do_sample=True,
+        temperature=0.1,
+        top_p=0.9,
+        top_k=top_k,
+    )
+
+    output_text_tokens = output_text_tokens.cpu()[0]
+    output_text_tokens = output_text_tokens[output_text_tokens < tokenizer(start_audio_token)["input_ids"][-1]]
+    decoded_text = tokenizer.decode(output_text_tokens, skip_special_tokens=True)
+
+    return decoded_text
+
+
+device = "cuda"
+
+n_codebooks_tts = 3
+n_codebooks_asr = 1
+
+start_audio_token = "<|start_of_audio|>"
+end_audio_token = "<|end_of_audio|>"
+end_sequence_token = "<|end_of_text|>"
+
+base_model = "Vikhrmodels/salt-asr_speech_1_wav_1_tts_speech_3_text-10k"
+
+
+if __name__ == "__main__":
+    tokenizer = AutoTokenizer.from_pretrained(base_model, cache_dir=".")
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model,
+        cache_dir=".",
+        torch_dtype=torch.bfloat16,
+        attn_implementation="sdpa",
+        device_map={"": 0}
+    )
+
+    quantizer_speech = SpeechTokenizer.load_from_checkpoint("speechtokenizer/config.json",
+                                                            "speechtokenizer/SpeechTokenizer.pt")
+    quantizer_speech = quantizer_speech.eval().to(device)
+    codebook_size = quantizer_speech.quantizer.bins
+
+    quantizer_wav = WavTokenizer.from_pretrained0802("wavtokenizer/config.yaml",
+                                                     "wavtokenizer/WavTokenizer_small_600_24k_4096.ckpt")
+    quantizer_wav = quantizer_wav.to(device)
+
+    text = ("Say 'COUNT NUMBERS FROM ONE TO TEN' with a male speaker delivers a very monotone and "
+            "low-pitched speech with a moderate speed in a setting with almost no noise, "
+            "creating a clear and quiet recording.")
+
+    audio_signal = infer_text_to_audio(text, model, tokenizer, quantizer_speech, top_k=60)
+    audio_signal.write("output.wav")
+
+    audio_path = "./input.wav"
+    generated_text = infer_audio_to_text(audio_path, model, tokenizer, quantizer_speech, quantizer_wav, top_k=10)
+    print(generated_text)
+
+