add HF inference script, enrich readme w/ egs
This commit is contained in:
slSeanWU
@@ -0,0 +1,474 @@
|
||||
#!/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...")
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user