deepgeek/CosyVoice
2
0
Fork 0
mirror of https://github.com/FunAudioLLM/CosyVoice.git synced 2026-02-04 17:39:25 +08:00
Code Issues Packages Projects Releases Wiki Activity

use thread pool in tools

Browse Source
This commit is contained in:
lyuxiang.lx
2024-09-18 17:41:15 +08:00
parent 2665b06e95
commit ff8e63567a
2 changed files with 61 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files

77
tools/extract_embedding.py
View File

@@ -13,74 +13,39 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import os
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import ThreadPoolExecutor, as_completed
import onnxruntime
import torch
import torchaudio
import torchaudio.compliance.kaldi as kaldi
from tqdm import tqdm
from itertools import repeat
def extract_embedding(utt: str, wav_file: str, ort_session: onnxruntime.InferenceSession):
audio, sample_rate = torchaudio.load(wav_file)
def single_job(utt):
audio, sample_rate = torchaudio.load(utt2wav[utt])
if sample_rate != 16000:
audio = torchaudio.transforms.Resample(
orig_freq=sample_rate, new_freq=16000
)(audio)
feat = kaldi.fbank(audio, num_mel_bins=80, dither=0, sample_frequency=16000)
audio = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(audio)
feat = kaldi.fbank(audio,
num_mel_bins=80,
dither=0,
sample_frequency=16000)
feat = feat - feat.mean(dim=0, keepdim=True)
embedding = ort_session.run(None, {ort_session.get_inputs()[0].name: feat.unsqueeze(dim=0).cpu().numpy()})[0].flatten().tolist()
return (utt, embedding)
return utt, embedding
def main(args):
utt2wav, utt2spk = {}, {}
with open("{}/wav.scp".format(args.dir)) as f:
for l in f:
l = l.replace("\n", "").split()
utt2wav[l[0]] = l[1]
with open("{}/utt2spk".format(args.dir)) as f:
for l in f:
l = l.replace("\n", "").split()
utt2spk[l[0]] = l[1]
assert os.path.exists(args.onnx_path), "onnx_path not exists"
option = onnxruntime.SessionOptions()
option.graph_optimization_level = (
onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
)
option.intra_op_num_threads = 1
providers = ["CPUExecutionProvider"]
ort_session = onnxruntime.InferenceSession(
args.onnx_path, sess_options=option, providers=providers
)
all_utt = utt2wav.keys()
with ThreadPoolExecutor(max_workers=args.num_thread) as executor:
results = list(
tqdm(
executor.map(extract_embedding, all_utt, [utt2wav[utt] for utt in all_utt], repeat(ort_session)),
total=len(utt2wav),
desc="Process data: "
)
)
all_task = [executor.submit(single_job, utt) for utt in utt2wav.keys()]
utt2embedding, spk2embedding = {}, {}
for utt, embedding in results:
for future in tqdm(as_completed(all_task)):
utt, embedding = future.result()
utt2embedding[utt] = embedding
spk = utt2spk[utt]
if spk not in spk2embedding:
spk2embedding[spk] = []
spk2embedding[spk].append(embedding)
for k, v in spk2embedding.items():
spk2embedding[k] = torch.tensor(v).mean(dim=0).tolist()
torch.save(utt2embedding, "{}/utt2embedding.pt".format(args.dir))
torch.save(spk2embedding, "{}/spk2embedding.pt".format(args.dir))
@@ -91,4 +56,22 @@ if __name__ == "__main__":
parser.add_argument("--onnx_path", type=str)
parser.add_argument("--num_thread", type=int, default=8)
args = parser.parse_args()
utt2wav, utt2spk = {}, {}
with open('{}/wav.scp'.format(args.dir)) as f:
for l in f:
l = l.replace('\n', '').split()
utt2wav[l[0]] = l[1]
with open('{}/utt2spk'.format(args.dir)) as f:
for l in f:
l = l.replace('\n', '').split()
utt2spk[l[0]] = l[1]
option = onnxruntime.SessionOptions()
option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
option.intra_op_num_threads = 1
providers = ["CPUExecutionProvider"]
ort_session = onnxruntime.InferenceSession(args.onnx_path, sess_options=option, providers=providers)
executor = ThreadPoolExecutor(max_workers=args.num_thread)
main(args)