diff --git a/scripts/process_old_dataset.py b/scripts/process_dataset.py similarity index 100% rename from scripts/process_old_dataset.py rename to scripts/process_dataset.py diff --git a/scripts/transcribe_dataset.py b/scripts/transcribe_dataset.py new file mode 100644 index 0000000..ca8dcfa --- /dev/null +++ b/scripts/transcribe_dataset.py @@ -0,0 +1,79 @@ +import os +import json +import torch +import torchaudio +import whisperx + +from tqdm.auto import tqdm +from pathlib import Path + +device = "cuda" +batch_size = 16 +dtype = "float16" +model_size = "large-v2" + +input_audio = "voice" +output_dataset = "metadata" +skip_existing = True + +model = whisperx.load_model(model_size, device, compute_type=dtype) + +align_model, align_model_metadata, align_model_language = (None, None, None) + +def pad(num, zeroes): + return str(num).zfill(zeroes+1) + +for dataset_name in os.listdir(f'./{input_audio}/'): + if not os.path.isdir(f'./{input_audio}/{dataset_name}/'): + print("Is not dir:", f'./{input_audio}/{dataset_name}/') + continue + + for speaker_id in tqdm(os.listdir(f'./{input_audio}/{dataset_name}/'), desc="Processing speaker"): + if not os.path.isdir(f'./{input_audio}/{dataset_name}/{speaker_id}'): + print("Is not dir:", f'./{input_audio}/{dataset_name}/{speaker_id}') + continue + + outpath = Path(f'./{output_dataset}/{dataset_name}/{speaker_id}/whisper.json') + + if outpath.exists(): + metadata = json.loads(open(outpath, 'r', encoding='utf-8').read()) + else: + os.makedirs(f'./{output_dataset}/{dataset_name}/{speaker_id}/', exist_ok=True) + metadata = {} + + for filename in tqdm(os.listdir(f'./{input_audio}/{dataset_name}/{speaker_id}/'), desc=f"Processing speaker: {speaker_id}"): + + if skip_existing and filename in metadata: + continue + + inpath = f'./{input_audio}/{dataset_name}/{speaker_id}/{filename}' + + metadata[filename] = { + "segments": [], + "language": "", + "text": [], + } + + audio = whisperx.load_audio(inpath) + result = model.transcribe(audio, batch_size=batch_size) + language = result["language"] + + if align_model_language != language: + tqdm.write(f'Loading language: {language}') + align_model, align_model_metadata = whisperx.load_align_model(language_code=language, device=device) + align_model_language = language + + result = whisperx.align(result["segments"], align_model, align_model_metadata, audio, device, return_char_alignments=False) + + metadata[filename]["segments"] = result["segments"] + metadata[filename]["language"] = language + + text = [] + for segment in result["segments"]: + id = len(text) + text.append( segment["text"] ) + metadata[filename]["segments"][id]["id"] = id + + metadata[filename]["text"] = " ".join(text).strip() + + open(outpath, 'w', encoding='utf-8').write(json.dumps(metadata)) \ No newline at end of file diff --git a/vall_e/data.py b/vall_e/data.py index 434b3d3..cf512be 100755 --- a/vall_e/data.py +++ b/vall_e/data.py @@ -177,10 +177,8 @@ def _get_phones(path, language="en"): content = metadata["phonemes"] else: content = open(_get_phone_path(path), "r", encoding="utf-8").read().split(" ") - content = _cleanup_phones( content ) return "".join(content) - #return [""] + [ " " if not p else p for p in content ] + [""] def _interleaved_reorder(l, fn): groups = defaultdict(list) @@ -431,8 +429,6 @@ class Dataset(_Dataset): text = cfg.hdf5[key]["text"][:] resps = cfg.hdf5[key]["audio"][:, :] - - text = np.array( _cleanup_phones( text, targets=[ self.phone_symmap[" "] ] ) ) text = torch.from_numpy(text).to(self.text_dtype) resps = torch.from_numpy(resps).to(torch.int16) @@ -455,12 +451,13 @@ class Dataset(_Dataset): txt = cfg.hdf5[key]["text"][:] qnt = cfg.hdf5[key]["audio"][:, :] - txt = np.array( _cleanup_phones( txt, targets=[ self.phone_symmap[" "] ] ) ) + txt = np.array( txt ) txt = torch.from_numpy(txt).to(self.text_dtype) qnt = torch.from_numpy(qnt).to(torch.int16) else: - txt = torch.tensor([*map(self.phone_symmap.get, _get_phones(sampled_path))]).to(self.text_dtype) + #txt = torch.tensor([*map(self.phone_symmap.get, _get_phones(sampled_path))]).to(self.text_dtype) + txt = torch.tensor(tokenize(_get_phones(sampled_path))).to(self.text_dtype) qnt = _load_quants(sampled_path) # [original text] [new text]