add initial vllm inf script
This commit is contained in:
slSeanWU
@@ -208,3 +208,5 @@ __marimo__/
|
||||
|
||||
# artifacts
|
||||
soundfonts/
|
||||
generated_outputs/
|
||||
test_outputs/
|
||||
@@ -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.
|
||||
@@ -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()
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user