Spaces:
Sleeping
Sleeping
# --coding:utf-8-- | |
import os | |
from encoder.utils import convert_audio | |
import torchaudio | |
import torch | |
from decoder.pretrained import WavTokenizer | |
import time | |
import logging | |
device1=torch.device('cuda:0') | |
device2=torch.device('cpu') | |
input_path = "./WavTokenizer/data/infer/lirbitts_testclean" | |
out_folder = './WavTokenizer/result/infer' | |
# os.system("rm -r %s"%(out_folder)) | |
# os.system("mkdir -p %s"%(out_folder)) | |
# ll="libritts_testclean500_large" | |
ll="wavtokenizer_smalldata_frame40_3s_nq1_code4096_dim512_kmeans200_attn_testclean_epoch34" | |
tmptmp=out_folder+"/"+ll | |
os.system("rm -r %s"%(tmptmp)) | |
os.system("mkdir -p %s"%(tmptmp)) | |
# 自己数据模型加载 | |
config_path = "./WavTokenizer/configs/wavtokenizer_smalldata_frame40_3s_nq1_code4096_dim512_kmeans200_attn.yaml" | |
model_path = "./WavTokenizer/result/train/wavtokenizer_smalldata_frame40_3s_nq1_code4096_dim512_kmeans200_attn/lightning_logs/version_3/checkpoints/wavtokenizer_checkpoint_epoch=24_step=137150_val_loss=5.6731.ckpt" | |
wavtokenizer = WavTokenizer.from_pretrained0802(config_path, model_path) | |
wavtokenizer = wavtokenizer.to(device1) | |
# wavtokenizer = wavtokenizer.to(device2) | |
with open(input_path,'r') as fin: | |
x=fin.readlines() | |
x = [i.strip() for i in x] | |
# 完成一些加速处理 | |
features_all=[] | |
for i in range(len(x)): | |
wav, sr = torchaudio.load(x[i]) | |
# print("***:",x[i]) | |
# wav = convert_audio(wav, sr, 24000, 1) # (1,131040) | |
bandwidth_id = torch.tensor([0]) | |
wav=wav.to(device1) | |
print(i) | |
features,discrete_code= wavtokenizer.encode_infer(wav, bandwidth_id=bandwidth_id) | |
features_all.append(features.cpu()) | |
wavtokenizer = wavtokenizer.to(device2) | |
for i in range(len(x)): | |
bandwidth_id = torch.tensor([0]) | |
print(i) | |
audio_out = wavtokenizer.decode(features_all[i], bandwidth_id=bandwidth_id) | |
# print(i,time.time()) | |
# breakpoint() # (1, 131200) | |
audio_path = out_folder + '/' + ll + '/' + x[i].split('/')[-1] | |
# os.makedirs(out_folder + '/' + ll, exist_ok=True) | |
torchaudio.save(audio_path, audio_out, sample_rate=24000, encoding='PCM_S', bits_per_sample=16) | |