run_onnx_inference.py

import numpy as np
import onnxruntime as ort
import torch
import torch.nn.functional as F
from pathlib import Path

from src.s02_tokenizer import MusicTokenizer, BOS_TOKEN, EOS_TOKEN
from src.s06_generator import top_k_top_p_filter, apply_repetition_penalty

def generate_with_onnx(onnx_path: str, max_tokens: int = 256, temp: float = 0.85) -> list[int]:
    """Generates music tokens using the exported ONNX model."""
    
    # 1. Load ONNX model session
    # Selects CUDA (GPU) if available, falls back to CPU
    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
    print(f"Loading ONNX session from {onnx_path}...")
    session = ort.InferenceSession(onnx_path, providers=providers)
    print(f"Active Execution Providers: {session.get_providers()}")

    # 2. Setup prompt (starts with Beginning of Sequence token)
    generated = [BOS_TOKEN]
    
    # 3. Autoregressive loop
    print(f"Generating up to {max_tokens} tokens...")
    for step in range(max_tokens - len(generated)):
        # Format input data: must be int64 numpy array of shape (batch_size, sequence_length)
        input_data = np.array([generated], dtype=np.int64)
        
        # Run model inference
        # input_ids -> output name "logits"
        outputs = session.run(["logits"], {"input_ids": input_data})
        
        # Get logits for the last token position
        logits = outputs[0][0, -1, :] # shape: (vocab_size,)
        logits_tensor = torch.tensor(logits).unsqueeze(0) / max(temp, 1e-8)
        
        # Apply sampling filters
        logits_tensor = apply_repetition_penalty(logits_tensor, generated, penalty=1.15)
        logits_tensor = top_k_top_p_filter(logits_tensor, top_k=40, top_p=0.92)
        
        # Sample next token
        probs = F.softmax(logits_tensor, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1).item()
        
        generated.append(next_token)
        
        if next_token == EOS_TOKEN:
            break

    return generated

def main():
    onnx_path = "checkpoints/model.onnx"
    output_path = "output/onnx_generated.mid"
    
    # Generate tokens
    tokens = generate_with_onnx(onnx_path, max_tokens=512, temp=0.85)
    print(f"Generated {len(tokens)} tokens.")
    
    # Decode back to MIDI file
    tokenizer = MusicTokenizer()
    midi = tokenizer.tokens_to_midi(tokens)
    
    # Save output
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    midi.write(output_path)
    print(f"Saved ONNX generated MIDI to: {output_path}")

if __name__ == "__main__":
    main()