File size: 5,659 Bytes
359afe5 45089ef 359afe5 dba24db 359afe5 45089ef 359afe5 aa67214 359afe5 aa67214 359afe5 8df6804 359afe5 99d9342 aa67214 99d9342 aa67214 359afe5 8df6804 359afe5 aa67214 359afe5 aa67214 359afe5 aa67214 359afe5 aa67214 359afe5 dba24db 45089ef aa67214 45089ef aa67214 45089ef aa67214 45089ef aa67214 45089ef |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 |
import os
import time
import torch
import torchaudio
from transformers import (
Wav2Vec2Processor, HubertForCTC,
WhisperProcessor, WhisperForConditionalGeneration, Wav2Vec2ForCTC, AutoProcessor, AutoModelForCTC
)
from .cmu_process import text_to_phoneme, cmu_to_ipa, clean_cmu
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)
# === Helper: move all tensors to model device ===
def to_device(batch, device):
if isinstance(batch, dict):
return {k: v.to(device) for k, v in batch.items()}
elif isinstance(batch, torch.Tensor):
return batch.to(device)
return batch
# === Setup: Load all 3 models ===
# 1. Base HuBERT
base_proc = Wav2Vec2Processor.from_pretrained("facebook/hubert-large-ls960-ft")
base_model = HubertForCTC.from_pretrained("facebook/hubert-large-ls960-ft").to(device).eval()
# 2. Whisper + phonemizer
whisper_proc = WhisperProcessor.from_pretrained("openai/whisper-base")
whisper_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-base").to(device).eval()
# 3. My Hubert Model (optional HF token via env)
HF_TOKEN = os.environ.get("HF_TOKEN", None)
# print(HF_TOKEN)
proc = Wav2Vec2Processor.from_pretrained("tecasoftai/hubert-finetune", token=HF_TOKEN)
model = HubertForCTC.from_pretrained("tecasoftai/hubert-finetune", token=HF_TOKEN).to(device).eval()
# 4. wav2vec2-xls-r-300m-timit-phoneme
# load model and processor
timit_proc = Wav2Vec2Processor.from_pretrained("vitouphy/wav2vec2-xls-r-300m-timit-phoneme")
timit_model = Wav2Vec2ForCTC.from_pretrained("vitouphy/wav2vec2-xls-r-300m-timit-phoneme").to(device).eval()
# 5 bookbot/wav2vec2-ljspeech-gruut
gruut_processor = AutoProcessor.from_pretrained("bookbot/wav2vec2-ljspeech-gruut")
gruut_model = AutoModelForCTC.from_pretrained("bookbot/wav2vec2-ljspeech-gruut").to(device).eval()
# 6 microsoft/wavlm-large-english-phoneme
wavlm_proc = AutoProcessor.from_pretrained("speech31/wavlm-large-english-phoneme")
wavlm_model = AutoModelForCTC.from_pretrained("speech31/wavlm-large-english-phoneme").to(device).eval()
# === Inference functions ===
def run_hubert_base(wav):
start = time.time()
inputs = base_proc(wav, sampling_rate=16000, return_tensors="pt", padding=True).input_values
inputs = inputs.to(device)
with torch.no_grad():
logits = base_model(inputs).logits
ids = torch.argmax(logits, dim=-1)
text = base_proc.batch_decode(ids)[0]
# Convert to phonemes (CMU-like string without stresses)
phonemes = text_to_phoneme(text)
phonemes = cmu_to_ipa(phonemes)
return phonemes.strip(), time.time() - start
def run_whisper(wav):
start = time.time()
inputs = whisper_proc(wav, sampling_rate=16000, return_tensors="pt")
input_features = inputs.input_features.to(device)
attention_mask = inputs.get("attention_mask", None)
gen_kwargs = {"language": "en"}
if attention_mask is not None:
gen_kwargs["attention_mask"] = attention_mask.to(device)
with torch.no_grad():
pred_ids = whisper_model.generate(input_features, **gen_kwargs)
text = whisper_proc.batch_decode(pred_ids, skip_special_tokens=True)[0]
phonemes = text_to_phoneme(text)
phonemes = cmu_to_ipa(phonemes)
return phonemes.strip(), time.time() - start
def run_model(wav):
start = time.time()
# Prepare input (BatchEncoding supports .to(device))
inputs = proc(wav, sampling_rate=16000, return_tensors="pt").to(device)
# Forward pass
with torch.no_grad():
logits = model(**inputs).logits
# Greedy decode
ids = torch.argmax(logits, dim=-1)
phonemes = proc.batch_decode(ids)[0]
phonemes = cmu_to_ipa(phonemes)
return phonemes.strip(), time.time() - start
def run_timit(wav):
start = time.time()
# Read and process the input
inputs = timit_proc(wav, sampling_rate=16_000, return_tensors="pt", padding=True)
inputs = inputs.to(device)
# Forward pass
with torch.no_grad():
logits = timit_model(inputs.input_values, attention_mask=inputs.attention_mask).logits
# Decode id into string
predicted_ids = torch.argmax(logits, axis=-1)
phonemes = timit_proc.batch_decode(predicted_ids)
phonemes = "".join(phonemes)
return phonemes.strip(), time.time() - start
def run_gruut(wav):
start = time.time()
# Preprocess waveform → model input
inputs = gruut_processor(
wav,
sampling_rate=16000,
return_tensors="pt",
padding=True
).to(device)
# Forward pass
with torch.no_grad():
logits = gruut_model(**inputs).logits
# Greedy decode → IPA phonemes
pred_ids = torch.argmax(logits, dim=-1)
phonemes = gruut_processor.batch_decode(pred_ids)[0]
phonemes = "".join(phonemes)
return phonemes.strip(), time.time() - start
def run_wavlm_large_phoneme(wav):
start = time.time()
# Preprocess waveform → model input
inputs = wavlm_proc(
wav,
sampling_rate=16000,
return_tensors="pt",
padding=True
).to(device)
input_values = inputs.input_values
attention_mask = inputs.get("attention_mask", None)
# Forward pass
with torch.no_grad():
logits = wavlm_model(input_values, attention_mask=attention_mask).logits
# Greedy decode → phoneme tokens
pred_ids = torch.argmax(logits, dim=-1)
phonemes = wavlm_proc.batch_decode(pred_ids)[0]
phonemes = "".join(phonemes)
return phonemes.strip(), time.time() - start |