cnywt/SyncTalk

1from transformers import Wav2Vec2Processor, HubertModel2import soundfile as sf3import numpy as np4import torch5 6print("Loading the Wav2Vec2 Processor...")7wav2vec2_processor = Wav2Vec2Processor.from_pretrained("facebook/hubert-large-ls960-ft")8print("Loading the HuBERT Model...")9hubert_model = HubertModel.from_pretrained("facebook/hubert-large-ls960-ft")10 11 12def get_hubert_from_16k_wav(wav_16k_name):13    speech_16k, _ = sf.read(wav_16k_name)14    hubert = get_hubert_from_16k_speech(speech_16k)15    return hubert16 17@torch.no_grad()18def get_hubert_from_16k_speech(speech, device="cuda:0"):19    global hubert_model20    hubert_model = hubert_model.to(device)21    if speech.ndim ==2:22        speech = speech[:, 0] # [T, 2] ==> [T,]23    input_values_all = wav2vec2_processor(speech, return_tensors="pt", sampling_rate=16000).input_values # [1, T]24    input_values_all = input_values_all.to(device)25    # For long audio sequence, due to the memory limitation, we cannot process them in one run26    # HuBERT process the wav with a CNN of stride [5,2,2,2,2,2], making a stride of 32027    # Besides, the kernel is [10,3,3,3,3,2,2], making 400 a fundamental unit to get 1 time step.28    # So the CNN is euqal to a big Conv1D with kernel k=400 and stride s=32029    # We have the equation to calculate out time step: T = floor((t-k)/s)30    # To prevent overlap, we set each clip length of (K+S*(N-1)), where N is the expected length T of this clip31    # The start point of next clip should roll back with a length of (kernel-stride) so it is stride * N32    kernel = 40033    stride = 32034    clip_length = stride * 100035    num_iter = input_values_all.shape[1] // clip_length36    expected_T = (input_values_all.shape[1] - (kernel-stride)) // stride37    res_lst = []38    for i in range(num_iter):39        if i == 0:40            start_idx = 041            end_idx = clip_length - stride + kernel42        else:43            start_idx = clip_length * i44            end_idx = start_idx + (clip_length - stride + kernel)45        input_values = input_values_all[:, start_idx: end_idx]46        hidden_states = hubert_model.forward(input_values).last_hidden_state # [B=1, T=pts//320, hid=1024]47        res_lst.append(hidden_states[0])48    if num_iter > 0:49        input_values = input_values_all[:, clip_length * num_iter:]50    else:51        input_values = input_values_all52    # if input_values.shape[1] != 0:53    if input_values.shape[1] >= kernel: # if the last batch is shorter than kernel_size, skip it            54        hidden_states = hubert_model(input_values).last_hidden_state # [B=1, T=pts//320, hid=1024]55        res_lst.append(hidden_states[0])56    ret = torch.cat(res_lst, dim=0).cpu() # [T, 1024]57    # assert ret.shape[0] == expected_T58    assert abs(ret.shape[0] - expected_T) <= 159    if ret.shape[0] < expected_T:60        ret = torch.nn.functional.pad(ret, (0,0,0,expected_T-ret.shape[0]))61    else:62        ret = ret[:expected_T]63    return ret64 65def make_even_first_dim(tensor):66    size = list(tensor.size())67    if size[0] % 2 == 1:68        size[0] -= 169        return tensor[:size[0]]70    return tensor71 72import soundfile as sf73import numpy as np74import torch75from argparse import ArgumentParser76import librosa77 78parser = ArgumentParser()79parser.add_argument('--wav', type=str, help='')80args = parser.parse_args()81 82wav_name = args.wav83 84speech, sr = sf.read(wav_name)85speech_16k = librosa.resample(speech, orig_sr=sr, target_sr=16000)86print("SR: {} to {}".format(sr, 16000))87# print(speech.shape, speech_16k.shape)88 89hubert_hidden = get_hubert_from_16k_speech(speech_16k)90hubert_hidden = make_even_first_dim(hubert_hidden).reshape(-1, 2, 1024)91np.save(wav_name.replace('.wav', '_hu.npy'), hubert_hidden.detach().numpy())92print(hubert_hidden.detach().numpy().shape)