Upload run_onnx_inference.py with huggingface_hub

run_onnx_inference.py

@@ -0,0 +1,70 @@
+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()