From 7f54f09a0d247a68a783b12b6075167198a5fc3c Mon Sep 17 00:00:00 2001
From: James Betker <jbetker@gmail.com>
Date: Tue, 29 Mar 2022 13:59:59 -0600
Subject: [PATCH] Add sweeper script for finding optimal generation
 hyperparameters.

---
 sweep.py | 61 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 61 insertions(+)
 create mode 100644 sweep.py

diff --git a/sweep.py b/sweep.py
new file mode 100644
index 0000000..07f9dfc
--- /dev/null
+++ b/sweep.py
@@ -0,0 +1,61 @@
+import os
+from random import shuffle
+
+import torchaudio
+
+from api import TextToSpeech
+from utils.audio import load_audio
+
+
+def permutations(args):
+    res = []
+    k = next(iter(args.keys()))
+    vals = args[k]
+    del args[k]
+    if not args:
+        return [{k: v} for v in vals]
+    lower = permutations(args)
+    for v in vals:
+        for l in lower:
+            lc = l.copy()
+            lc[k] = v
+            res.append(lc)
+    return res
+
+
+if __name__ == '__main__':
+    fname = 'Y:\\libritts\\test-clean\\transcribed-brief-w2v.tsv'
+    outpath_base = 'D:\\tmp\\tortoise-tts-eval\\std_sweep_diffusion'
+    outpath_real = 'D:\\tmp\\tortoise-tts-eval\\real'
+
+    arg_ranges = {
+        'diffusion_temperature': [.5, .7, 1],
+        'cond_free_k': [.5, 1, 2],
+    }
+    cfgs = permutations(arg_ranges)
+    shuffle(cfgs)
+
+    for cfg in cfgs:
+        outpath = os.path.join(outpath_base, f'{cfg["cond_free_k"]}_{cfg["diffusion_temperature"]}')
+        os.makedirs(outpath, exist_ok=True)
+        os.makedirs(outpath_real, exist_ok=True)
+        with open(fname, 'r', encoding='utf-8') as f:
+            lines = [l.strip().split('\t') for l in f.readlines()]
+
+        recorder = open(os.path.join(outpath, 'transcript.tsv'), 'w', encoding='utf-8')
+        tts = TextToSpeech()
+        for e, line in enumerate(lines):
+            transcript = line[0]
+            if len(transcript) > 120:
+                continue  # We need to support this, but cannot yet.
+            path = os.path.join(os.path.dirname(fname), line[1])
+            cond_audio = load_audio(path, 22050)
+            torchaudio.save(os.path.join(outpath_real, os.path.basename(line[1])), cond_audio, 22050)
+            sample = tts.tts(transcript, [cond_audio, cond_audio], num_autoregressive_samples=256, k=1, diffusion_iterations=200, cond_free=False,
+                             repetition_penalty=1.5, length_penalty=2, temperature=.9, top_p=.9)
+            down = torchaudio.functional.resample(sample, 24000, 22050)
+            fout_path = os.path.join(outpath, os.path.basename(line[1]))
+            torchaudio.save(fout_path, down.squeeze(0), 22050)
+            recorder.write(f'{transcript}\t{fout_path}\n')
+            recorder.flush()
+        recorder.close()
\ No newline at end of file