slSeanWU
474 lines
16 KiB
Python
474 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
MIDI-LLM: Text-to-MIDI Generation using HuggingFace Transformers
|
|
|
|
This script generates MIDI files from text prompts using the MIDI-LLM model with HuggingFace backend.
|
|
Simpler to set up than vLLM but slower for inference.
|
|
"""
|
|
|
|
import json
|
|
import time
|
|
import argparse
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from typing import List, Optional
|
|
|
|
import torch
|
|
import tqdm
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
|
|
# Import helper functions and constants
|
|
from midi_llm.utils import (
|
|
save_generation,
|
|
AMT_GPT2_BOS_ID,
|
|
LLAMA_VOCAB_SIZE,
|
|
LLAMA_MODEL_NAME,
|
|
)
|
|
|
|
# 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_hf_model(model_path: str):
|
|
"""
|
|
Initialize HuggingFace model in BFloat16.
|
|
|
|
Args:
|
|
model_path: Path to model checkpoint
|
|
|
|
Returns:
|
|
model
|
|
"""
|
|
print(f"\n{'='*70}")
|
|
print("Model Configuration")
|
|
print(f"{'='*70}")
|
|
print(f"Model path: {model_path}")
|
|
print(f"Precision: BFloat16")
|
|
print(f"{'='*70}\n")
|
|
|
|
# Load model in BF16
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_path,
|
|
dtype=torch.bfloat16,
|
|
trust_remote_code=True
|
|
).to(device="cuda")
|
|
|
|
model.eval()
|
|
print(f"✓ Model loaded successfully\n")
|
|
|
|
return model
|
|
|
|
|
|
def generate_from_prompts_hf(
|
|
model,
|
|
tokenizer: AutoTokenizer,
|
|
prompts: List[str],
|
|
output_dir: Path,
|
|
model_path: str,
|
|
soundfont_path: Optional[str] = None,
|
|
synthesize: bool = False,
|
|
temperature: float = DEFAULT_TEMPERATURE,
|
|
top_p: float = DEFAULT_TOP_P,
|
|
max_tokens: int = DEFAULT_MAX_TOKENS,
|
|
n_outputs: int = DEFAULT_N_OUTPUTS,
|
|
system_prompt: Optional[str] = None
|
|
) -> dict:
|
|
"""
|
|
Generate MIDI from text prompts using HuggingFace model.
|
|
|
|
Args:
|
|
model: HuggingFace model
|
|
tokenizer: HuggingFace tokenizer
|
|
prompts: List of text prompts
|
|
output_dir: Base output directory
|
|
model_path: Path to model (to check for with_edits)
|
|
soundfont_path: Path to SoundFont file
|
|
synthesize: Whether to synthesize to audio
|
|
temperature: Sampling temperature
|
|
top_p: Nucleus sampling threshold
|
|
max_tokens: Maximum tokens to generate
|
|
n_outputs: Number of outputs per prompt
|
|
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": []
|
|
}
|
|
|
|
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
|
|
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, return_tensors="pt", padding=False)
|
|
input_ids = llama_input["input_ids"]
|
|
|
|
# Add MIDI BOS token
|
|
midi_bos = torch.tensor([[AMT_GPT2_BOS_ID + LLAMA_VOCAB_SIZE]])
|
|
input_ids = torch.cat([input_ids, midi_bos], dim=1)
|
|
|
|
# Move to device
|
|
device = next(model.parameters()).device
|
|
input_ids = input_ids.to(device)
|
|
|
|
# Generate multiple outputs
|
|
start_time = time.time()
|
|
|
|
with torch.no_grad():
|
|
outputs = model.generate(
|
|
input_ids=input_ids,
|
|
do_sample=True,
|
|
max_new_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
num_return_sequences=n_outputs,
|
|
pad_token_id=tokenizer.pad_token_id,
|
|
)
|
|
|
|
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")
|
|
|
|
# Extract only the generated tokens (remove prompt)
|
|
prompt_len = input_ids.shape[1]
|
|
outputs = outputs[:, prompt_len:]
|
|
|
|
# Shift tokens back to MIDI vocab range
|
|
outputs = outputs - LLAMA_VOCAB_SIZE
|
|
outputs = outputs.cpu().tolist()
|
|
|
|
# Save all outputs for this prompt
|
|
successful_outputs = 0
|
|
prompt_files = []
|
|
|
|
for output_idx, midi_tokens in enumerate(outputs):
|
|
# 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 HuggingFace",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
Examples:
|
|
# Generate from a single prompt (there will be 4 outputs by default)
|
|
python generate_transformers.py --model path/to/checkpoint \\
|
|
--prompt "A cheerful piano melody"
|
|
|
|
# Generate single output without synthesis
|
|
python generate_transformers.py --model path/to/checkpoint \\
|
|
--prompt "A relaxing jazz piece" \\
|
|
--n_outputs 1 \\
|
|
--no-synthesize
|
|
|
|
# Interactive mode (with initial prompt)
|
|
python generate_transformers.py --model path/to/checkpoint \\
|
|
--prompt "A cheerful melody" \\
|
|
--interactive
|
|
|
|
# Interactive-only mode (no initial prompt)
|
|
python generate_transformers.py --model path/to/checkpoint \\
|
|
--interactive
|
|
|
|
# Generate from prompts file
|
|
python generate_transformers.py --model path/to/checkpoint \\
|
|
--prompts_file prompts.txt
|
|
"""
|
|
)
|
|
|
|
# Required arguments
|
|
parser.add_argument(
|
|
"--model",
|
|
type=str,
|
|
default="slseanwu/MIDI-LLM_Llama-3.2-1B",
|
|
help="Path to MIDI-LLM model checkpoint, can be HuggingFace model ID or local path (default: slseanwu/MIDI-LLM_Llama-3.2-1B)"
|
|
)
|
|
|
|
# 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(
|
|
"--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")
|
|
import sys
|
|
sys.exit(1)
|
|
|
|
# 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 from the model checkpoint
|
|
print("Loading tokenizer and model...")
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
args.model,
|
|
cache_dir=args.cache_dir,
|
|
pad_token="<|eot_id|>",
|
|
)
|
|
|
|
# Load model
|
|
model = prepare_hf_model(model_path=args.model)
|
|
|
|
# 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_hf(
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
prompts=prompts,
|
|
output_dir=output_dir,
|
|
model_path=args.model,
|
|
soundfont_path=args.soundfont if args.synthesize else None,
|
|
synthesize=args.synthesize,
|
|
temperature=args.temperature,
|
|
top_p=args.top_p,
|
|
max_tokens=args.max_tokens,
|
|
n_outputs=args.n_outputs
|
|
)
|
|
|
|
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,
|
|
}
|
|
}, 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_hf(
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
prompts=[user_prompt],
|
|
output_dir=output_dir,
|
|
model_path=args.model,
|
|
soundfont_path=args.soundfont if args.synthesize else None,
|
|
synthesize=args.synthesize,
|
|
temperature=args.temperature,
|
|
top_p=args.top_p,
|
|
max_tokens=args.max_tokens,
|
|
n_outputs=args.n_outputs
|
|
)
|
|
|
|
# Print input prompt
|
|
print(f"Input prompt: {user_prompt}")
|
|
|
|
# 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()
|