mrq/vall-e
5
9
Fork 0
Code Issues 8 Pull Requests Packages Projects Releases Wiki Activity

moved transcribe and process dataset scripts to vall_e/emb within the module itself, argparse-ified transcription script

Browse Source
This commit is contained in:
mrq 2024-08-05 19:40:50 -05:00
parent 7cdfa3dc0c
commit 597441e48b
3 changed files with 197 additions and 130 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

103
scripts/transcribe_dataset.py
View File

@ -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))

77
scripts/process_dataset.py → vall_e/emb/process.py
View File

@ -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()

147
vall_e/emb/transcribe.py Normal file
View File

@ -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()