enable model selection

app.py

@@ -7,42 +7,75 @@ from anticipation.tokenize import extract_instruments
 import torch
 from pyharp import *
 
+# === Define AMT Model Checkpoints ===
+SMALL_MODEL = "stanford-crfm/music-small-800k"     # Faster inference, worse quality
+MEDIUM_MODEL = "stanford-crfm/music-medium-800k"   # Slower inference, better quality
+LARGE_MODEL = "stanford-crfm/music-large-800k"     # Slowest inference, best quality
+
 # Define the model card for PyHARP
 model_card = ModelCard(
     name="Anticipatory Music Transformer",
     description="Using Anticipatory Music Transformer (AMT) to generate accompaniment for a given MIDI file.",
     author="John Thickstun, David Hall, Chris Donahue, Percy Liang",
     tags=["midi", "generation", "accompaniment"],
-    midi_in=True,  # PyHARP will automatically handle MIDI input
+    midi_in=True,  # PyHARP automatically handles MIDI input
     midi_out=True
 )
 
-# Load the AMT model
-model_name = "stanford-crfm/music-medium-800k"
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
+# === Function to Load AMT Model Based on Selection ===
+def load_amt_model(model_choice):
+    """Loads the selected AMT model."""
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return AutoModelForCausalLM.from_pretrained(model_choice).to(device)
+
+# === Function to Detect the Melody Program Automatically ===
+def detect_melody_program(midi_file):
+    """Automatically detects the only MIDI program in the input file."""
+    events = midi_to_events(midi_file.name)
+    instrument_programs = list(ops.get_instruments(events).keys())
+
+    if len(instrument_programs) == 1:
+        return instrument_programs[0]  # Return the only available program
+    elif len(instrument_programs) > 1:
+        return min(instrument_programs)  # Pick the lowest-numbered program
+    else:
+        return 0  # Default to Acoustic Grand Piano if no program found
 
+# === Function to Generate Accompaniment ===
+def generate_accompaniment(midi_file, model_choice, start_time, end_time):
+    """Generates accompaniment using the selected AMT model."""
+    # Load selected AMT model
+    model = load_amt_model(model_choice)
 
-# Function to generate accompaniment
-def generate_accompaniment(midi_file, selected_midi_program, start_time, end_time):
     # Convert MIDI to events
     events = midi_to_events(midi_file.name)
+    
+    # Automatically detect the melody program
+    melody_program = detect_melody_program(midi_file)
+
     # Clip events based on the selected time range
     clipped_events = ops.clip(events, start_time, end_time)
+
     # Normalize timeline (start at 0)
     clipped_events = ops.translate(clipped_events, -ops.min_time(clipped_events, seconds=False))
-    # Extract the chosen melody instrument
-    clipped_events, melody = extract_instruments(clipped_events, [selected_midi_program])
+
+    # Extract the melody instrument automatically
+    clipped_events, melody = extract_instruments(clipped_events, [melody_program])
+
     # Prepare history (first 5 seconds of the segment)
     history = ops.clip(clipped_events, 0, 5, clip_duration=False)
-    # Generate accompaniment using the AMT model
+
+    # Generate accompaniment using AMT
     accompaniment = generate(
         model, 0, end_time - start_time, inputs=history, controls=melody, top_p=0.98
     )
+
     # Normalize generated accompaniment
     accompaniment = ops.translate(accompaniment, -ops.min_time(accompaniment, seconds=False))
+
     # Combine accompaniment with melody
     output_events = ops.clip(ops.combine(accompaniment, melody), 0, end_time - start_time, clip_duration=True)
+
     # Convert back to MIDI
     output_midi = "generated_accompaniment.midi"
     mid = events_to_midi(output_events)
@@ -50,17 +83,20 @@ def generate_accompaniment(midi_file, selected_midi_program, start_time, end_tim
 
     return output_midi
 
-
-# PyHARP process function
-def process_fn(input_midi, selected_midi_program, start_time, end_time):
-    output_midi = generate_accompaniment(input_midi, selected_midi_program, start_time, end_time)
+# === PyHARP Process Function ===
+def process_fn(input_midi, model_choice, start_time, end_time):
+    """Processes the input and runs AMT with selected model."""
+    output_midi = generate_accompaniment(input_midi, model_choice, start_time, end_time)
     return output_midi, LabelList()
 
-
-# Build Gradio interface wrapped in PyHARP
+# === Build Gradio Interface with Model Selection ===
 with gr.Blocks() as demo:
     components = [
-        gr.Slider(0, 127, step=1, value=53, label="Select Melody Instrument (MIDI Program Number)"),
+        gr.Dropdown(
+            choices=[SMALL_MODEL, MEDIUM_MODEL, LARGE_MODEL], 
+            value=MEDIUM_MODEL, 
+            label="Select AMT Model (Faster vs. Higher Quality)"
+        ),
         gr.Slider(0, 30, step=1, label="Start Time (seconds)"),
         gr.Slider(0, 30, step=1, label="End Time (seconds)"),
     ]
@@ -72,4 +108,4 @@ with gr.Blocks() as demo:
     )
 
 demo.queue()
-demo.launch(share=True, show_error=True)
+demo.launch(share = True,show_error=True)