mv batch option to argparse

app.py

@@ -20,7 +20,6 @@ from midi_model import MIDIModel, MIDIModelConfig
 from midi_synthesizer import MidiSynthesizer
 
 MAX_SEED = np.iinfo(np.int32).max
-OUTPUT_BATCH_SIZE = 8
 in_space = os.getenv("SYSTEM") == "spaces"
 
 
@@ -305,7 +304,10 @@ def load_javascript(dir="javascript"):
     javascript = ""
     for path in scripts_list:
         with open(path, "r", encoding="utf8") as jsfile:
-            javascript += f"\n<!-- {path} --><script>{jsfile.read()}</script>"
+            js_content = jsfile.read()
+            js_content = js_content.replace("const MIDI_OUTPUT_BATCH_SIZE=4;",
+                                            f"const MIDI_OUTPUT_BATCH_SIZE={OUTPUT_BATCH_SIZE};")
+            javascript += f"\n<!-- {path} --><script>{js_content}</script>"
     template_response_ori = gr.routes.templates.TemplateResponse
 
     def template_response(*args, **kwargs):
@@ -344,8 +346,10 @@ if __name__ == "__main__":
     parser.add_argument("--share", action="store_true", default=False, help="share gradio app")
     parser.add_argument("--port", type=int, default=7860, help="gradio server port")
     parser.add_argument("--device", type=str, default="cuda", help="device to run model")
+    parser.add_argument("--batch", type=int, default=8, help="batch size")
     parser.add_argument("--max-gen", type=int, default=1024, help="max")
     opt = parser.parse_args()
+    OUTPUT_BATCH_SIZE = opt.batch
     soundfont_path = hf_hub_download_retry(repo_id="skytnt/midi-model", filename="soundfont.sf2")
     thread_pool = ThreadPoolExecutor(max_workers=OUTPUT_BATCH_SIZE)
     synthesizer = MidiSynthesizer(soundfont_path)

javascript/app.js

@@ -1,4 +1,5 @@
-const MIDI_OUTPUT_BATCH_SIZE= 8;
+const MIDI_OUTPUT_BATCH_SIZE=4;
+//Do not change MIDI_OUTPUT_BATCH_SIZE. It will be automatically replaced.
 
 /**
  * 自动绕过 shadowRoot 的 querySelector