File size: 3,496 Bytes
30e9297 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | """
XML Configuration-driven Inference for AriaLM.
Allows specifying prompt notes, generation length, and sampling settings in an XML file.
"""
import logging
import xml.etree.ElementTree as ET
from pathlib import Path
import pretty_midi
import torch
from src.s01_config import GenConfig, ModelConfig, PathConfig
from src.s02_tokenizer import MusicTokenizer
from src.s04_model import MusicTransformer
from src.s06_generator import generate
logger = logging.getLogger(__name__)
def parse_xml_config(xml_path: Path) -> tuple[GenConfig, list[int]]:
"""
Parse generation settings and starting prompt notes from an XML file.
"""
tree = ET.parse(xml_path)
root = tree.getroot()
# 1. Parse Settings
settings = root.find("settings")
temp = float(settings.find("temperature").text or 0.85)
top_k = int(settings.find("top_k").text or 40)
top_p = float(settings.find("top_p").text or 0.92)
max_tokens = int(settings.find("max_tokens").text or 512)
rep_penalty = float(settings.find("repetition_penalty").text or 1.15)
seed = int(settings.find("seed").text or 42)
gen_config = GenConfig(
temperature=temp,
top_k=top_k,
top_p=top_p,
max_tokens=max_tokens,
repetition_penalty=rep_penalty,
seed=seed,
)
# 2. Parse Prompt Notes
prompt_node = root.find("prompt")
prompt_tokens = []
if prompt_node is not None:
tokenizer = MusicTokenizer()
pm = pretty_midi.PrettyMIDI()
instrument = pretty_midi.Instrument(program=0) # Default piano
current_time = 0.0
for note_el in prompt_node.findall("note"):
pitch = int(note_el.attrib["pitch"])
velocity = int(note_el.attrib["velocity"])
duration = float(note_el.attrib["duration_ms"]) / 1000.0
delay = float(note_el.attrib.get("delay_ms", 0)) / 1000.0
current_time += delay
note = pretty_midi.Note(
velocity=velocity,
pitch=pitch,
start=current_time,
end=current_time + duration,
)
instrument.notes.append(note)
current_time += duration
pm.instruments.append(instrument)
# Convert prompt MIDI to token IDs
prompt_tokens = tokenizer.midi_to_tokens(pm)
return gen_config, prompt_tokens
def generate_from_xml(xml_path: Path, model_path: Path, output_path: Path):
"""Run generation using parameters specified in an XML file."""
# Load config and prompt
gen_config, prompt_tokens = parse_xml_config(xml_path)
# Load Model
tokenizer = MusicTokenizer()
model_config = ModelConfig(vocab_size=tokenizer.vocab_size)
model = MusicTransformer.from_config(model_config)
ckpt = torch.load(model_path, map_location="cpu", weights_only=False)
model.load_state_dict(ckpt["model_state_dict"])
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
print(f"Generating from XML Config: {xml_path}")
print(f"Sampling details: Temp={gen_config.temperature}, Seed={gen_config.seed}, Length={gen_config.max_tokens}")
# Generate
tokens = generate(model, tokenizer, gen_config, prompt_tokens, device=device)
# Save output
midi = tokenizer.tokens_to_midi(tokens)
output_path.parent.mkdir(parents=True, exist_ok=True)
midi.write(str(output_path))
print(f"Saved generated MIDI to: {output_path}")
|