add initial vllm inf script

.gitignore

@@ -208,3 +208,5 @@ __marimo__/
 
 # artifacts
 soundfonts/
+generated_outputs/
+test_outputs/

assets/example_prompts.txt

@@ -0,0 +1,5 @@
+Rock piece suitable for a hero's epic battle against a dark villain. The music is rich with synth leads, brass section, and heavy drums.
+A sad and emotional song featuring relentless, heartbreaking acoustic guitars.
+Slow and loving music that is best for a romantic film. The piano and string ensemble give a classical vibe.
+An energetic and motivating pop song you love to hear on a long road trip.
+Upbeat and playful jazz music with lively saxophones, like you're going out on a Sunday picnic.

generate_vllm.py

@@ -0,0 +1,497 @@
+#!/usr/bin/env python3
+"""
+This script generates MIDI files from text prompts using the MIDI-LLM model with vLLM backend.
+vLLM provides faster inference compared to standard HuggingFace model.generate() mixin.
+"""
+
+import json
+import time
+import argparse
+from pathlib import Path
+from datetime import datetime
+from typing import List, Optional
+
+import torch
+import tqdm
+from vllm import LLM, SamplingParams, TokensPrompt
+from transformers import AutoTokenizer
+
+# Import helper functions and constants
+from midi_llm.utils import (
+    save_generation,
+    synthesize_midi_to_audio,
+    has_excessive_notes_at_any_time,
+    AMT_GPT2_BOS_ID,
+    LLAMA_VOCAB_SIZE,
+    LLAMA_MODEL_NAME,
+    ALLOWED_TOKEN_IDS,
+    SYNTHESIS_AVAILABLE,
+)
+
+# Default generation parameters
+DEFAULT_TEMPERATURE = 1.0
+DEFAULT_TOP_P = 0.98
+DEFAULT_MAX_TOKENS = 2046
+DEFAULT_N_OUTPUTS = 4  # give more outputs for variability
+
+
+def prepare_vllm_model(
+    model_path: str,
+    temperature: float,
+    top_p: float,
+    max_tokens: int,
+    n_outputs: int,
+    do_fp8_quantization: bool = False,
+    gpu_memory_utilization: float = 0.9
+) -> tuple:
+    """
+    Initialize vLLM model and sampling parameters.
+    
+    Args:
+        model_path: Path to model checkpoint
+        temperature: Sampling temperature
+        top_p: Nucleus sampling parameter
+        max_tokens: Maximum tokens to generate
+        n_outputs: Number of outputs per prompt
+        do_fp8_quantization: Whether to use FP8 quantization
+        gpu_memory_utilization: Fraction of GPU memory to use
+        
+    Returns:
+        Tuple of (model, sampling_params)
+    """
+    sampling_params = SamplingParams(
+        temperature=temperature,
+        top_p=top_p,
+        n=n_outputs,
+        max_tokens=max_tokens,
+        allowed_token_ids=ALLOWED_TOKEN_IDS,
+    )
+    
+    print(f"\n{'='*70}")
+    print("Model Configuration")
+    print(f"{'='*70}")
+    print(f"Model path: {model_path}")
+    print(f"Quantization: {'FP8' if do_fp8_quantization else 'None (BF16)'}")
+    print(f"GPU memory utilization: {gpu_memory_utilization:.1%}")
+    print(f"\nSampling Parameters:")
+    print(f"  Temperature: {temperature}")
+    print(f"  Top-p: {top_p}")
+    print(f"  Max tokens: {max_tokens}")
+    print(f"  Outputs per prompt: {n_outputs}")
+    print(f"{'='*70}\n")
+    
+    model = LLM(
+        model=model_path,
+        quantization="fp8" if do_fp8_quantization else None,
+        gpu_memory_utilization=gpu_memory_utilization,
+        trust_remote_code=True,
+    )
+    
+    print(f"✓ Model loaded successfully\n")
+    
+    return model, sampling_params
+
+
+def generate_from_prompts(
+    model: LLM,
+    tokenizer: AutoTokenizer,
+    prompts: List[str],
+    sampling_params: SamplingParams,
+    output_dir: Path,
+    soundfont_path: Optional[str] = None,
+    synthesize: bool = False,
+    system_prompt: Optional[str] = None
+) -> dict:
+    """
+    Generate MIDI from text prompts and save results.
+    
+    Args:
+        model: vLLM model
+        tokenizer: HuggingFace tokenizer
+        prompts: List of text prompts
+        sampling_params: vLLM sampling parameters
+        output_dir: Base output directory (timestamped subdirs will be created inside)
+        soundfont_path: Path to SoundFont file
+        synthesize: Whether to synthesize to audio
+        system_prompt: Optional system prompt prefix
+        
+    Returns:
+        Dictionary with generation statistics and output files
+    """
+    # Default system prompt
+    if system_prompt is None:
+        system_prompt = "You are a world-class composer. Please compose some music according to the following description: "
+    
+    stats = {
+        "total_prompts": len(prompts),
+        "successful_generations": 0,
+        "failed_generations": 0,
+        "generation_times": [],
+        "output_files": []  # Track all generated files
+    }
+    
+    for idx, prompt in enumerate(tqdm.tqdm(prompts, desc="Generating")):
+        print(f"\n[{idx+1}/{len(prompts)}] Prompt: {prompt}")
+        
+        # Create output directory for this prompt with timestamp
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        prompt_output_dir = output_dir / f"{timestamp}_prompt_{idx+1}"
+        
+        # Prepare full prompt
+        # add space to the end of each prompt to match training
+        full_prompt = system_prompt + prompt + " "
+        
+        # Tokenize
+        llama_input = tokenizer(full_prompt, padding=False)
+        input_ids = llama_input["input_ids"]
+        
+        # Add MIDI BOS token (AMT_GPT2_BOS_ID in extended vocab)
+        input_ids.append(AMT_GPT2_BOS_ID + LLAMA_VOCAB_SIZE)
+        
+        # Generate
+        start_time = time.time()
+        vllm_input = [TokensPrompt(prompt_token_ids=input_ids)]
+        outputs = model.generate(vllm_input, sampling_params)
+        generation_time = time.time() - start_time
+        
+        if idx > 0:  # Skip first generation for timing (warmup)
+            stats["generation_times"].append(generation_time)
+        
+        print(f"Generation time: {generation_time:.2f}s")
+        
+        # Save all outputs for this prompt
+        n_outputs = len(outputs[0].outputs)
+        successful_outputs = 0
+        prompt_files = []
+        
+        for output_idx in range(n_outputs):
+            # Extract tokens and shift back to MIDI vocab range
+            token_ids = outputs[0].outputs[output_idx].token_ids
+            midi_tokens = [t - LLAMA_VOCAB_SIZE for t in token_ids]
+            
+            # Save generation
+            success = save_generation(
+                tokens=midi_tokens,
+                prompt=prompt,
+                output_dir=prompt_output_dir,
+                generation_idx=output_idx + 1,
+                soundfont_path=soundfont_path,
+                synthesize=synthesize
+            )
+            
+            if success:
+                successful_outputs += 1
+                # Track output files
+                midi_file = prompt_output_dir / f"gen_{output_idx + 1}.mid"
+                prompt_files.append(str(midi_file))
+                if synthesize and soundfont_path:
+                    mp3_file = prompt_output_dir / f"gen_{output_idx + 1}.mp3"
+                    if mp3_file.exists():
+                        prompt_files.append(str(mp3_file))
+        
+        print(f"Successfully saved {successful_outputs}/{n_outputs} outputs")
+        stats["successful_generations"] += successful_outputs
+        stats["failed_generations"] += (n_outputs - successful_outputs)
+        stats["output_files"].extend(prompt_files)
+    
+    return stats
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate MIDI files from text prompts using MIDI-LLM with vLLM",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Generate from a single prompt (there will be 4 outputs by default)
+  python generate_vllm.py --model path/to/checkpoint \\
+      --prompt "A cheerful piano melody"
+  
+  # Generate single output without synthesis
+  python generate_vllm.py --model path/to/checkpoint \\
+      --prompt "A relaxing jazz piece" \\
+      --n_outputs 1 \\
+      --no-synthesize
+  
+  # Interactive mode (with initial prompt)
+  python generate_vllm.py --model path/to/checkpoint \\
+      --prompt "A cheerful melody" \\
+      --interactive
+  
+  # Interactive-only mode (no initial prompt)
+  python generate_vllm.py --model path/to/checkpoint \\
+      --interactive
+  
+  # Generate from prompts file with FP8 quantization
+  python generate_vllm.py --model path/to/checkpoint \\
+      --prompts_file prompts.txt \\
+      --fp8 \\
+      --temperature 1.0 \\
+      --top_p 0.98
+        """
+    )
+    
+    # Required arguments
+    parser.add_argument(
+        "--model",
+        type=str,
+        required=True,
+        help="Path to MIDI-LLM model checkpoint"
+    )
+    
+    # Input arguments (not required if using --interactive only)
+    input_group = parser.add_mutually_exclusive_group(required=False)
+    input_group.add_argument(
+        "--prompt",
+        type=str,
+        help="Single text prompt for generation"
+    )
+    input_group.add_argument(
+        "--prompts_file",
+        type=str,
+        help="Path to file containing prompts (one per line)"
+    )
+    
+    # Output arguments
+    parser.add_argument(
+        "--output_root",
+        type=str,
+        default="./generated_outputs",
+        help="Root directory for outputs (timestamped subdirs will be created inside, default: ./generated_outputs)"
+    )
+    parser.add_argument(
+        "--n_outputs",
+        type=int,
+        default=DEFAULT_N_OUTPUTS,
+        help=f"Number of outputs to generate per prompt (default: {DEFAULT_N_OUTPUTS})"
+    )
+    
+    # Synthesis arguments
+    parser.add_argument(
+        "--no-synthesize",
+        dest="synthesize",
+        action="store_false",
+        help="Skip audio synthesis (only generate MIDI files)"
+    )
+    parser.set_defaults(synthesize=True)
+    parser.add_argument(
+        "--soundfont",
+        type=str,
+        default="./soundfonts/FluidR3_GM/FluidR3_GM.sf2",
+        help="Path to SoundFont file for synthesis (default: ./soundfonts/FluidR3_GM/FluidR3_GM.sf2)"
+    )
+    
+    # Generation parameters
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=DEFAULT_TEMPERATURE,
+        help=f"Sampling temperature (default: {DEFAULT_TEMPERATURE})"
+    )
+    parser.add_argument(
+        "--top_p",
+        type=float,
+        default=DEFAULT_TOP_P,
+        help=f"Nucleus sampling threshold (default: {DEFAULT_TOP_P})"
+    )
+    parser.add_argument(
+        "--max_tokens",
+        type=int,
+        default=DEFAULT_MAX_TOKENS,
+        help=f"Maximum tokens to generate (default: {DEFAULT_MAX_TOKENS})"
+    )
+    
+    # Model arguments
+    parser.add_argument(
+        "--fp8",
+        action="store_true",
+        help="Use FP8 quantization for faster inference (requires compatible GPU)"
+    )
+    parser.add_argument(
+        "--gpu_memory",
+        type=float,
+        default=0.9,
+        help="Fraction of GPU memory to use (default: 0.9)"
+    )
+    parser.add_argument(
+        "--cache_dir",
+        type=str,
+        default=None,
+        help="HuggingFace cache directory (default: $HF_HOME or ~/.cache/huggingface)"
+    )
+    parser.add_argument(
+        "--interactive",
+        action="store_true",
+        help="Enter interactive mode after initial generation (keep generating until empty prompt)"
+    )
+    
+    args = parser.parse_args()
+    
+    # Validate that either prompts are provided or interactive mode is enabled
+    if not args.prompt and not args.prompts_file and not args.interactive:
+        parser.error("Either --prompt, --prompts_file, or --interactive must be specified")
+    
+    # Load prompts (if provided)
+    prompts = []
+    if args.prompt:
+        prompts = [args.prompt]
+    elif args.prompts_file:
+        with open(args.prompts_file, "r") as f:
+            prompts = [line.strip() for line in f if line.strip()]
+        print(f"Loaded {len(prompts)} prompts from {args.prompts_file}")
+    
+    # Check synthesis requirements
+    if args.synthesize:
+        soundfont_path = Path(args.soundfont)
+        if not soundfont_path.exists():
+            print(f"Error: SoundFont not found at {soundfont_path}")
+            print("Please download a SoundFont or disable synthesis with --no-synthesize")
+            sys.exit(1)
+        
+        if not SYNTHESIS_AVAILABLE:
+            print("Warning: Audio synthesis libraries not available.")
+            print("Synthesis will be skipped. Install dependencies:")
+            print("  conda install conda-forge::fluidsynth conda-forge::ffmpeg")
+            print("  pip install midi2audio librosa soundfile")
+            args.synthesize = False
+    
+    # Create output root directory with timestamp
+    output_root = Path(args.output_root)
+    session_timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S")
+    output_dir = output_root / session_timestamp
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    print(f"Output directory: {output_dir.absolute()}\n")
+    
+    # Load tokenizer
+    print("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        LLAMA_MODEL_NAME,
+        cache_dir=args.cache_dir,
+        pad_token="<|eot_id|>",
+    )
+    
+    # Load model
+    model, sampling_params = prepare_vllm_model(
+        model_path=args.model,
+        temperature=args.temperature,
+        top_p=args.top_p,
+        max_tokens=args.max_tokens,
+        n_outputs=args.n_outputs,
+        do_fp8_quantization=args.fp8,
+        gpu_memory_utilization=args.gpu_memory
+    )
+    
+    # Generate from initial prompts (if provided)
+    if prompts:
+        print(f"Starting generation for {len(prompts)} prompt(s)...\n")
+        start_time = time.time()
+        
+        stats = generate_from_prompts(
+            model=model,
+            tokenizer=tokenizer,
+            prompts=prompts,
+            sampling_params=sampling_params,
+            output_dir=output_dir,
+            soundfont_path=args.soundfont if args.synthesize else None,
+            synthesize=args.synthesize
+        )
+        
+        total_time = time.time() - start_time
+        
+        # Print summary
+        print(f"\n{'='*70}")
+        print("Generation Summary")
+        print(f"{'='*70}")
+        print(f"Total prompts: {stats['total_prompts']}")
+        print(f"Successful generations: {stats['successful_generations']}")
+        print(f"Failed generations: {stats['failed_generations']}")
+        print(f"Total time: {total_time:.2f}s")
+        
+        if stats['generation_times']:
+            avg_time = sum(stats['generation_times']) / len(stats['generation_times'])
+            print(f"Average generation time: {avg_time:.2f}s (excluding warmup)")
+        
+        print(f"\nOutputs saved to: {output_dir.absolute()}")
+        
+        # Print generated files
+        if stats['output_files']:
+            print(f"\nGenerated files:")
+            for file_path in stats['output_files']:
+                file_type = "🎵 MIDI" if file_path.endswith('.mid') else "🎧 Audio"
+                print(f"  {file_type}: {file_path}")
+        
+        print(f"{'='*70}\n")
+        
+        # Save stats to JSON
+        stats_file = output_dir / "generation_stats.json"
+        with open(stats_file, "w") as f:
+            json.dump({
+                **stats,
+                "total_time": total_time,
+                "average_time": sum(stats['generation_times']) / len(stats['generation_times']) if stats['generation_times'] else 0,
+                "config": {
+                    "model": args.model,
+                    "temperature": args.temperature,
+                    "top_p": args.top_p,
+                    "max_tokens": args.max_tokens,
+                    "n_outputs": args.n_outputs,
+                    "fp8": args.fp8,
+                }
+            }, f, indent=2)
+    else:
+        print(f"No initial prompts provided. Starting in interactive mode...\n")
+    
+    # Interactive mode
+    if args.interactive:
+        print(f"\n{'='*70}")
+        print("Interactive Mode")
+        print(f"{'='*70}")
+        print("Enter prompts to generate more MIDI files.")
+        print("Press Enter with empty prompt to exit.\n")
+        
+        while True:
+            try:
+                # Get user input
+                user_prompt = input("Prompt: ").strip()
+                
+                # Exit if empty
+                if not user_prompt:
+                    print("\nExiting interactive mode. Goodbye!")
+                    break
+                
+                # Generate from the new prompt
+                print()
+                interactive_stats = generate_from_prompts(
+                    model=model,
+                    tokenizer=tokenizer,
+                    prompts=[user_prompt],
+                    sampling_params=sampling_params,
+                    output_dir=output_dir,
+                    soundfont_path=args.soundfont if args.synthesize else None,
+                    synthesize=args.synthesize
+                )
+                
+                # Print mini summary
+                print(f"\n✓ Generated {interactive_stats['successful_generations']}/{args.n_outputs} outputs")
+                if interactive_stats['generation_times']:
+                    print(f"  Generation time: {interactive_stats['generation_times'][0]:.2f}s")
+                
+                # Print file paths
+                if interactive_stats['output_files']:
+                    for file_path in interactive_stats['output_files']:
+                        file_type = "🎵" if file_path.endswith('.mid') else "🎧"
+                        print(f"  {file_type} {file_path}")
+                print()
+                
+            except KeyboardInterrupt:
+                print("\n\nInterrupted. Exiting interactive mode.")
+                break
+            except EOFError:
+                print("\n\nExiting interactive mode.")
+                break
+
+
+if __name__ == "__main__":
+    main()
+

midi_llm/__init__.py

@@ -0,0 +1,9 @@
+"""
+MIDI-LLM: Utility functions and helpers.
+
+This package contains supporting code for MIDI-LLM generation, training, and data processing.
+Users new to the codebase can skip this directory - start with generate_vllm.py or train.py instead.
+"""
+
+__version__ = "0.1.0"
+

midi_llm/utils.py

@@ -0,0 +1,220 @@
+"""
+Utility functions for MIDI-LLM.
+
+This module contains helper functions for audio synthesis, MIDI conversion,
+and other supporting operations. Users can safely skip this file when learning
+the codebase - start with generate_vllm.py or train.py instead.
+"""
+
+import os
+import sys
+from pathlib import Path
+from typing import List, Optional, Union
+
+import torch
+
+# Core dependency - required
+try:
+    from anticipation.convert import events_to_midi
+except ImportError:
+    print("Error: anticipation package not found. Please install it for MIDI conversion.")
+    print("Install with: pip install anticipation")
+    sys.exit(1)
+
+# Optional dependencies for audio synthesis
+SYNTHESIS_AVAILABLE = False
+try:
+    import midi2audio
+    import librosa
+    import librosa.effects
+    import soundfile as sf
+    SYNTHESIS_AVAILABLE = True
+except ImportError:
+    pass
+
+
+# ============================================================================
+# Constants
+# ============================================================================
+
+AMT_GPT2_BOS_ID = 55026
+LLAMA_VOCAB_SIZE = 128256
+LLAMA_MODEL_NAME = "meta-llama/Llama-3.2-1B"
+
+# MIDI tokens are in the extended vocabulary range
+ALLOWED_TOKEN_IDS = list(range(LLAMA_VOCAB_SIZE, LLAMA_VOCAB_SIZE + AMT_GPT2_BOS_ID))
+
+
+# ============================================================================
+# Validation
+# ============================================================================
+
+def has_excessive_notes_at_any_time(
+    tokens: Union[torch.Tensor, List[int]], 
+    max_notes_per_time: int = 64
+) -> bool:
+    """
+    Check if generated MIDI has excessive simultaneous notes at any time point.
+    
+    This validation helps filter out invalid or unrealistic generations that have
+    too many notes playing at once, which can indicate a failure mode.
+    
+    Args:
+        tokens: Token sequence (torch.Tensor or list of ints)
+        max_notes_per_time: Maximum allowed notes at any single time point
+        
+    Returns:
+        True if excessive notes detected, False otherwise
+    """
+    # Convert to tensor if needed
+    if isinstance(tokens, list):
+        tokens = torch.tensor(tokens)
+    
+    # Extract time tokens (every 3rd token in the sequence: time, duration, note)
+    times = tokens[::3]
+    
+    # Use torch.bincount for efficient counting
+    # bincount returns counts for indices 0 to max_value
+    counts = torch.bincount(times)
+    
+    # Check if any time has more than max_notes_per_time notes
+    return torch.any(counts > max_notes_per_time).item()
+
+
+# ============================================================================
+# Audio Synthesis
+# ============================================================================
+
+def synthesize_midi_to_audio(
+    midi_path: str, 
+    soundfont_path: str,
+    save_mp3: bool = True,
+    samplerate: Optional[int] = None
+) -> bool:
+    """
+    Synthesize MIDI file to audio (WAV/MP3) using FluidSynth.
+    
+    Args:
+        midi_path: Path to MIDI file
+        soundfont_path: Path to SoundFont (.sf2) file
+        save_mp3: If True, convert to MP3 and delete WAV
+        samplerate: Optional sample rate for audio
+        
+    Returns:
+        True if successful, False otherwise
+    """
+    if not SYNTHESIS_AVAILABLE:
+        print("Warning: Audio synthesis libraries not available. Skipping synthesis.")
+        print("Install with: conda install conda-forge::fluidsynth conda-forge::ffmpeg")
+        print("              pip install midi2audio librosa soundfile")
+        return False
+    
+    try:
+        wav_path = midi_path.replace(".mid", ".wav")
+        
+        # Initialize FluidSynth
+        fs = midi2audio.FluidSynth(soundfont_path)
+        if samplerate is not None:
+            fs.sample_rate = samplerate
+        
+        # Synthesize MIDI to WAV
+        fs.midi_to_audio(midi_path, wav_path)
+        
+        # Trim silence from audio
+        wav, sr = librosa.load(wav_path)
+        wav, _ = librosa.effects.trim(wav, top_db=30)
+        sf.write(wav_path, wav, sr)
+        
+        if save_mp3:
+            # Convert WAV to MP3 using ffmpeg
+            mp3_path = midi_path.replace(".mid", ".mp3")
+            if samplerate is None:
+                cmd = f"ffmpeg -i {wav_path} -codec:a libmp3lame -qscale:a 2 {mp3_path} -y >/dev/null 2>&1"
+            else:
+                cmd = f"ffmpeg -i {wav_path} -codec:a libmp3lame -qscale:a 2 -ar {samplerate} {mp3_path} -y >/dev/null 2>&1"
+            
+            os.system(cmd)
+            
+            # Remove WAV file
+            if os.path.exists(wav_path):
+                os.remove(wav_path)
+        
+        return True
+    
+    except Exception as e:
+        print(f"Error synthesizing MIDI to audio: {e}")
+        return False
+
+
+# ============================================================================
+# MIDI Generation and Saving
+# ============================================================================
+
+def save_generation(
+    tokens: List[int],
+    prompt: str,
+    output_dir: Path,
+    generation_idx: int,
+    soundfont_path: Optional[str] = None,
+    synthesize: bool = False,
+    validate: bool = True
+) -> bool:
+    """
+    Save generated tokens as MIDI file (and optionally audio).
+    
+    Args:
+        tokens: List of generated token IDs (already shifted from LLAMA vocab)
+        prompt: Original text prompt
+        output_dir: Directory to save outputs
+        generation_idx: Index of this generation (for multiple outputs)
+        soundfont_path: Path to SoundFont file for synthesis
+        synthesize: Whether to synthesize to audio
+        validate: Whether to validate tokens before saving (checks for excessive notes)
+        
+    Returns:
+        True if successful, False otherwise
+    """
+    try:
+        # Validate tokens before saving
+        if validate:
+            if has_excessive_notes_at_any_time(tokens, max_notes_per_time=64):
+                print(f"  ✗ Generation {generation_idx}: Failed validation (excessive simultaneous notes)")
+                return False
+        
+        # Create output directory
+        output_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Save prompt text
+        prompt_file = output_dir / "prompt.txt"
+        with open(prompt_file, "w") as f:
+            f.write(prompt)
+        
+        # Save token sequence
+        tokens_file = output_dir / f"gen_{generation_idx}_tokens.txt"
+        with open(tokens_file, "w") as f:
+            for token in tokens:
+                f.write(f"{token}\n")
+        
+        # Convert tokens to MIDI
+        midi_obj = events_to_midi(tokens)
+        midi_file = output_dir / f"gen_{generation_idx}.mid"
+        midi_obj.save(str(midi_file))
+        
+        print(f"  ✓ Saved MIDI: {midi_file}")
+        
+        # Optionally synthesize to audio
+        if synthesize and soundfont_path:
+            success = synthesize_midi_to_audio(
+                str(midi_file),
+                soundfont_path,
+                save_mp3=True
+            )
+            if success:
+                print(f"  ✓ Synthesized audio: {midi_file.with_suffix('.mp3')}")
+        
+        return True
+    
+    except Exception as e:
+        print(f"  ✗ Error saving generation {generation_idx}: {e}")
+        return False
+

requirements.txt

@@ -6,6 +6,7 @@ anticipation @ git+https://github.com/jthickstun/anticipation.git@af37397922665a
 anyio==4.11.0
 astor==0.8.1
 attrs==25.4.0
+audioread==3.1.0
 blake3==1.0.8
 cachetools==6.2.1
 cbor2==5.7.0
@@ -16,6 +17,7 @@ click==8.2.1
 cloudpickle==3.1.1
 compressed-tensors==0.11.0
 cupy-cuda12x==13.6.0
+decorator==5.2.1
 depyf==0.19.0
 dill==0.4.0
 diskcache==5.6.3
@@ -42,9 +44,12 @@ idna==3.11
 interegular==0.3.3
 Jinja2==3.1.6
 jiter==0.11.1
+joblib==1.5.2
 jsonschema==4.25.1
 jsonschema-specifications==2025.9.1
 lark==1.2.2
+lazy_loader==0.4
+librosa==0.11.0
 llguidance==0.7.30
 llvmlite==0.44.0
 lm-format-enforcer==0.11.3
@@ -84,6 +89,8 @@ outlines_core==0.2.11
 packaging==25.0
 partial-json-parser==0.2.1.1.post6
 pillow==11.3.0
+platformdirs==4.5.0
+pooch==1.8.2
 prometheus-fastapi-instrumentator==7.1.0
 prometheus_client==0.23.1
 propcache==0.4.1
@@ -111,6 +118,7 @@ rich-toolkit==0.15.1
 rignore==0.7.1
 rpds-py==0.28.0
 safetensors==0.6.2
+scikit-learn==1.7.2
 scipy==1.16.2
 sentencepiece==0.2.1
 sentry-sdk==2.42.1
@@ -121,6 +129,7 @@ soundfile==0.13.1
 soxr==1.0.0
 starlette==0.48.0
 sympy==1.14.0
+threadpoolctl==3.6.0
 tiktoken==0.12.0
 tokenizers==0.22.1
 torch==2.8.0