cleanup loop, save files while generating a batch in the event it crashes midway through

This commit is contained in:
mrq 2023-02-12 01:15:22 +00:00
parent 1b55730e67
commit ddd0c4ccf8

View File

@ -143,12 +143,10 @@ def generate(
volume_adjust = torchaudio.transforms.Vol(gain=args.output_volume, gain_type="amplitude") if args.output_volume != 1 else None volume_adjust = torchaudio.transforms.Vol(gain=args.output_volume, gain_type="amplitude") if args.output_volume != 1 else None
idx = 0 idx = 1
for i, file in enumerate(os.listdir(outdir)): for i, file in enumerate(os.listdir(outdir)):
if file[-5:] == ".json": if file[-5:] == ".json":
idx = idx + 1 idx = idx + 1
if idx:
idx = idx + 1
# reserve, if for whatever reason you manage to concurrently generate # reserve, if for whatever reason you manage to concurrently generate
with open(f'{outdir}/input_{idx}.json', 'w', encoding="utf-8") as f: with open(f'{outdir}/input_{idx}.json', 'w', encoding="utf-8") as f:
@ -180,24 +178,23 @@ def generate(
run_time = time.time()-start_time run_time = time.time()-start_time
print(f"Generating line took {run_time} seconds") print(f"Generating line took {run_time} seconds")
if isinstance(gen, list): if not isinstance(gen, list):
for j, g in enumerate(gen): gen = [gen]
name = get_name(line=line, candidate=j)
audio_cache[name] = { for j, g in enumerate(gen):
'audio': g, audio = g.squeeze(0).cpu()
'text': cut_text, name = get_name(line=line, candidate=j)
'time': run_time
}
else:
name = get_name(line=line)
audio_cache[name] = { audio_cache[name] = {
'audio': gen, 'audio': audio,
'text': cut_text, 'text': cut_text,
'time': run_time, 'time': run_time
} }
# save here in case some error happens mid-batch
torchaudio.save(f'{outdir}/{voice}_{name}.wav', audio, args.output_sample_rate)
for k in audio_cache: for k in audio_cache:
audio = audio_cache[k]['audio'].squeeze(0).cpu() audio = audio_cache[k]['audio']
if resampler is not None: if resampler is not None:
audio = resampler(audio) audio = resampler(audio)
if volume_adjust is not None: if volume_adjust is not None: