feat: better ui

app.py

@@ -1,459 +1,53 @@
-from typing import Tuple
 from prelude import prelude
 
 prelude()
 
-import os
-import traceback
-import numpy as np
-from sklearn.cluster import MiniBatchKMeans
-from random import shuffle
 import gradio as gr
-import zipfile
-import tempfile
-import shutil
-import faiss
-from glob import glob
-from infer.modules.train.preprocess import PreProcess
-from infer.modules.train.extract.extract_f0_rmvpe import FeatureInput
-from infer.modules.train.extract_feature_print import HubertFeatureExtractor
-from infer.modules.train.train import train
-from infer.lib.train.process_ckpt import extract_small_model
-from infer.modules.vc.modules import VC
-from configs.config import Config
-import demucs.separate
-import soundfile as sf
-from zero import zero
-from model import device
+from app.setup import SetupTab
+from app.extract import FeatureExtractionTab
+from app.train import TrainTab
+from app.export import ExportTab
+from app.infer import InferenceTab
 
 
-def extract_audio_files(zip_file: str, target_dir: str) -> list[str]:
-    with zipfile.ZipFile(zip_file, "r") as zip_ref:
-        zip_ref.extractall(target_dir)
-
-    audio_files = [
-        os.path.join(target_dir, f)
-        for f in os.listdir(target_dir)
-        if f.endswith((".wav", ".mp3", ".ogg"))
-    ]
-    if not audio_files:
-        raise gr.Error("No audio files found at the top level of the zip file")
-
-    return audio_files
-
-
-def preprocess(zip_file: str) -> str:
-    temp_dir = tempfile.mkdtemp()
-    print(f"Using exp dir: {temp_dir}")
-
-    data_dir = os.path.join(temp_dir, "_data")
-    os.makedirs(data_dir)
-    audio_files = extract_audio_files(zip_file, data_dir)
-
-    pp = PreProcess(40000, temp_dir, 3.0, False)
-    pp.pipeline_mp_inp_dir(data_dir, 4)
-
-    pp.logfile.seek(0)
-    log = pp.logfile.read()
-
-    return temp_dir, f"Preprocessed {len(audio_files)} audio files.\n{log}"
-
-
-@zero(duration=300)
-def extract_features(exp_dir: str) -> str:
-    err = None
-    fi = FeatureInput(exp_dir)
-    try:
-        fi.run()
-    except Exception as e:
-        err = e
-
-    fi.logfile.seek(0)
-    log = fi.logfile.read()
-
-    if err:
-        log = f"Error: {err}\n{log}"
-        return log
-
-    hfe = HubertFeatureExtractor(exp_dir)
-    try:
-        hfe.run()
-    except Exception as e:
-        err = e
-
-    hfe.logfile.seek(0)
-    log += hfe.logfile.read()
-
-    if err:
-        log = f"Error: {err}\n{log}"
-
-    return log
-
-
-def write_filelist(exp_dir: str) -> None:
-    if_f0_3 = True
-    spk_id5 = 0
-    gt_wavs_dir = "%s/0_gt_wavs" % (exp_dir)
-    feature_dir = "%s/3_feature768" % (exp_dir)
-
-    if if_f0_3:
-        f0_dir = "%s/2a_f0" % (exp_dir)
-        f0nsf_dir = "%s/2b-f0nsf" % (exp_dir)
-        names = (
-            set([name.split(".")[0] for name in os.listdir(gt_wavs_dir)])
-            & set([name.split(".")[0] for name in os.listdir(feature_dir)])
-            & set([name.split(".")[0] for name in os.listdir(f0_dir)])
-            & set([name.split(".")[0] for name in os.listdir(f0nsf_dir)])
-        )
-    else:
-        names = set([name.split(".")[0] for name in os.listdir(gt_wavs_dir)]) & set(
-            [name.split(".")[0] for name in os.listdir(feature_dir)]
-        )
-    opt = []
-    for name in names:
-        if if_f0_3:
-            opt.append(
-                "%s/%s.wav|%s/%s.npy|%s/%s.wav.npy|%s/%s.wav.npy|%s"
-                % (
-                    gt_wavs_dir.replace("\\", "\\\\"),
-                    name,
-                    feature_dir.replace("\\", "\\\\"),
-                    name,
-                    f0_dir.replace("\\", "\\\\"),
-                    name,
-                    f0nsf_dir.replace("\\", "\\\\"),
-                    name,
-                    spk_id5,
-                )
-            )
-        else:
-            opt.append(
-                "%s/%s.wav|%s/%s.npy|%s"
-                % (
-                    gt_wavs_dir.replace("\\", "\\\\"),
-                    name,
-                    feature_dir.replace("\\", "\\\\"),
-                    name,
-                    spk_id5,
-                )
-            )
-    fea_dim = 768
-
-    now_dir = os.getcwd()
-    sr2 = "40k"
-    if if_f0_3:
-        for _ in range(2):
-            opt.append(
-                "%s/logs/mute/0_gt_wavs/mute%s.wav|%s/logs/mute/3_feature%s/mute.npy|%s/logs/mute/2a_f0/mute.wav.npy|%s/logs/mute/2b-f0nsf/mute.wav.npy|%s"
-                % (now_dir, sr2, now_dir, fea_dim, now_dir, now_dir, spk_id5)
-            )
-    else:
-        for _ in range(2):
-            opt.append(
-                "%s/logs/mute/0_gt_wavs/mute%s.wav|%s/logs/mute/3_feature%s/mute.npy|%s"
-                % (now_dir, sr2, now_dir, fea_dim, spk_id5)
-            )
-    shuffle(opt)
-    with open("%s/filelist.txt" % exp_dir, "w") as f:
-        f.write("\n".join(opt))
-
-
-@zero(duration=300)
-def train_model(exp_dir: str) -> str:
-    shutil.copy("config.json", exp_dir)
-    write_filelist(exp_dir)
-    train(exp_dir)
-
-    models = glob(f"{exp_dir}/G_*.pth")
-    print(models)
-    if not models:
-        raise gr.Error("No model found")
-
-    latest_model = max(models, key=os.path.getctime)
-    return latest_model
-
-
-def download_weight(exp_dir: str) -> str:
-    models = glob(f"{exp_dir}/G_*.pth")
-    if not models:
-        raise gr.Error("No model found")
-
-    latest_model = max(models, key=os.path.getctime)
-    print(f"Latest model: {latest_model}")
-
-    name = os.path.basename(exp_dir)
-    out = os.path.join(exp_dir, f"{name}.pth")
-    extract_small_model(
-        latest_model, out, "40k", True, "Model trained by ZeroGPU.", "v2"
-    )
-
-    return out
-
-
-def train_index(exp_dir: str) -> str:
-    feature_dir = "%s/3_feature768" % (exp_dir)
-    if not os.path.exists(feature_dir):
-        raise gr.Error("Please extract features first.")
-    listdir_res = list(os.listdir(feature_dir))
-    if len(listdir_res) == 0:
-        raise gr.Error("Please extract features first.")
-    npys = []
-    for name in sorted(listdir_res):
-        phone = np.load("%s/%s" % (feature_dir, name))
-        npys.append(phone)
-    big_npy = np.concatenate(npys, 0)
-    big_npy_idx = np.arange(big_npy.shape[0])
-    np.random.shuffle(big_npy_idx)
-    big_npy = big_npy[big_npy_idx]
-    if big_npy.shape[0] > 2e5:
-        print("Trying doing kmeans %s shape to 10k centers." % big_npy.shape[0])
-        try:
-            big_npy = (
-                MiniBatchKMeans(
-                    n_clusters=10000,
-                    verbose=True,
-                    batch_size=256 * 8,
-                    compute_labels=False,
-                    init="random",
-                )
-                .fit(big_npy)
-                .cluster_centers_
-            )
-        except:
-            info = traceback.format_exc()
-            print(info)
-            raise gr.Error(info)
-
-    np.save("%s/total_fea.npy" % exp_dir, big_npy)
-    n_ivf = min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39)
-    print("%s,%s" % (big_npy.shape, n_ivf))
-    index = faiss.index_factory(768, "IVF%s,Flat" % n_ivf)
-    # index = faiss.index_factory(256if version19=="v1"else 768, "IVF%s,PQ128x4fs,RFlat"%n_ivf)
-    print("training")
-    index_ivf = faiss.extract_index_ivf(index)  #
-    index_ivf.nprobe = 1
-    index.train(big_npy)
-    faiss.write_index(
-        index,
-        "%s/trained_IVF%s_Flat_nprobe_%s.index" % (exp_dir, n_ivf, index_ivf.nprobe),
-    )
-    print("adding")
-    batch_size_add = 8192
-    for i in range(0, big_npy.shape[0], batch_size_add):
-        index.add(big_npy[i : i + batch_size_add])
-    faiss.write_index(
-        index,
-        "%s/added_IVF%s_Flat_nprobe_%s.index" % (exp_dir, n_ivf, index_ivf.nprobe),
-    )
-    print("built added_IVF%s_Flat_nprobe_%s.index" % (n_ivf, index_ivf.nprobe))
-
-    return "%s/added_IVF%s_Flat_nprobe_%s.index" % (exp_dir, n_ivf, index_ivf.nprobe)
-
-
-def download_expdir(exp_dir: str) -> str:
-    shutil.make_archive(exp_dir, "zip", exp_dir)
-    return f"{exp_dir}.zip"
-
-
-def restore_expdir(zip: str) -> str:
-    exp_dir = tempfile.mkdtemp()
-    shutil.unpack_archive(zip, exp_dir)
-    return exp_dir
-
-
-@zero(duration=120)
-def infer(exp_dir: str, original_audio: str, f0add: int) -> Tuple[int, np.ndarray]:
-    name = os.path.basename(exp_dir)
-    model = os.path.join(exp_dir, f"{name}.pth")
-    if not os.path.exists(model):
-        raise gr.Error("Model not found")
-
-    index = glob(f"{exp_dir}/added_*.index")
-    if not index:
-        raise gr.Error("Index not found")
-
-    base = os.path.basename(original_audio)
-    base = os.path.splitext(base)[0]
-    demucs.separate.main(
-        ["--two-stems", "vocals", "-d", str(device), "-n", "htdemucs", original_audio]
-    )
-    out = os.path.join("separated", "htdemucs", base, "vocals.wav")
-
-    cfg = Config()
-    vc = VC(cfg)
-    vc.get_vc(model)
-    _, wav_opt = vc.vc_single(
-        0,
-        out,
-        f0add,
-        None,
-        "rmvpe",
-        index,
-        None,
-        0.5,
-        3,
-        0,
-        1,
-        0.33,
-    )
-
-    sr = wav_opt[0]
-    data = wav_opt[1]
-
-    return sr, data
-
-
-def merge(exp_dir: str, original_audio: str, vocal: Tuple[int, np.ndarray]) -> str:
-    base = os.path.basename(original_audio)
-    base = os.path.splitext(base)[0]
-    music = os.path.join("separated", "htdemucs", base, "no-vocals.wav")
-
-    tmp = os.path.join(exp_dir, "tmp.wav")
-    sf.write(tmp, vocal[1], vocal[0])
-
-    os.system(
-        f"ffmpeg -i {music} -i {tmp} -filter_complex '[1]volume=2[a];[0][a]amix=inputs=2:duration=first:dropout_transition=2' {tmp}.merged.mp3"
+with gr.Blocks() as app:
+    gr.Markdown("# ZeroRVC")
+    gr.Markdown(
+        "Run Retrieval-based Voice Conversion training and inference on HuggingFace ZeroGPU."
     )
 
-    return f"{tmp}.merged.mp3"
-
-
-with gr.Blocks() as app:
-    # allow user to manually select the experiment directory
     exp_dir = gr.Textbox(
-        label="Experiment directory (don't touch it unless you know what you are doing)",
+        label="Experiment directory",
         visible=True,
-        interactive=True,
+        interactive=False,
     )
 
-    with gr.Tabs():
-        with gr.Tab(label="New / Restore"):
-            with gr.Row():
-                with gr.Column():
-                    zip_file = gr.File(
-                        label="Upload a zip file containing audio files for training",
-                        file_types=["zip"],
-                    )
-                    preprocess_output = gr.Textbox(
-                        label="Preprocessing output", lines=5
-                    )
-
-                preprocess_btn = gr.Button(
-                    value="Start New Experiment", variant="primary"
-                )
+    setup = SetupTab()
+    feature_extraction = FeatureExtractionTab()
+    training = TrainTab()
+    export = ExportTab()
+    inferencing = InferenceTab()
 
-            with gr.Row():
-                restore_zip_file = gr.File(
-                    label="Upload the experiment directory zip file",
-                    file_types=["zip"],
-                )
-                restore_btn = gr.Button(value="Restore Experiment", variant="primary")
+    with gr.Tabs():
+        with gr.Tab(label="Setup"):
+            setup.ui()
 
-        with gr.Tab(label="Extract features"):
-            with gr.Row():
-                extract_features_btn = gr.Button(
-                    value="Extract features", variant="primary"
-                )
-            with gr.Row():
-                extract_features_output = gr.Textbox(
-                    label="Feature extraction output", lines=10
-                )
+        with gr.Tab(label="Feature Extraction"):
+            feature_extraction.ui()
 
-        with gr.Tab(label="Train"):
-            with gr.Row():
-                train_btn = gr.Button(value="Train", variant="primary")
-                latest_model = gr.File(label="Latest checkpoint")
-            with gr.Row():
-                train_index_btn = gr.Button(value="Train index", variant="primary")
-                trained_index = gr.File(label="Trained index")
+        with gr.Tab(label="Training"):
+            training.ui()
 
         with gr.Tab(label="Download"):
-            with gr.Row():
-                download_weight_btn = gr.Button(
-                    value="Download latest model", variant="primary"
-                )
-                download_weight_output = gr.File(label="Download latest model")
-
-            with gr.Row():
-                download_expdir_btn = gr.Button(
-                    value="Download experiment directory", variant="primary"
-                )
-                download_expdir_output = gr.File(label="Download experiment directory")
+            export.ui()
 
         with gr.Tab(label="Inference"):
-            with gr.Row():
-                original_audio = gr.Audio(
-                    label="Upload original audio",
-                    type="filepath",
-                    show_download_button=True,
-                )
-                f0add = gr.Slider(
-                    label="F0 add",
-                    minimum=-16,
-                    maximum=16,
-                    step=1,
-                    value=0,
-                )
-                infer_btn = gr.Button(value="Infer", variant="primary")
-            with gr.Row():
-                infer_output = gr.Audio(label="Inferred audio")
-            with gr.Row():
-                merge_output = gr.Audio(label="Merged audio")
-
-    preprocess_btn.click(
-        fn=preprocess,
-        inputs=[zip_file],
-        outputs=[exp_dir, preprocess_output],
-    )
-
-    extract_features_btn.click(
-        fn=extract_features,
-        inputs=[exp_dir],
-        outputs=[extract_features_output],
-    )
-
-    train_btn.click(
-        fn=train_model,
-        inputs=[exp_dir],
-        outputs=[latest_model],
-    ).success(
-        fn=train_model,
-        inputs=[exp_dir],
-        outputs=[latest_model],
-    )
+            inferencing.ui()
 
-    train_index_btn.click(
-        fn=train_index,
-        inputs=[exp_dir],
-        outputs=[trained_index],
-    )
-
-    download_weight_btn.click(
-        fn=download_weight,
-        inputs=[exp_dir],
-        outputs=[download_weight_output],
-    )
-
-    download_expdir_btn.click(
-        fn=download_expdir,
-        inputs=[exp_dir],
-        outputs=[download_expdir_output],
-    )
-
-    restore_btn.click(
-        fn=restore_expdir,
-        inputs=[restore_zip_file],
-        outputs=[exp_dir],
-    )
-
-    infer_btn.click(
-        fn=infer,
-        inputs=[exp_dir, original_audio, f0add],
-        outputs=[infer_output],
-    ).success(
-        fn=merge,
-        inputs=[exp_dir, original_audio, infer_output],
-        outputs=[merge_output],
-    )
+    setup.build(exp_dir)
+    feature_extraction.build(exp_dir)
+    training.build(exp_dir)
+    export.build(exp_dir)
+    inferencing.build(exp_dir)
 
     app.launch()

app/export.py

@@ -0,0 +1,79 @@
+from glob import glob
+import os
+import shutil
+import gradio as gr
+from infer.lib.train.process_ckpt import extract_small_model
+
+
+def download_weight(exp_dir: str) -> str:
+    models = glob(f"{exp_dir}/G_*.pth")
+    if not models:
+        raise gr.Error("No model found")
+
+    latest_model = max(models, key=os.path.getctime)
+    print(f"Latest model: {latest_model}")
+
+    name = os.path.basename(exp_dir)
+    out = os.path.join(exp_dir, f"{name}.pth")
+    extract_small_model(
+        latest_model, out, "40k", True, "Model trained by ZeroGPU.", "v2"
+    )
+
+    return out
+
+
+def download_expdir(exp_dir: str) -> str:
+    shutil.make_archive(exp_dir, "zip", exp_dir)
+    return f"{exp_dir}.zip"
+
+
+def remove_expdir(exp_dir: str) -> str:
+    shutil.rmtree(exp_dir)
+    return ""
+
+
+class ExportTab:
+    def __init__(self):
+        pass
+
+    def ui(self):
+        gr.Markdown("# Download Model or Experiment Directory")
+        gr.Markdown(
+            "You can download the latest model or the entire experiment directory here."
+        )
+
+        with gr.Row():
+            self.download_weight_btn = gr.Button(
+                value="Latest model (for inferencing)", variant="primary"
+            )
+            self.download_weight_output = gr.File(label="Prune latest model")
+
+        with gr.Row():
+            self.download_expdir_btn = gr.Button(
+                value="Download experiment directory", variant="primary"
+            )
+            self.download_expdir_output = gr.File(label="Archive experiment directory")
+
+        with gr.Row():
+            self.remove_expdir_btn = gr.Button(
+                value="REMOVE experiment directory", variant="stop"
+            )
+
+    def build(self, exp_dir: gr.Textbox):
+        self.download_weight_btn.click(
+            fn=download_weight,
+            inputs=[exp_dir],
+            outputs=[self.download_weight_output],
+        )
+
+        self.download_expdir_btn.click(
+            fn=download_expdir,
+            inputs=[exp_dir],
+            outputs=[self.download_expdir_output],
+        )
+
+        self.remove_expdir_btn.click(
+            fn=remove_expdir,
+            inputs=[exp_dir],
+            outputs=[exp_dir],
+        )

app/extract.py

@@ -0,0 +1,64 @@
+import gradio as gr
+from infer.modules.train.extract.extract_f0_rmvpe import FeatureInput
+from infer.modules.train.extract_feature_print import HubertFeatureExtractor
+from zero import zero
+
+
+@zero(duration=300)
+def extract_features(exp_dir: str) -> str:
+    err = None
+    fi = FeatureInput(exp_dir)
+    try:
+        fi.run()
+    except Exception as e:
+        err = e
+
+    fi.logfile.seek(0)
+    log = fi.logfile.read()
+
+    if err:
+        log = f"Error: {err}\n{log}"
+        return log
+
+    hfe = HubertFeatureExtractor(exp_dir)
+    try:
+        hfe.run()
+    except Exception as e:
+        err = e
+
+    hfe.logfile.seek(0)
+    log += hfe.logfile.read()
+
+    if err:
+        log = f"Error: {err}\n{log}"
+
+    return log
+
+
+class FeatureExtractionTab:
+    def __init__(self):
+        pass
+
+    def ui(self):
+        gr.Markdown("# Feature Extraction")
+        gr.Markdown(
+            "Before training, you need to extract features from the audio files. "
+            "This process may take a while, depending on the number of audio files. "
+            "Under the hood, this process extracts speech features using HuBERT and extracts F0 by RMVPE."
+        )
+
+        with gr.Row():
+            self.extract_features_btn = gr.Button(
+                value="Extract features", variant="primary"
+            )
+        with gr.Row():
+            self.extract_features_log = gr.Textbox(
+                label="Feature extraction log", lines=10
+            )
+
+    def build(self, exp_dir: gr.Textbox):
+        self.extract_features_btn.click(
+            fn=extract_features,
+            inputs=[exp_dir],
+            outputs=[self.extract_features_log],
+        )

app/infer.py

@@ -0,0 +1,110 @@
+from glob import glob
+import os
+from typing import Tuple
+import demucs
+import gradio as gr
+import numpy as np
+import soundfile as sf
+from configs.config import Config
+from infer.modules.vc.modules import VC
+from zero import zero
+from model import device
+
+
+@zero(duration=120)
+def infer(exp_dir: str, original_audio: str, f0add: int) -> Tuple[int, np.ndarray]:
+    name = os.path.basename(exp_dir)
+    model = os.path.join(exp_dir, f"{name}.pth")
+    if not os.path.exists(model):
+        raise gr.Error("Model not found")
+
+    index = glob(f"{exp_dir}/added_*.index")
+    if not index:
+        raise gr.Error("Index not found")
+
+    base = os.path.basename(original_audio)
+    base = os.path.splitext(base)[0]
+    demucs.separate.main(
+        ["--two-stems", "vocals", "-d", str(device), "-n", "htdemucs", original_audio]
+    )
+    out = os.path.join("separated", "htdemucs", base, "vocals.wav")
+
+    cfg = Config()
+    vc = VC(cfg)
+    vc.get_vc(model)
+    _, wav_opt = vc.vc_single(
+        0,
+        out,
+        f0add,
+        None,
+        "rmvpe",
+        index,
+        None,
+        0.5,
+        3,
+        0,
+        1,
+        0.33,
+    )
+
+    sr = wav_opt[0]
+    data = wav_opt[1]
+
+    return sr, data
+
+
+def merge(exp_dir: str, original_audio: str, vocal: Tuple[int, np.ndarray]) -> str:
+    base = os.path.basename(original_audio)
+    base = os.path.splitext(base)[0]
+    music = os.path.join("separated", "htdemucs", base, "no-vocals.wav")
+
+    tmp = os.path.join(exp_dir, "tmp.wav")
+    sf.write(tmp, vocal[1], vocal[0])
+
+    os.system(
+        f"ffmpeg -i {music} -i {tmp} -filter_complex '[1]volume=2[a];[0][a]amix=inputs=2:duration=first:dropout_transition=2' {tmp}.merged.mp3"
+    )
+
+    return f"{tmp}.merged.mp3"
+
+
+class InferenceTab:
+    def __init__(self):
+        pass
+
+    def ui(self):
+        gr.Markdown("# Inference")
+        gr.Markdown(
+            "After trained model is pruned, you can use it to infer on new music. \n"
+            "Upload the original audio and adjust the F0 add value to generate the inferred audio."
+        )
+
+        with gr.Row():
+            self.original_audio = gr.Audio(
+                label="Upload original audio",
+                type="filepath",
+                show_download_button=True,
+            )
+            self.f0add = gr.Slider(
+                label="F0 add",
+                minimum=-16,
+                maximum=16,
+                step=1,
+                value=0,
+            )
+            self.infer_btn = gr.Button(value="Infer", variant="primary")
+        with gr.Row():
+            self.infer_output = gr.Audio(label="Inferred audio")
+        with gr.Row():
+            self.merge_output = gr.Audio(label="Merged audio")
+
+    def build(self, exp_dir: gr.Textbox):
+        self.infer_btn.click(
+            fn=infer,
+            inputs=[exp_dir, self.original_audio, self.f0add],
+            outputs=[self.infer_output],
+        ).success(
+            fn=merge,
+            inputs=[exp_dir, self.original_audio, self.infer_output],
+            outputs=[self.merge_output],
+        )

app/setup.py

@@ -0,0 +1,110 @@
+import os
+import shutil
+import gradio as gr
+import zipfile
+import tempfile
+from infer.modules.train.preprocess import PreProcess
+from typing import Tuple
+
+
+def extract_audio_files(zip_file: str, target_dir: str) -> list[str]:
+    with zipfile.ZipFile(zip_file, "r") as zip_ref:
+        zip_ref.extractall(target_dir)
+
+    audio_files = [
+        os.path.join(target_dir, f)
+        for f in os.listdir(target_dir)
+        if f.endswith((".wav", ".mp3", ".ogg"))
+    ]
+    if not audio_files:
+        raise gr.Error("No audio files found at the top level of the zip file")
+
+    return audio_files
+
+
+def create_new_expdir(zip_file: str) -> Tuple[str, str]:
+    temp_dir = tempfile.mkdtemp()
+    print(f"Using exp dir: {temp_dir}")
+
+    data_dir = os.path.join(temp_dir, "_data")
+    os.makedirs(data_dir)
+    audio_files = extract_audio_files(zip_file, data_dir)
+
+    pp = PreProcess(40000, temp_dir, 3.0, False)
+    pp.pipeline_mp_inp_dir(data_dir, 4)
+
+    pp.logfile.seek(0)
+    log = pp.logfile.read()
+
+    return temp_dir, f"Preprocessed {len(audio_files)} audio files.\n{log}"
+
+
+def restore_expdir(zip: str) -> str:
+    exp_dir = tempfile.mkdtemp()
+    shutil.unpack_archive(zip, exp_dir)
+    return exp_dir
+
+
+def set_dir(dir_val: str) -> str:
+    if not dir_val.startswith("/tmp/"):
+        dir_val = os.path.join("/tmp", dir_val)
+    if not os.path.isdir(dir_val):
+        raise gr.Error("Directory does not exist")
+
+    return dir_val
+
+
+class SetupTab:
+    def __init__(self):
+        pass
+
+    def ui(self):
+        gr.Markdown("# Setup Experiment")
+        gr.Markdown(
+            "You can upload a zip file containing audio files to start a new experiment, or upload an experiment directory zip file to restore an existing experiment."
+        )
+
+        with gr.Row():
+            with gr.Column():
+                self.zip_file = gr.File(
+                    label="Upload a zip file containing audio files for training",
+                    file_types=["zip"],
+                )
+                self.preprocess_log = gr.Textbox(label="Log", lines=5)
+
+            self.preprocess_btn = gr.Button(
+                value="Start New Experiment", variant="primary"
+            )
+
+        with gr.Row():
+            self.restore_zip_file = gr.File(
+                label="Upload the experiment directory zip file",
+                file_types=["zip"],
+            )
+            self.restore_btn = gr.Button(value="Restore Experiment", variant="primary")
+
+        with gr.Row():
+            self.dir_val = gr.Textbox(
+                label="Manually set the experiment directory (don't touch it unless you know what you are doing)",
+                placeholder="/tmp/...",
+            )
+            self.set_dir_btn = gr.Button(value="Set Directory")
+
+    def build(self, exp_dir: gr.Textbox):
+        self.preprocess_btn.click(
+            fn=create_new_expdir,
+            inputs=[self.zip_file],
+            outputs=[exp_dir, self.preprocess_log],
+        )
+
+        self.restore_btn.click(
+            fn=restore_expdir,
+            inputs=[self.restore_zip_file],
+            outputs=[exp_dir],
+        )
+
+        self.set_dir_btn.click(
+            fn=set_dir,
+            inputs=[self.dir_val],
+            outputs=[exp_dir],
+        )

app/train.py

@@ -0,0 +1,191 @@
+import os
+import shutil
+import traceback
+import faiss
+import gradio as gr
+import numpy as np
+from sklearn.cluster import MiniBatchKMeans
+from random import shuffle
+from glob import glob
+from infer.modules.train.train import train
+from zero import zero
+
+
+def write_filelist(exp_dir: str) -> None:
+    if_f0_3 = True
+    spk_id5 = 0
+    gt_wavs_dir = "%s/0_gt_wavs" % (exp_dir)
+    feature_dir = "%s/3_feature768" % (exp_dir)
+
+    if if_f0_3:
+        f0_dir = "%s/2a_f0" % (exp_dir)
+        f0nsf_dir = "%s/2b-f0nsf" % (exp_dir)
+        names = (
+            set([name.split(".")[0] for name in os.listdir(gt_wavs_dir)])
+            & set([name.split(".")[0] for name in os.listdir(feature_dir)])
+            & set([name.split(".")[0] for name in os.listdir(f0_dir)])
+            & set([name.split(".")[0] for name in os.listdir(f0nsf_dir)])
+        )
+    else:
+        names = set([name.split(".")[0] for name in os.listdir(gt_wavs_dir)]) & set(
+            [name.split(".")[0] for name in os.listdir(feature_dir)]
+        )
+    opt = []
+    for name in names:
+        if if_f0_3:
+            opt.append(
+                "%s/%s.wav|%s/%s.npy|%s/%s.wav.npy|%s/%s.wav.npy|%s"
+                % (
+                    gt_wavs_dir.replace("\\", "\\\\"),
+                    name,
+                    feature_dir.replace("\\", "\\\\"),
+                    name,
+                    f0_dir.replace("\\", "\\\\"),
+                    name,
+                    f0nsf_dir.replace("\\", "\\\\"),
+                    name,
+                    spk_id5,
+                )
+            )
+        else:
+            opt.append(
+                "%s/%s.wav|%s/%s.npy|%s"
+                % (
+                    gt_wavs_dir.replace("\\", "\\\\"),
+                    name,
+                    feature_dir.replace("\\", "\\\\"),
+                    name,
+                    spk_id5,
+                )
+            )
+    fea_dim = 768
+
+    now_dir = os.getcwd()
+    sr2 = "40k"
+    if if_f0_3:
+        for _ in range(2):
+            opt.append(
+                "%s/logs/mute/0_gt_wavs/mute%s.wav|%s/logs/mute/3_feature%s/mute.npy|%s/logs/mute/2a_f0/mute.wav.npy|%s/logs/mute/2b-f0nsf/mute.wav.npy|%s"
+                % (now_dir, sr2, now_dir, fea_dim, now_dir, now_dir, spk_id5)
+            )
+    else:
+        for _ in range(2):
+            opt.append(
+                "%s/logs/mute/0_gt_wavs/mute%s.wav|%s/logs/mute/3_feature%s/mute.npy|%s"
+                % (now_dir, sr2, now_dir, fea_dim, spk_id5)
+            )
+    shuffle(opt)
+    with open("%s/filelist.txt" % exp_dir, "w") as f:
+        f.write("\n".join(opt))
+
+
+@zero(duration=300)
+def train_model(exp_dir: str) -> str:
+    shutil.copy("config.json", exp_dir)
+    write_filelist(exp_dir)
+    train(exp_dir)
+
+    models = glob(f"{exp_dir}/G_*.pth")
+    print(models)
+    if not models:
+        raise gr.Error("No model found")
+
+    latest_model = max(models, key=os.path.getctime)
+    return latest_model
+
+
+def train_index(exp_dir: str) -> str:
+    feature_dir = "%s/3_feature768" % (exp_dir)
+    if not os.path.exists(feature_dir):
+        raise gr.Error("Please extract features first.")
+    listdir_res = list(os.listdir(feature_dir))
+    if len(listdir_res) == 0:
+        raise gr.Error("Please extract features first.")
+    npys = []
+    for name in sorted(listdir_res):
+        phone = np.load("%s/%s" % (feature_dir, name))
+        npys.append(phone)
+    big_npy = np.concatenate(npys, 0)
+    big_npy_idx = np.arange(big_npy.shape[0])
+    np.random.shuffle(big_npy_idx)
+    big_npy = big_npy[big_npy_idx]
+    if big_npy.shape[0] > 2e5:
+        print("Trying doing kmeans %s shape to 10k centers." % big_npy.shape[0])
+        try:
+            big_npy = (
+                MiniBatchKMeans(
+                    n_clusters=10000,
+                    verbose=True,
+                    batch_size=256 * 8,
+                    compute_labels=False,
+                    init="random",
+                )
+                .fit(big_npy)
+                .cluster_centers_
+            )
+        except:
+            info = traceback.format_exc()
+            print(info)
+            raise gr.Error(info)
+
+    np.save("%s/total_fea.npy" % exp_dir, big_npy)
+    n_ivf = min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39)
+    print("%s,%s" % (big_npy.shape, n_ivf))
+    index = faiss.index_factory(768, "IVF%s,Flat" % n_ivf)
+    # index = faiss.index_factory(256if version19=="v1"else 768, "IVF%s,PQ128x4fs,RFlat"%n_ivf)
+    print("training")
+    index_ivf = faiss.extract_index_ivf(index)  #
+    index_ivf.nprobe = 1
+    index.train(big_npy)
+    faiss.write_index(
+        index,
+        "%s/trained_IVF%s_Flat_nprobe_%s.index" % (exp_dir, n_ivf, index_ivf.nprobe),
+    )
+    print("adding")
+    batch_size_add = 8192
+    for i in range(0, big_npy.shape[0], batch_size_add):
+        index.add(big_npy[i : i + batch_size_add])
+    faiss.write_index(
+        index,
+        "%s/added_IVF%s_Flat_nprobe_%s.index" % (exp_dir, n_ivf, index_ivf.nprobe),
+    )
+    print("built added_IVF%s_Flat_nprobe_%s.index" % (n_ivf, index_ivf.nprobe))
+
+    return "%s/added_IVF%s_Flat_nprobe_%s.index" % (exp_dir, n_ivf, index_ivf.nprobe)
+
+
+class TrainTab:
+    def __init__(self):
+        pass
+
+    def ui(self):
+        gr.Markdown("# Training")
+        gr.Markdown(
+            "You can start training the model by clicking the button below. "
+            "Each time you click the button, the model will train for 20 epochs, which takes about 10 minutes on ZeroGPU (A100). "
+            "Tha latest *training checkpoint* will be avaible below."
+        )
+
+        with gr.Row():
+            self.train_btn = gr.Button(value="Train", variant="primary")
+            self.latest_checkpoint = gr.File(label="Latest checkpoint")
+        with gr.Row():
+            self.train_index_btn = gr.Button(value="Train index", variant="primary")
+            self.trained_index = gr.File(label="Trained index")
+
+    def build(self, exp_dir: gr.Textbox):
+        self.train_btn.click(
+            fn=train_model,
+            inputs=[exp_dir],
+            outputs=[self.latest_checkpoint],
+        ).success(
+            fn=train_model,
+            inputs=[exp_dir],
+            outputs=[self.latest_checkpoint],
+        )
+
+        self.train_index_btn.click(
+            fn=train_index,
+            inputs=[exp_dir],
+            outputs=[self.trained_index],
+        )

configs/config.py

@@ -132,7 +132,7 @@ class Config:
             if self.has_xpu():
                 self.device = self.instead = "xpu:0"
                 self.is_half = True
-            i_device = int(self.device.split(":")[-1])
+            i_device = int(0)
             self.gpu_name = torch.cuda.get_device_name(i_device)
             if (
                 ("16" in self.gpu_name and "V100" not in self.gpu_name.upper())

infer/modules/vc/modules.py

@@ -129,16 +129,16 @@ class VC:
 
         self.pipeline = Pipeline(self.tgt_sr, self.config)
         n_spk = self.cpt["config"][-3]
-        index = {"value": get_index_path_from_model(sid), "__type__": "update"}
-        logger.info("Select index: " + index["value"])
+        # index = {"value": get_index_path_from_model(sid), "__type__": "update"}
+        # logger.info("Select index: " + index["value"])
 
         return (
             (
                 {"visible": True, "maximum": n_spk, "__type__": "update"},
                 to_return_protect0,
                 to_return_protect1,
-                index,
-                index,
+                # index,
+                # index,
             )
             if to_return_protect
             else {"visible": True, "maximum": n_spk, "__type__": "update"}