amt context history

app.py

@@ -1,5 +1,5 @@
 import os
-import spaces  # Enables ZeroGPU on Hugging Face
+#import spaces  # Enables ZeroGPU on Hugging Face
 import gradio as gr
 import torch
 from dataclasses import asdict
@@ -26,8 +26,8 @@ model_card = ModelCard(
     name="Anticipatory Music Transformer",
    description=(
     "Generate musical accompaniment for your existing vamp using the Anticipatory Music Transformer. "
-    "Input: a MIDI file with a short accompaniment followed by a melody line. "
-    "Output: a new MIDI file with extended accompaniment matching the melody. "
+    "Input: a MIDI file with a short accompaniment (vamp) followed by a melody line. "
+    "Output: a new MIDI file with extended accompaniment matching the melody continuation. "
     "Use the sliders to choose model size and how much of the song is used as context."
 ),
     author="John Thickstun, David Hall, Chris Donahue, Percy Liang",
@@ -188,7 +188,7 @@ def auto_extract_melody(mid, debug=False):
 
     return events, melody
 
-@spaces.GPU
+#@spaces.GPU
 # Core generation 
 def generate_accompaniment(midi_path: str, model_choice: str, history_length: float):
     """
@@ -207,19 +207,23 @@ def generate_accompaniment(midi_path: str, model_choice: str, history_length: fl
         print("No melody detected; using all events")
         melody = all_events
 
-    melody_time = mid.length  # Mido computes total duration in seconds
-    total_time = round(mid.length or ops.max_time(all_events, seconds=True))
+    ## Compute total time
+    mid_time = mid.length or 0
+    ops_time = ops.max_time(all_events, seconds=True)
+    total_time = round(max(mid_time, ops_time))
 
     # History portion
-    history = ops.clip(all_events, 0, history_length, clip_duration=False)
+    #history = ops.clip(all_events, 0, history_length, clip_duration=False)
+    melody_history = ops.clip(all_events, 0, history_length, clip_duration=False)
+    melody_future  = ops.clip(melody, history_length, total_time, clip_duration=False)
 
     # Generate accompaniment for the remaining duration
     accompaniment = generate(
         model,
         start_time=history_length,   # start after history
         end_time=total_time,         # go to end
-        inputs=history,
-        controls=melody,
+        inputs=melody_history,
+        controls=melody_future,
         top_p=0.95,
         debug=False
     )
@@ -232,6 +236,8 @@ def generate_accompaniment(midi_path: str, model_choice: str, history_length: fl
         clip_duration=True
     )
 
+    print(f"[DEBUG] Context events: {len(melody_history)}, Future melody events: {len(melody_future)}")
+
     print(f"Generating from {history_length:.2f}s -> {total_time:.2f}s "
       f"(duration = {total_time - history_length:.2f}s)")