diff --git a/.gitignore b/.gitignore index d737031..36fb90f 100644 --- a/.gitignore +++ b/.gitignore @@ -207,4 +207,6 @@ marimo/_lsp/ __marimo__/ # artifacts -soundfonts/ \ No newline at end of file +soundfonts/ +generated_outputs/ +test_outputs/ \ No newline at end of file diff --git a/assets/example_prompts.txt b/assets/example_prompts.txt new file mode 100644 index 0000000..571ef5f --- /dev/null +++ b/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. diff --git a/generate_vllm.py b/generate_vllm.py new file mode 100644 index 0000000..7986f04 --- /dev/null +++ b/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() + diff --git a/midi_llm/__init__.py b/midi_llm/__init__.py new file mode 100644 index 0000000..f9aa2ec --- /dev/null +++ b/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" + diff --git a/midi_llm/utils.py b/midi_llm/utils.py new file mode 100644 index 0000000..55293ba --- /dev/null +++ b/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 + diff --git a/requirements.txt b/requirements.txt index 0699482..d90542c 100644 --- a/requirements.txt +++ b/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