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

scripts/transcribe_dataset.py

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

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.audio_backend = audio_backend # "encodec"
-	cfg.inference.weight_dtype = args.dtype # "bfloat16"
+	cfg.inference.weight_dtype = dtype # "bfloat16"
-	cfg.inference.amp = args.amp # False
+	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 .g2p import encode as phonemize
-	from vall_e.emb.qnt import encode as valle_quantize, _replace_file_extension
+	from .qnt import encode as quantize, _replace_file_extension
 
-	input_audio = args.input_audio # "voice""
+	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"
-	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"
 
 	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)
+						phones = phonemize(text, language=language)
-						qnt = valle_quantize(waveform, sr=sample_rate, device=device)
+						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("--input-metadata", type=str, default="training/metadata")
-	parser.add_argument("--output_group", type=str, default="training")
+	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()

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