tf32

app.py

@@ -189,24 +189,22 @@ def run(model_name, tab, mid_seq, continuation_state, instruments, drum_kit, bpm
         init_msgs += [create_msg("visualizer_clear", tokenizer.version),
                       create_msg("visualizer_append", events)]
     yield mid_seq, continuation_state, None, None, seed, send_msgs(init_msgs)
-    ctx = torch.amp.autocast(device_type=opt.device, dtype=torch.bfloat16, enabled=opt.device != "cpu")
-    with ctx:
-        midi_generator = generate(model, mid, max_len=max_len, temp=temp, top_p=top_p, top_k=top_k,
-                                  disable_patch_change=disable_patch_change, disable_control_change=not allow_cc,
-                                  disable_channels=disable_channels, generator=generator)
-        events = []
-        t = time.time() + 1
-        for i, token_seq in enumerate(midi_generator):
-            token_seq = token_seq.tolist()
-            mid_seq.append(token_seq)
-            events.append(tokenizer.tokens2event(token_seq))
-            ct = time.time()
-            if ct - t > 0.5:
-                yield (mid_seq, continuation_state, None, None, seed,
-                       send_msgs([create_msg("visualizer_append", events),
-                                  create_msg("progress", [i + 1, gen_events])]))
-                t = ct
-                events = []
+    midi_generator = generate(model, mid, max_len=max_len, temp=temp, top_p=top_p, top_k=top_k,
+                              disable_patch_change=disable_patch_change, disable_control_change=not allow_cc,
+                              disable_channels=disable_channels, generator=generator)
+    events = []
+    t = time.time() + 1
+    for i, token_seq in enumerate(midi_generator):
+        token_seq = token_seq.tolist()
+        mid_seq.append(token_seq)
+        events.append(tokenizer.tokens2event(token_seq))
+        ct = time.time()
+        if ct - t > 0.5:
+            yield (mid_seq, continuation_state, None, None, seed,
+                   send_msgs([create_msg("visualizer_append", events),
+                              create_msg("progress", [i + 1, gen_events])]))
+            t = ct
+            events = []
 
     events = [tokenizer.tokens2event(tokens) for tokens in mid_seq]
     mid = tokenizer.detokenize(mid_seq)
@@ -307,6 +305,10 @@ if __name__ == "__main__":
                    }
     models = {}
     if opt.device == "cuda":
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
         torch.backends.cuda.enable_mem_efficient_sdp(True)
         torch.backends.cuda.enable_flash_sdp(True)
     for name, (repo_id, path, config) in models_info.items():
@@ -315,7 +317,7 @@ if __name__ == "__main__":
         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="cpu", dtype=torch.bfloat16 if opt.device == "cuda" else torch.float32)
+        model.to(device="cpu", dtype=torch.float32)
         models[name] = model
 
     load_javascript()