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