zerogpu

app.py

@@ -1,3 +1,4 @@
+import spaces
 import argparse
 import glob
 import json
@@ -94,11 +95,12 @@ def create_msg(name, data):
 def send_msgs(msgs):
     return json.dumps(msgs)
 
-
+@spaces.GPU()
 def run(model_name, tab, mid_seq, continuation_state, instruments, drum_kit, bpm, time_sig, key_sig, mid, midi_events,
         reduce_cc_st, remap_track_channel, add_default_instr, remove_empty_channels, seed, seed_rand,
         gen_events, temp, top_p, top_k, allow_cc):
     model = models[model_name]
+    model.to(device=opt.device, dtype=torch.bfloat16 if opt.device == "cuda" else torch.float32).eval()
     tokenizer = model.tokenizer
     bpm = int(bpm)
     if time_sig == "auto":
@@ -300,10 +302,9 @@ if __name__ == "__main__":
     for name, (repo_id, path, config) in models_info.items():
         model_path = hf_hub_download_retry(repo_id=repo_id, filename=f"{path}model.ckpt")
         model = MIDIModel(config=MIDIModelConfig.from_name(config))
-        ckpt = torch.load(model_path, map_location="cpu")
+        ckpt = torch.load(model_path, map_location="cpu", weights_only=True)
         state_dict = ckpt.get("state_dict", ckpt)
         model.load_state_dict(state_dict, strict=False)
-        model.to(device=opt.device, dtype=torch.bfloat16 if opt.device == "cuda" else torch.float32).eval()
         models[name] = model
 
     load_javascript()