Upload generate.py

generate.py

@@ -0,0 +1,153 @@
+"""
+MuseMorphic Inference Script
+=============================
+
+Generate MIDI music from a trained MuseMorphic model.
+
+Usage:
+    python generate.py --checkpoint ./checkpoints/musemorphic_model.pt \
+                       --output generated.mid \
+                       --n_phrases 32 \
+                       --temperature 0.7
+"""
+
+import argparse
+import os
+import sys
+import torch
+import torch.nn.functional as F
+
+sys.path.insert(0, os.path.dirname(__file__))
+from musemorphic.model import MuseMorphicConfig, MuseMorphic, PhraseVAE, LatentMamba
+from musemorphic.tokenizer import REMIPlusTokenizer, notes_to_midi_file
+
+
+def load_model(checkpoint_path: str, device: torch.device):
+    """Load trained model from checkpoint."""
+    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    
+    config = MuseMorphicConfig(**ckpt['config'])
+    
+    vae = PhraseVAE(config).to(device)
+    mamba = LatentMamba(config).to(device)
+    
+    vae.load_state_dict(ckpt['vae_state_dict'])
+    mamba.load_state_dict(ckpt['mamba_state_dict'])
+    
+    model = MuseMorphic(config)
+    model.phrase_vae = vae
+    model.latent_mamba = mamba
+    model.to(device)
+    model.eval()
+    
+    return model, config
+
+
+def generate_midi(
+    model: MuseMorphic,
+    config: MuseMorphicConfig,
+    tokenizer: REMIPlusTokenizer,
+    n_phrases: int = 32,
+    temperature: float = 0.7,
+    max_decode_len: int = 128,
+    device: torch.device = torch.device('cpu'),
+) -> list:
+    """
+    Generate MIDI notes from the model.
+    
+    Returns list of note dicts: {pitch, start, duration, velocity}
+    """
+    with torch.no_grad():
+        # Stage 2: Generate latent phrase sequence
+        z_generated = model.latent_mamba.generate(
+            n_phrases=n_phrases,
+            temperature=temperature,
+            batch_size=1,
+        )
+        
+        # Stage 1: Decode each phrase latent to tokens
+        all_tokens = []
+        for t in range(z_generated.shape[1]):
+            z = z_generated[:, t]
+            
+            # Autoregressive decode
+            generated_ids = [config.bos_token_id]
+            
+            for _ in range(max_decode_len):
+                input_tensor = torch.tensor([generated_ids], dtype=torch.long, device=device)
+                logits = model.phrase_vae.decode(z, input_tensor)
+                
+                next_logits = logits[0, -1] / max(temperature, 0.1)
+                probs = F.softmax(next_logits, dim=-1)
+                next_token = torch.multinomial(probs, 1).item()
+                generated_ids.append(next_token)
+                
+                if next_token == config.eos_token_id:
+                    break
+            
+            phrase_tokens = tokenizer.decode(generated_ids)
+            all_tokens.extend(phrase_tokens)
+    
+    # Convert tokens to MIDI notes
+    notes = tokenizer.tokens_to_midi_notes(all_tokens)
+    return notes, all_tokens
+
+
+def main():
+    parser = argparse.ArgumentParser(description='MuseMorphic MIDI Generator')
+    parser.add_argument('--checkpoint', type=str, required=True, help='Path to model checkpoint')
+    parser.add_argument('--output', type=str, default='generated.mid', help='Output MIDI file path')
+    parser.add_argument('--n_phrases', type=int, default=32, help='Number of phrases to generate')
+    parser.add_argument('--temperature', type=float, default=0.7, help='Sampling temperature')
+    parser.add_argument('--tempo', type=float, default=120.0, help='Output MIDI tempo')
+    parser.add_argument('--device', type=str, default='auto', help='Device (auto/cuda/cpu)')
+    
+    args = parser.parse_args()
+    
+    # Device
+    if args.device == 'auto':
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    else:
+        device = torch.device(args.device)
+    
+    print(f'Device: {device}')
+    
+    # Load model
+    print(f'Loading model from {args.checkpoint}...')
+    model, config = load_model(args.checkpoint, device)
+    
+    params = sum(p.numel() for p in model.parameters())
+    print(f'Model parameters: {params:,} ({params/1e6:.2f}M)')
+    
+    # Tokenizer
+    tokenizer = REMIPlusTokenizer()
+    
+    # Generate
+    print(f'Generating {args.n_phrases} phrases at temperature {args.temperature}...')
+    notes, tokens = generate_midi(
+        model, config, tokenizer,
+        n_phrases=args.n_phrases,
+        temperature=args.temperature,
+        device=device,
+    )
+    
+    print(f'Generated {len(notes)} notes, {len(tokens)} tokens')
+    
+    if notes:
+        # Write MIDI
+        success = notes_to_midi_file(notes, args.output, tempo=args.tempo)
+        if success:
+            print(f'\n🎵 MIDI saved to: {args.output}')
+            total_dur = max(n['start'] + n['duration'] for n in notes)
+            print(f'  Duration: ~{total_dur/480:.1f} beats ({total_dur/480/args.tempo*60:.1f} seconds at {args.tempo} BPM)')
+            pitches = [n['pitch'] for n in notes]
+            print(f'  Pitch range: {min(pitches)}-{max(pitches)}')
+            print(f'  Note count: {len(notes)}')
+        else:
+            print('Failed to write MIDI. Install midiutil: pip install midiutil')
+    else:
+        print('No notes generated. Model may need more training.')
+
+
+if __name__ == '__main__':
+    main()