feat: restore exp

app.py

@@ -1,6 +1,6 @@
 import os
 
-os.environ["PYTORCH_JIT"] = "0"
+os.environ["PYTORCH_JIT"] = "0v"
 
 from random import shuffle
 import gradio as gr
@@ -15,6 +15,22 @@ from infer.modules.train.train import train
 from infer.lib.train.process_ckpt import extract_small_model
 from zero import zero
 
+# patch for jit script
+# if we find `def expand_2d_or_3d_tensor(x,` in /usr/local/lib/python3.10/site-packages/fairseq/models/model_utils.py
+# patch it with `def expand_2d_or_3d_tensor(x: Tensor,`
+FAIRSEQ_CODE = "/usr/local/lib/python3.10/site-packages/fairseq/models/model_utils.py"
+if os.path.exists(FAIRSEQ_CODE):
+    with open(FAIRSEQ_CODE, "r") as f:
+        lines = f.readlines()
+    with open(FAIRSEQ_CODE, "w") as f:
+        for line in lines:
+            if "def expand_2d_or_3d_tensor(x, trg_dim: int, padding_idx: int):" in line:
+                f.write(
+                    "def expand_2d_or_3d_tensor(x: Tensor, trg_dim: int, padding_idx: int) -> Tensor:\n"
+                )
+            else:
+                f.write(line)
+
 
 def extract_audio_files(zip_file: str, target_dir: str) -> list[str]:
     with zipfile.ZipFile(zip_file, "r") as zip_ref:
@@ -186,6 +202,11 @@ def download_expdir(exp_dir: str) -> str:
     return f"{exp_dir}.zip"
 
 
+def restore_expdir(zip: str, exp_dir: str) -> str:
+    shutil.unpack_archive(zip, exp_dir)
+    return exp_dir
+
+
 with gr.Blocks() as app:
     with gr.Row():
         with gr.Column():
@@ -198,6 +219,13 @@ with gr.Blocks() as app:
         with gr.Column():
             preprocess_output = gr.Textbox(label="Preprocessing output", lines=5)
 
+    with gr.Row():
+        restore_zip_file = gr.File(
+            label="Upload the experiment directory zip file",
+            file_types=["zip"],
+        )
+        restore_btn = gr.Button(value="Preprocess")
+
     with gr.Row():
         with gr.Column():
             extract_features_btn = gr.Button(
@@ -260,4 +288,10 @@ with gr.Blocks() as app:
         outputs=[download_expdir_output],
     )
 
+    restore_btn.click(
+        fn=restore_expdir,
+        inputs=[restore_zip_file, exp_dir],
+        outputs=[exp_dir],
+    )
+
     app.launch()

infer/modules/train/train.py

@@ -122,9 +122,6 @@ def run(rank, n_gpus, hps, logger: logging.Logger, state):
         writer = SummaryWriter(log_dir=hps.model_dir)
         writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))
 
-    dist.init_process_group(
-        backend="gloo", init_method="env://", world_size=n_gpus, rank=rank
-    )
     torch.manual_seed(hps.train.seed)
     if torch.cuda.is_available():
         torch.cuda.set_device(rank)
@@ -190,16 +187,6 @@ def run(rank, n_gpus, hps, logger: logging.Logger, state):
         betas=hps.train.betas,
         eps=hps.train.eps,
     )
-    # net_g = DDP(net_g, device_ids=[rank], find_unused_parameters=True)
-    # net_d = DDP(net_d, device_ids=[rank], find_unused_parameters=True)
-    if hasattr(torch, "xpu") and torch.xpu.is_available():
-        pass
-    elif torch.cuda.is_available():
-        net_g = DDP(net_g, device_ids=[rank])
-        net_d = DDP(net_d, device_ids=[rank])
-    else:
-        net_g = DDP(net_g)
-        net_d = DDP(net_d)
 
     try:  # 如果能加载自动resume
         _, _, _, epoch_str = utils.load_checkpoint(