Fix and add quick option

app.py

@@ -24,24 +24,25 @@ cfg.config = "fusion_stage3"
 cfg.print_config = False
 cfg.data_config = None
 cfg.phase = "inference"
-cfg.weight = None
 cfg.num_workers = 1
 
 @spaces.GPU
 @torch.inference_mode()
-def predict_mos(audio_path: str, domain: str) -> float:
+def predict_mos(audio_path: str, domain: str, quick: bool) -> float:
     data = pd.DataFrame({"file_path": [audio_path]})
     data["dataset"] = domain
-    data['mos'] = 0
-
+    data["mos"] = 0
     preds = 0.0
     for fold in range(5):
         cfg.now_fold = fold
+        cfg.weight = f"models/fusion_stage3/fold{fold}_s42_best_model.pth"
         model = get_model(cfg, device).eval()
         for _ in range(5):
             test_dataset = get_dataset(cfg, data, "test")
             p = model(*[torch.tensor(t).unsqueeze(0).to(device) for t in test_dataset[0][:-1]])
-            preds += p.cpu().numpy()[0]
+            preds += p.cpu().numpy()[0][0]
+            if quick:
+                return preds
     preds /= 25.0
     return preds
 
@@ -65,12 +66,22 @@ with gr.Blocks() as demo:
                     "blizzard2011",
                 ],
                 label="Data-domain ID for the MOS prediction",
-                value="sarulab"
+                value="sarulab",
+            )
+            quick = gr.Checkbox(
+                label="Quick prediction",
+                value=True,
+                info=(
+                    "UTMOSv2 makes predictions repeatedly for five randomly selected frames "
+                    "of the input speech waveform for all five folds. "
+                    "To make quick predictions by reducing this to a single repetition, "
+                    "check this checkbox:",
+                ),
             )
             submit = gr.Button(value="Submit")
 
         with gr.Column():
             output = gr.Textbox(label="Predicted MOS", type="text")
-    submit.click(fn=predict_mos, inputs=[audio, domain], outputs=[output])
+    submit.click(fn=predict_mos, inputs=[audio, domain, quick], outputs=[output])
 
 demo.queue().launch()