From 597441e48b43524c362c9cc9f7fe758c79a6ea68 Mon Sep 17 00:00:00 2001 From: mrq Date: Mon, 5 Aug 2024 19:40:50 -0500 Subject: [PATCH] moved transcribe and process dataset scripts to vall_e/emb within the module itself, argparse-ified transcription script --- scripts/transcribe_dataset.py | 103 ------------ .../emb/process.py | 77 +++++---- vall_e/emb/transcribe.py | 147 ++++++++++++++++++ 3 files changed, 197 insertions(+), 130 deletions(-) delete mode 100644 scripts/transcribe_dataset.py rename scripts/process_dataset.py => vall_e/emb/process.py (78%) create mode 100644 vall_e/emb/transcribe.py diff --git a/scripts/transcribe_dataset.py b/scripts/transcribe_dataset.py deleted file mode 100644 index 657e527..0000000 --- a/scripts/transcribe_dataset.py +++ /dev/null @@ -1,103 +0,0 @@ -import os -import json -import torch -import torchaudio -import whisperx - -from tqdm.auto import tqdm -from pathlib import Path - -# to-do: use argparser -batch_size = 16 -device = "cuda" -dtype = "float16" -model_name = "large-v3" - -input_audio = "voices" -output_dataset = "training/metadata" - -skip_existing = True -diarize = False - -# -model = whisperx.load_model(model_name, device, compute_type=dtype) -align_model, align_model_metadata, align_model_language = (None, None, None) -if diarize: - diarize_model = whisperx.DiarizationPipeline(device=device) -else: - diarize_model = 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}/'): - 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}'): - 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 - - if ".json" in filename: - continue - - inpath = f'./{input_audio}/{dataset_name}/{speaker_id}/{filename}' - - if os.path.isdir(inpath): - continue - - metadata[filename] = { - "segments": [], - "language": "", - "text": "", - "start": 0, - "end": 0, - } - - audio = whisperx.load_audio(inpath) - result = model.transcribe(audio, batch_size=batch_size) - language = result["language"] - - if language[:2] not in ["ja"]: - language = "en" - - 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 - - if diarize_model is not None: - diarize_segments = diarize_model(audio) - result = whisperx.assign_word_speakers(diarize_segments, result) - - text = [] - start = 0 - end = 0 - for segment in result["segments"]: - text.append( segment["text"] ) - start = min( start, segment["start"] ) - end = max( end, segment["end"] ) - - metadata[filename]["text"] = " ".join(text).strip() - metadata[filename]["start"] = start - metadata[filename]["end"] = end - - open(outpath, 'w', encoding='utf-8').write(json.dumps(metadata)) \ No newline at end of file diff --git a/scripts/process_dataset.py b/vall_e/emb/process.py similarity index 78% rename from scripts/process_dataset.py rename to vall_e/emb/process.py index 19cec86..357f38f 100644 --- a/scripts/process_dataset.py +++ b/vall_e/emb/process.py @@ -1,3 +1,8 @@ +""" +# Handles processing audio provided through --input-audio of adequately annotated transcriptions provided through --input-metadata (through transcribe.py) +# Outputs NumPy objects containing quantized audio and adequate metadata for use of loading in the trainer through --output-dataset +""" + import os import json import argparse @@ -7,8 +12,8 @@ import numpy as np from tqdm.auto import tqdm from pathlib import Path -from vall_e.config import cfg +from ..config import cfg def pad(num, zeroes): return str(num).zfill(zeroes+1) @@ -17,14 +22,26 @@ def process_items( items, stride=0 ): items = sorted( items ) return items if stride == 0 else [ item for i, item in enumerate( items ) if i % stride == 0 ] -def process_dataset( args ): +def process( + audio_backend="encodec", + input_audio="voices", + input_metadata="metadata", + output_dataset="training", + raise_exceptions=False, + stride=0, + slice="auto", + + device="cuda", + dtype="float16", + amp=False, + ): # encodec / vocos - if args.audio_backend in ["encodec", "vocos"]: + if audio_backend in ["encodec", "vocos"]: audio_extension = ".enc" cfg.sample_rate = 24_000 cfg.model.resp_levels = 8 - elif args.audio_backend == "dac": + elif audio_backend == "dac": audio_extension = ".dac" cfg.sample_rate = 44_100 cfg.model.resp_levels = 9 @@ -33,24 +50,18 @@ def process_dataset( args ): audio_extension = ".dec" cfg.model.resp_levels = 8 # ? else: - raise Exception(f"Unknown audio backend: {args.audio_backend}") + raise Exception(f"Unknown audio backend: {audio_backend}") # prepare from args - cfg.audio_backend = args.audio_backend # "encodec" - cfg.inference.weight_dtype = args.dtype # "bfloat16" - cfg.inference.amp = args.amp # False + cfg.audio_backend = audio_backend # "encodec" + cfg.inference.weight_dtype = dtype # "bfloat16" + cfg.inference.amp = amp # False # import after because we've overriden the config above - from vall_e.emb.g2p import encode as valle_phonemize - from vall_e.emb.qnt import encode as valle_quantize, _replace_file_extension + from .g2p import encode as phonemize + from .qnt import encode as quantize, _replace_file_extension - input_audio = args.input_audio # "voice"" - input_metadata = args.input_metadata # "metadata" - output_group = f"{args.output_group}-{'2' if cfg.sample_rate == 24_000 else '4'}{'8' if cfg.sample_rate == 48_000 else '4'}KHz-{cfg.audio_backend}" # "training" - device = args.device # "cuda" - raise_exceptions = args.raise_exceptions # False - stride = args.stride # 0 - slice = args.slice # "auto" + output_dataset = f"{output_dataset}/{'2' if cfg.sample_rate == 24_000 else '4'}{'8' if cfg.sample_rate == 48_000 else '4'}KHz-{cfg.audio_backend}" # "training" language_map = {} # k = group, v = language @@ -88,18 +99,18 @@ def process_dataset( args ): if only_speakers and speaker_id not in only_speakers: continue - os.makedirs(f'./{output_group}/{group_name}/{speaker_id}/', exist_ok=True) + os.makedirs(f'./{output_dataset}/{group_name}/{speaker_id}/', exist_ok=True) if speaker_id == "Noise": for filename in sorted(os.listdir(f'./{input_audio}/{group_name}/{speaker_id}/')): inpath = Path(f'./{input_audio}/{group_name}/{speaker_id}/{filename}') - outpath = Path(f'./{output_group}/{group_name}/{speaker_id}/{filename}') + outpath = Path(f'./{output_dataset}/{group_name}/{speaker_id}/{filename}') if _replace_file_extension(outpath, audio_extension).exists(): continue waveform, sample_rate = torchaudio.load(inpath) - qnt = valle_quantize(waveform, sr=sample_rate, device=device) + qnt = quantize(waveform, sr=sample_rate, device=device) if cfg.audio_backend == "dac": np.save(open(_replace_file_extension(outpath, audio_extension), "wb"), { @@ -158,7 +169,7 @@ def process_dataset( args ): language = language_map[group_name] if group_name in language_map else (metadata[filename]["language"] if "language" in metadata[filename] else "en") if len(metadata[filename]["segments"]) == 0 or not use_slices: - outpath = Path(f'./{output_group}/{group_name}/{speaker_id}/{fname}.{extension}') + outpath = Path(f'./{output_dataset}/{group_name}/{speaker_id}/{fname}.{extension}') text = metadata[filename]["text"] if len(text) == 0: @@ -185,7 +196,7 @@ def process_dataset( args ): id = pad(i, 4) i = i + 1 - outpath = Path(f'./{output_group}/{group_name}/{speaker_id}/{fname}_{id}.{extension}') + outpath = Path(f'./{output_dataset}/{group_name}/{speaker_id}/{fname}_{id}.{extension}') text = segment["text"] if len(text) == 0: @@ -223,8 +234,8 @@ def process_dataset( args ): try: outpath, text, language, waveform, sample_rate = job - phones = valle_phonemize(text, language=language) - qnt = valle_quantize(waveform, sr=sample_rate, device=device) + phones = phonemize(text, language=language) + qnt = quantize(waveform, sr=sample_rate, device=device) if cfg.audio_backend == "dac": @@ -273,8 +284,8 @@ def main(): parser.add_argument("--dtype", type=str, default="bfloat16") parser.add_argument("--amp", action="store_true") parser.add_argument("--input-audio", type=str, default="voices") - parser.add_argument("--input-metadata", type=str, default="metadata") - parser.add_argument("--output_group", type=str, default="training") + parser.add_argument("--input-metadata", type=str, default="training/metadata") + parser.add_argument("--output-dataset", type=str, default="training/dataset") parser.add_argument("--device", type=str, default="cuda") parser.add_argument("--raise-exceptions", action="store_true") parser.add_argument("--stride", type=int, default=0) @@ -282,7 +293,19 @@ def main(): args = parser.parse_args() - process_dataset( args ) + process( + audio_backend=args.audio_backend, + input_audio=args.input_audio, + input_metadata=args.input_metadata, + output_dataset=args.output_dataset, + raise_exceptions=args.raise_exceptions, + stride=args.stride, + slice=args.slice, + + device=args.device, + dtype=args.dtype, + amp=args.amp, + ) if __name__ == "__main__": main() \ No newline at end of file diff --git a/vall_e/emb/transcribe.py b/vall_e/emb/transcribe.py new file mode 100644 index 0000000..ab9aad7 --- /dev/null +++ b/vall_e/emb/transcribe.py @@ -0,0 +1,147 @@ +""" +# Handles transcribing audio provided through --input-audio +""" + +import os +import json +import argparse + +import torch +import torchaudio + +import whisperx + +from tqdm.auto import tqdm +from pathlib import Path + +def pad(num, zeroes): + return str(num).zfill(zeroes+1) + +def transcribe( + input_audio = "voices", + output_metadata = "training/metadata", + model_name = "large-v3", + + skip_existing = True, + diarize = False, + + batch_size = 16, + device = "cuda", + dtype = "float16", +): + # + model = whisperx.load_model(model_name, device, compute_type=dtype) + align_model, align_model_metadata, align_model_language = (None, None, None) + if diarize: + diarize_model = whisperx.DiarizationPipeline(device=device) + else: + diarize_model = None + + + for dataset_name in os.listdir(f'./{input_audio}/'): + if not os.path.isdir(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}'): + continue + + outpath = Path(f'./{output_metadata}/{dataset_name}/{speaker_id}/whisper.json') + + if outpath.exists(): + metadata = json.loads(open(outpath, 'r', encoding='utf-8').read()) + else: + os.makedirs(f'./{output_metadata}/{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 + + if ".json" in filename: + continue + + inpath = f'./{input_audio}/{dataset_name}/{speaker_id}/{filename}' + + if os.path.isdir(inpath): + continue + + metadata[filename] = { + "segments": [], + "language": "", + "text": "", + "start": 0, + "end": 0, + } + + audio = whisperx.load_audio(inpath) + result = model.transcribe(audio, batch_size=batch_size) + language = result["language"] + + """ + if language[:2] not in ["ja"]: + language = "en" + """ + + 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 + + if diarize_model is not None: + diarize_segments = diarize_model(audio) + result = whisperx.assign_word_speakers(diarize_segments, result) + + text = [] + start = 0 + end = 0 + for segment in result["segments"]: + text.append( segment["text"] ) + start = min( start, segment["start"] ) + end = max( end, segment["end"] ) + + metadata[filename]["text"] = " ".join(text).strip() + metadata[filename]["start"] = start + metadata[filename]["end"] = end + + open(outpath, 'w', encoding='utf-8').write(json.dumps(metadata)) + +def main(): + parser = argparse.ArgumentParser() + + parser.add_argument("--input-audio", type=str, default="voices") + parser.add_argument("--output-metadata", type=str, default="training/metadata") + + parser.add_argument("--model-name", type=str, default="large-v3") + parser.add_argument("--skip-existing", action="store_true") + parser.add_argument("--diarize", action="store_true") + parser.add_argument("--batch-size", type=int, default=16) + + parser.add_argument("--device", type=str, default="cuda") + parser.add_argument("--dtype", type=str, default="bfloat16") + parser.add_argument("--amp", action="store_true") + # parser.add_argument("--raise-exceptions", action="store_true") + + args = parser.parse_args() + + transcribe( + input_audio = args.input_audio, + output_metadata = args.output_metadata, + model_name = args.model_name, + + skip_existing = args.skip_existing, + diarize = args.diarize, + + batch_size = args.batch_size, + device = args.device, + dtype = args.dtype, + ) + +if __name__ == "__main__": + main() \ No newline at end of file