Update app.py

app.py

@@ -19,41 +19,7 @@ import TMIDIX
 import matplotlib.pyplot as plt
 
 in_space = os.getenv("SYSTEM") == "spaces"
-
-# =================================================================================================
-
-@spaces.GPU
-def generate_drums(notes_times,
-                   max_drums_limit = 8,
-                   num_memory_tokens = 4096,
-                   temperature=0.9):
-
-    ctx = torch.amp.autocast(device_type='cuda', dtype=torch.float16)
-
-    x = torch.tensor([notes_times] * 1, dtype=torch.long, device='cuda')
-
-    o = 128
-
-    ncount = 0
-
-    while o > 127 and ncount < max_drums_limit:
-      with ctx:
-        out = model.generate(x[-num_memory_tokens:],
-                            1,
-                            temperature=temperature,
-                            return_prime=False,
-                            verbose=False)
-
-      o = out.tolist()[0][0]
-
-      if 256 <= o < 384:
-        ncount += 1
-
-      if o > 127:
-        x = torch.cat((x, out), 1)
-
-    return x.tolist()[0][len(notes_times):]
-
+         
 # =================================================================================================
                        
 @spaces.GPU
@@ -150,15 +116,38 @@ def GenerateDrums(input_midi, input_num_tokens):
 
     output = []
 
+    temperature=0.9,
+    max_drums_limit=8,
+    num_memory_tokens=4096
+
     for c in comp_times[:input_num_tokens]:
-      output.append(c)
-    
-      out = generate_drums(output,
-                           temperature=0.9,
-                           max_drums_limit=8,
-                           num_memory_tokens=4096
-                           )
-      output.extend(out)
+        output.append(c)
+    
+        x = torch.tensor([output] * 1, dtype=torch.long, device=DEVICE)
+        
+        o = 128
+        
+        ncount = 0
+        
+        while o > 127 and ncount < max_drums_limit:
+          with ctx:
+            out = model.generate(x[-num_memory_tokens:],
+                                1,
+                                temperature=temperature,
+                                return_prime=False,
+                                verbose=False)
+        
+          o = out.tolist()[0][0]
+        
+          if 256 <= o < 384:
+            ncount += 1
+        
+          if o > 127:
+            x = torch.cat((x, out), 1)
+        
+        x_output = x.tolist()[0][len(notes_times):]
+        
+        output.extend(x_output1)
 
     print('=' * 70)
     print('Done!')