Upload MuseMorphic_Training.ipynb

MuseMorphic_Training.ipynb

@@ -0,0 +1,817 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "T4"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# 🎵 MuseMorphic: Lightweight MIDI Generator\n",
+        "\n",
+        "**A novel consumer-grade architecture for controllable, infinite-length MIDI generation.**\n",
+        "\n",
+        "Key features:\n",
+        "- **~33M parameters** — trains on free Colab T4, inference <1GB VRAM\n",
+        "- **O(n) complexity** — Mamba SSM backbone, no quadratic attention\n",
+        "- **Two-stage hierarchical** — PhraseVAE (compress) + LatentMamba (generate)\n",
+        "- **Music-native** — FME embeddings with harmonic awareness\n",
+        "- **Controllable** — tempo, key, density, style conditioning\n",
+        "- **Infinite generation** — fixed-size recurrent state, no memory growth\n",
+        "- **Training-stable by design** — σReparam + ZClip + Pre-LN + BF16\n",
+        "\n",
+        "📄 [Architecture Paper/README](https://huggingface.co/asdf98/MuseMorphic)\n",
+        "\n",
+        "---\n",
+        "\n",
+        "## Setup\n",
+        "\n",
+        "Run this cell first to install all dependencies."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ============================================================\n",
+        "# 1. Install Dependencies\n",
+        "# ============================================================\n",
+        "!pip install -q torch torchvision torchaudio\n",
+        "!pip install -q einops datasets pretty_midi midiutil\n",
+        "!pip install -q huggingface_hub\n",
+        "\n",
+        "# Clone MuseMorphic repo\n",
+        "!git clone https://huggingface.co/asdf98/MuseMorphic /content/MuseMorphic 2>/dev/null || (cd /content/MuseMorphic && git pull)\n",
+        "\n",
+        "import sys\n",
+        "sys.path.insert(0, '/content/MuseMorphic/musemorphic')\n",
+        "\n",
+        "print('✅ Dependencies installed!')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ============================================================\n",
+        "# 2. Check GPU & Hardware\n",
+        "# ============================================================\n",
+        "import torch\n",
+        "import os\n",
+        "\n",
+        "print(f'PyTorch version: {torch.__version__}')\n",
+        "print(f'CUDA available: {torch.cuda.is_available()}')\n",
+        "if torch.cuda.is_available():\n",
+        "    print(f'GPU: {torch.cuda.get_device_name(0)}')\n",
+        "    print(f'VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB')\n",
+        "    print(f'BF16 support: {torch.cuda.is_bf16_supported()}')\n",
+        "else:\n",
+        "    print('⚠️ No GPU detected. Training will be slow but functional on CPU.')\n",
+        "\n",
+        "# Auto-detect best dtype\n",
+        "if torch.cuda.is_available() and torch.cuda.is_bf16_supported():\n",
+        "    DTYPE = 'bf16'\n",
+        "    print('\\n✅ Using BFloat16 (optimal: no loss scaling needed)')\n",
+        "elif torch.cuda.is_available():\n",
+        "    DTYPE = 'fp16'\n",
+        "    print('\\n⚠️ Using Float16 (T4 GPU — BF16 not supported, using FP16 with gradient scaling)')\n",
+        "else:\n",
+        "    DTYPE = 'fp32'\n",
+        "    print('\\n📋 Using Float32 (CPU mode)')"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Configuration\n",
+        "\n",
+        "Adjust these settings based on your GPU and desired output quality."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ============================================================\n",
+        "# 3. Configuration\n",
+        "# ============================================================\n",
+        "\n",
+        "# ---- Model Size Presets ----\n",
+        "# 'tiny'  : ~8M params  — Fast experiments, lower quality\n",
+        "# 'small' : ~33M params — Default, good quality (recommended for Colab T4)\n",
+        "# 'medium': ~65M params — Better quality, needs more VRAM\n",
+        "\n",
+        "MODEL_SIZE = 'small'  # @param ['tiny', 'small', 'medium']\n",
+        "\n",
+        "# ---- Dataset ----\n",
+        "# 'auto'     : Auto-select best available\n",
+        "# 'maestro'  : Classical piano (MAESTRO)\n",
+        "# 'synthetic': Generated data (for testing)\n",
+        "DATASET = 'auto'  # @param ['auto', 'maestro', 'synthetic']\n",
+        "MAX_PIECES = 500  # @param {type: 'integer'}\n",
+        "\n",
+        "# ---- Training ----\n",
+        "VAE_EPOCHS = 15      # @param {type: 'integer'}\n",
+        "MAMBA_EPOCHS = 30    # @param {type: 'integer'}\n",
+        "BATCH_SIZE = 32      # @param {type: 'integer'}\n",
+        "LEARNING_RATE = 3e-4 # @param {type: 'number'}\n",
+        "\n",
+        "# ---- Output ----\n",
+        "OUTPUT_DIR = '/content/checkpoints'  # @param {type: 'string'}\n",
+        "PUSH_TO_HUB = False   # @param {type: 'boolean'}\n",
+        "HUB_MODEL_ID = ''     # @param {type: 'string'}\n",
+        "\n",
+        "print(f'Model size: {MODEL_SIZE}')\n",
+        "print(f'Dataset: {DATASET}')\n",
+        "print(f'Training: VAE {VAE_EPOCHS}ep + Mamba {MAMBA_EPOCHS}ep, batch={BATCH_SIZE}')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ============================================================\n",
+        "# 4. Build Model Configuration\n",
+        "# ============================================================\n",
+        "from model import MuseMorphicConfig, MuseMorphic, model_summary\n",
+        "\n",
+        "# Model size presets\n",
+        "SIZE_CONFIGS = {\n",
+        "    'tiny': MuseMorphicConfig(\n",
+        "        d_model=128,\n",
+        "        vae_encoder_layers=2,\n",
+        "        vae_decoder_layers=2,\n",
+        "        vae_n_heads=4,\n",
+        "        vae_d_ff=256,\n",
+        "        latent_dim=32,\n",
+        "        mamba_d_model=128,\n",
+        "        mamba_n_layers=4,\n",
+        "        mamba_d_state=8,\n",
+        "        mamba_expand=2,\n",
+        "    ),\n",
+        "    'small': MuseMorphicConfig(\n",
+        "        d_model=256,\n",
+        "        vae_encoder_layers=3,\n",
+        "        vae_decoder_layers=3,\n",
+        "        vae_n_heads=4,\n",
+        "        vae_d_ff=512,\n",
+        "        latent_dim=64,\n",
+        "        mamba_d_model=256,\n",
+        "        mamba_n_layers=8,\n",
+        "        mamba_d_state=16,\n",
+        "        mamba_expand=2,\n",
+        "    ),\n",
+        "    'medium': MuseMorphicConfig(\n",
+        "        d_model=384,\n",
+        "        vae_encoder_layers=4,\n",
+        "        vae_decoder_layers=4,\n",
+        "        vae_n_heads=6,\n",
+        "        vae_d_ff=768,\n",
+        "        latent_dim=96,\n",
+        "        mamba_d_model=384,\n",
+        "        mamba_n_layers=12,\n",
+        "        mamba_d_state=16,\n",
+        "        mamba_expand=2,\n",
+        "    ),\n",
+        "}\n",
+        "\n",
+        "config = SIZE_CONFIGS[MODEL_SIZE]\n",
+        "model = model_summary(config)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Data Preparation\n",
+        "\n",
+        "Automatically downloads and preprocesses MIDI data."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ============================================================\n",
+        "# 5. Load & Preprocess Dataset\n",
+        "# ============================================================\n",
+        "import logging\n",
+        "logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')\n",
+        "\n",
+        "from data_pipeline import prepare_training_data, auto_select_dataset, load_dataset_notes\n",
+        "from data_pipeline import preprocess_dataset, _generate_synthetic_dataset\n",
+        "from tokenizer import REMIPlusTokenizer\n",
+        "\n",
+        "# Select dataset\n",
+        "if DATASET == 'auto':\n",
+        "    dataset_name = auto_select_dataset()\n",
+        "elif DATASET == 'maestro':\n",
+        "    dataset_name = 'maestro_v1_sustain'\n",
+        "elif DATASET == 'synthetic':\n",
+        "    dataset_name = None\n",
+        "else:\n",
+        "    dataset_name = DATASET\n",
+        "\n",
+        "# Load and preprocess\n",
+        "tokenizer = REMIPlusTokenizer()\n",
+        "\n",
+        "if dataset_name is not None:\n",
+        "    try:\n",
+        "        pieces = load_dataset_notes(dataset_name, max_pieces=MAX_PIECES)\n",
+        "        print(f'✅ Loaded {len(pieces)} pieces from {dataset_name}')\n",
+        "    except Exception as e:\n",
+        "        print(f'⚠️ Failed to load {dataset_name}: {e}')\n",
+        "        print('Falling back to synthetic data...')\n",
+        "        pieces = _generate_synthetic_dataset(MAX_PIECES)\n",
+        "else:\n",
+        "    pieces = _generate_synthetic_dataset(MAX_PIECES)\n",
+        "    print(f'✅ Generated {len(pieces)} synthetic pieces')\n",
+        "\n",
+        "# Preprocess\n",
+        "phrases, controls = preprocess_dataset(pieces, tokenizer, max_phrase_len=config.vae_max_seq_len)\n",
+        "print(f'\\n📊 Dataset Summary:')\n",
+        "print(f'  Total phrases: {len(phrases)}')\n",
+        "print(f'  Avg phrase length: {sum(len(p) for p in phrases)/len(phrases):.1f} tokens')\n",
+        "print(f'  Vocab size: {tokenizer.vocab_size}')\n",
+        "print(f'  Sample phrase (first 10 tokens): {phrases[0][:10]}')"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Training\n",
+        "\n",
+        "Two-stage training with curriculum:\n",
+        "\n",
+        "**Stage 1 — PhraseVAE** (compress phrases to latent vectors):\n",
+        "- 1a. Span-infilling pretraining (learn REMI grammar)\n",
+        "- 1b. Autoencoder (pure reconstruction, KL=0)\n",
+        "- 1c. VAE fine-tuning (KL weight = 0.01)\n",
+        "\n",
+        "**Stage 2 — LatentMamba** (generate latent sequences):\n",
+        "- Predict next phrase latent from history, O(n) complexity"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ============================================================\n",
+        "# 6. Training — Stage 1: PhraseVAE\n",
+        "# ============================================================\n",
+        "import time\n",
+        "import random\n",
+        "import math\n",
+        "import numpy as np\n",
+        "import torch\n",
+        "import torch.nn.functional as F\n",
+        "from torch.utils.data import DataLoader, Dataset\n",
+        "from model import PhraseVAE, ZClip\n",
+        "\n",
+        "# Set seed\n",
+        "SEED = 42\n",
+        "random.seed(SEED)\n",
+        "np.random.seed(SEED)\n",
+        "torch.manual_seed(SEED)\n",
+        "if torch.cuda.is_available():\n",
+        "    torch.cuda.manual_seed_all(SEED)\n",
+        "\n",
+        "# Device & dtype\n",
+        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
+        "if DTYPE == 'bf16' and torch.cuda.is_available() and torch.cuda.is_bf16_supported():\n",
+        "    amp_dtype = torch.bfloat16\n",
+        "elif DTYPE == 'fp16' and torch.cuda.is_available():\n",
+        "    amp_dtype = torch.float16\n",
+        "else:\n",
+        "    amp_dtype = torch.float32\n",
+        "\n",
+        "# Dataset\n",
+        "class PhraseDS(Dataset):\n",
+        "    def __init__(self, phrases, max_len, pad_id=0):\n",
+        "        self.phrases = phrases\n",
+        "        self.max_len = max_len\n",
+        "        self.pad_id = pad_id\n",
+        "    def __len__(self):\n",
+        "        return len(self.phrases)\n",
+        "    def __getitem__(self, idx):\n",
+        "        ids = self.phrases[idx][:self.max_len]\n",
+        "        padded = ids + [self.pad_id] * (self.max_len - len(ids))\n",
+        "        return torch.tensor(padded, dtype=torch.long)\n",
+        "\n",
+        "train_ds = PhraseDS(phrases, config.vae_max_seq_len)\n",
+        "train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, \n",
+        "                          num_workers=2, pin_memory=True, drop_last=True)\n",
+        "\n",
+        "# Create VAE\n",
+        "vae = PhraseVAE(config).to(device)\n",
+        "vae_params = sum(p.numel() for p in vae.parameters())\n",
+        "print(f'PhraseVAE parameters: {vae_params:,} ({vae_params/1e6:.2f}M)')\n",
+        "\n",
+        "# Optimizer\n",
+        "optimizer = torch.optim.AdamW(vae.parameters(), lr=LEARNING_RATE, weight_decay=0.01)\n",
+        "scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=500, T_mult=2, eta_min=1e-6)\n",
+        "zclip = ZClip(z_thresh=2.5)\n",
+        "\n",
+        "# FP16 needs GradScaler, BF16 does not\n",
+        "use_scaler = (amp_dtype == torch.float16)\n",
+        "scaler = torch.amp.GradScaler() if use_scaler else None\n",
+        "\n",
+        "# Span masking helper\n",
+        "def apply_span_mask(token_ids, mask_prob=0.15, mask_id=3, span_len=3):\n",
+        "    masked = token_ids.clone()\n",
+        "    B, L = masked.shape\n",
+        "    for b in range(B):\n",
+        "        n_masks = max(1, int(L * mask_prob / span_len))\n",
+        "        for _ in range(n_masks):\n",
+        "            start = random.randint(1, max(1, L - span_len - 1))\n",
+        "            end = min(start + span_len, L)\n",
+        "            masked[b, start:end] = mask_id\n",
+        "    return masked\n",
+        "\n",
+        "# ---- Training Loop ----\n",
+        "print('\\n' + '='*60)\n",
+        "print('Starting PhraseVAE Training')\n",
+        "print('='*60)\n",
+        "\n",
+        "# Compute total epochs\n",
+        "pretrain_epochs = max(1, VAE_EPOCHS // 5)\n",
+        "ae_epochs = max(1, VAE_EPOCHS * 3 // 5)\n",
+        "vae_epochs = max(1, VAE_EPOCHS - pretrain_epochs - ae_epochs)\n",
+        "\n",
+        "stages = [\n",
+        "    ('1a-Pretrain', pretrain_epochs, 0.0, True),\n",
+        "    ('1b-AE', ae_epochs, 0.0, False),\n",
+        "    ('1c-VAE', vae_epochs, 0.01, False),\n",
+        "]\n",
+        "\n",
+        "global_step = 0\n",
+        "history = {'loss': [], 'recon': [], 'kl': []}\n",
+        "\n",
+        "for stage_name, n_epochs, kl_weight, use_masking in stages:\n",
+        "    print(f'\\n--- Stage {stage_name} ({n_epochs} epochs, KL={kl_weight}) ---')\n",
+        "    \n",
+        "    # Lower LR for VAE fine-tuning stage\n",
+        "    if stage_name == '1c-VAE':\n",
+        "        for pg in optimizer.param_groups:\n",
+        "            pg['lr'] = LEARNING_RATE * 0.1\n",
+        "    \n",
+        "    for epoch in range(n_epochs):\n",
+        "        vae.train()\n",
+        "        epoch_loss = 0\n",
+        "        n_batches = 0\n",
+        "        t0 = time.time()\n",
+        "        \n",
+        "        for batch in train_loader:\n",
+        "            token_ids = batch.to(device)\n",
+        "            \n",
+        "            # Apply masking for pretraining\n",
+        "            if use_masking:\n",
+        "                input_ids = apply_span_mask(token_ids)\n",
+        "            else:\n",
+        "                input_ids = token_ids\n",
+        "            \n",
+        "            optimizer.zero_grad()\n",
+        "            \n",
+        "            with torch.autocast(device_type=device.type, dtype=amp_dtype):\n",
+        "                outputs = vae(input_ids, target_tokens=token_ids, kl_weight=kl_weight)\n",
+        "                loss = outputs['loss']\n",
+        "            \n",
+        "            # NaN check\n",
+        "            if torch.isnan(loss) or torch.isinf(loss):\n",
+        "                print(f'⚠️ NaN/Inf at step {global_step}! Skipping...')\n",
+        "                optimizer.zero_grad()\n",
+        "                continue\n",
+        "            \n",
+        "            if use_scaler:\n",
+        "                scaler.scale(loss).backward()\n",
+        "                scaler.unscale_(optimizer)\n",
+        "                zclip(vae)\n",
+        "                scaler.step(optimizer)\n",
+        "                scaler.update()\n",
+        "            else:\n",
+        "                loss.backward()\n",
+        "                zclip(vae)\n",
+        "                optimizer.step()\n",
+        "            \n",
+        "            scheduler.step()\n",
+        "            \n",
+        "            epoch_loss += loss.item()\n",
+        "            n_batches += 1\n",
+        "            global_step += 1\n",
+        "            \n",
+        "            # Log\n",
+        "            history['loss'].append(loss.item())\n",
+        "            history['recon'].append(outputs['recon_loss'].item())\n",
+        "            history['kl'].append(outputs['kl_loss'].item())\n",
+        "        \n",
+        "        elapsed = time.time() - t0\n",
+        "        avg_loss = epoch_loss / max(n_batches, 1)\n",
+        "        lr = optimizer.param_groups[0]['lr']\n",
+        "        print(f'  Epoch {epoch+1}/{n_epochs} | Loss: {avg_loss:.4f} | '\n",
+        "              f'Recon: {outputs[\"recon_loss\"].item():.4f} | '\n",
+        "              f'KL: {outputs[\"kl_loss\"].item():.4f} | '\n",
+        "              f'LR: {lr:.2e} | Time: {elapsed:.1f}s')\n",
+        "\n",
+        "print(f'\\n✅ PhraseVAE training complete! ({global_step} total steps)')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ============================================================\n",
+        "# 7. Training — Stage 2: LatentMamba\n",
+        "# ============================================================\n",
+        "from model import LatentMamba\n",
+        "\n",
+        "# Freeze VAE encoder\n",
+        "vae.eval()\n",
+        "for p in vae.parameters():\n",
+        "    p.requires_grad = False\n",
+        "\n",
+        "# Encode all phrases to latent space\n",
+        "print('Encoding phrases to latent space...')\n",
+        "all_latents = []\n",
+        "encode_loader = DataLoader(train_ds, batch_size=64, shuffle=False, num_workers=2)\n",
+        "\n",
+        "with torch.no_grad():\n",
+        "    for batch in encode_loader:\n",
+        "        token_ids = batch.to(device)\n",
+        "        with torch.autocast(device_type=device.type, dtype=amp_dtype):\n",
+        "            z, _, _ = vae.encode(token_ids)\n",
+        "        all_latents.append(z.cpu())\n",
+        "\n",
+        "all_z = torch.cat(all_latents, dim=0)\n",
+        "print(f'Encoded {all_z.shape[0]} phrases to latent dim {all_z.shape[1]}')\n",
+        "\n",
+        "# Create latent sequences (group phrases into chunks)\n",
+        "SEQ_LEN = min(64, len(all_z) // 4)\n",
+        "latent_seqs = []\n",
+        "for i in range(0, len(all_z) - SEQ_LEN, SEQ_LEN // 2):  # 50% overlap\n",
+        "    latent_seqs.append(all_z[i:i+SEQ_LEN])\n",
+        "\n",
+        "print(f'Created {len(latent_seqs)} latent sequences of length {SEQ_LEN}')\n",
+        "\n",
+        "class LatentDS(Dataset):\n",
+        "    def __init__(self, seqs):\n",
+        "        self.seqs = seqs\n",
+        "    def __len__(self):\n",
+        "        return len(self.seqs)\n",
+        "    def __getitem__(self, idx):\n",
+        "        return self.seqs[idx]\n",
+        "\n",
+        "latent_ds = LatentDS(latent_seqs)\n",
+        "latent_loader = DataLoader(latent_ds, batch_size=min(BATCH_SIZE, len(latent_seqs)),\n",
+        "                           shuffle=True, drop_last=True)\n",
+        "\n",
+        "# Create LatentMamba\n",
+        "mamba = LatentMamba(config).to(device)\n",
+        "mamba_params = sum(p.numel() for p in mamba.parameters())\n",
+        "print(f'LatentMamba parameters: {mamba_params:,} ({mamba_params/1e6:.2f}M)')\n",
+        "\n",
+        "# Optimizer\n",
+        "mamba_optimizer = torch.optim.AdamW(mamba.parameters(), lr=LEARNING_RATE * 0.5, weight_decay=0.01)\n",
+        "mamba_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n",
+        "    mamba_optimizer, T_0=300, T_mult=2, eta_min=1e-6)\n",
+        "mamba_zclip = ZClip(z_thresh=2.5)\n",
+        "\n",
+        "mamba_scaler = torch.amp.GradScaler() if use_scaler else None\n",
+        "\n",
+        "# Training loop\n",
+        "print('\\n' + '='*60)\n",
+        "print('Starting LatentMamba Training')\n",
+        "print('='*60)\n",
+        "\n",
+        "mamba_history = {'mse': [], 'cos': []}\n",
+        "\n",
+        "for epoch in range(MAMBA_EPOCHS):\n",
+        "    mamba.train()\n",
+        "    epoch_loss = 0\n",
+        "    n_batches = 0\n",
+        "    t0 = time.time()\n",
+        "    \n",
+        "    for batch in latent_loader:\n",
+        "        z_seq = batch.to(device)\n",
+        "        z_input = z_seq[:, :-1]\n",
+        "        z_target = z_seq[:, 1:]\n",
+        "        \n",
+        "        mamba_optimizer.zero_grad()\n",
+        "        \n",
+        "        with torch.autocast(device_type=device.type, dtype=amp_dtype):\n",
+        "            z_pred = mamba(z_input)\n",
+        "            mse_loss = F.mse_loss(z_pred, z_target)\n",
+        "            cos_loss = 1.0 - F.cosine_similarity(\n",
+        "                z_pred.reshape(-1, z_pred.shape[-1]),\n",
+        "                z_target.reshape(-1, z_target.shape[-1]), dim=-1\n",
+        "            ).mean()\n",
+        "            loss = mse_loss + 0.1 * cos_loss\n",
+        "        \n",
+        "        if torch.isnan(loss) or torch.isinf(loss):\n",
+        "            print(f'⚠️ NaN/Inf at epoch {epoch}! Skipping...')\n",
+        "            mamba_optimizer.zero_grad()\n",
+        "            continue\n",
+        "        \n",
+        "        if use_scaler:\n",
+        "            mamba_scaler.scale(loss).backward()\n",
+        "            mamba_scaler.unscale_(mamba_optimizer)\n",
+        "            mamba_zclip(mamba)\n",
+        "            mamba_scaler.step(mamba_optimizer)\n",
+        "            mamba_scaler.update()\n",
+        "        else:\n",
+        "            loss.backward()\n",
+        "            mamba_zclip(mamba)\n",
+        "            mamba_optimizer.step()\n",
+        "        \n",
+        "        mamba_scheduler.step()\n",
+        "        \n",
+        "        epoch_loss += loss.item()\n",
+        "        n_batches += 1\n",
+        "        \n",
+        "        mamba_history['mse'].append(mse_loss.item())\n",
+        "        mamba_history['cos'].append(cos_loss.item())\n",
+        "    \n",
+        "    elapsed = time.time() - t0\n",
+        "    avg_loss = epoch_loss / max(n_batches, 1)\n",
+        "    lr = mamba_optimizer.param_groups[0]['lr']\n",
+        "    \n",
+        "    if (epoch + 1) % 5 == 0 or epoch == 0:\n",
+        "        print(f'  Epoch {epoch+1}/{MAMBA_EPOCHS} | Loss: {avg_loss:.6f} | '\n",
+        "              f'MSE: {mse_loss.item():.6f} | Cos: {cos_loss.item():.4f} | '\n",
+        "              f'LR: {lr:.2e} | Time: {elapsed:.1f}s')\n",
+        "\n",
+        "print(f'\\n✅ LatentMamba training complete!')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ============================================================\n",
+        "# 8. Plot Training Curves\n",
+        "# ============================================================\n",
+        "import matplotlib.pyplot as plt\n",
+        "\n",
+        "fig, axes = plt.subplots(1, 3, figsize=(18, 5))\n",
+        "\n",
+        "# VAE Loss\n",
+        "axes[0].plot(history['loss'], alpha=0.3, color='blue')\n",
+        "window = min(50, len(history['loss']) // 5) if len(history['loss']) > 10 else 1\n",
+        "if window > 1:\n",
+        "    smoothed = np.convolve(history['loss'], np.ones(window)/window, mode='valid')\n",
+        "    axes[0].plot(smoothed, color='blue', linewidth=2)\n",
+        "axes[0].set_title('PhraseVAE Loss')\n",
+        "axes[0].set_xlabel('Step')\n",
+        "axes[0].set_ylabel('Loss')\n",
+        "axes[0].grid(True, alpha=0.3)\n",
+        "\n",
+        "# VAE KL\n",
+        "axes[1].plot(history['kl'], alpha=0.5, color='red')\n",
+        "axes[1].set_title('KL Divergence')\n",
+        "axes[1].set_xlabel('Step')\n",
+        "axes[1].set_ylabel('KL')\n",
+        "axes[1].grid(True, alpha=0.3)\n",
+        "\n",
+        "# Mamba Loss\n",
+        "axes[2].plot(mamba_history['mse'], alpha=0.3, color='green')\n",
+        "window = min(50, len(mamba_history['mse']) // 5) if len(mamba_history['mse']) > 10 else 1\n",
+        "if window > 1:\n",
+        "    smoothed = np.convolve(mamba_history['mse'], np.ones(window)/window, mode='valid')\n",
+        "    axes[2].plot(smoothed, color='green', linewidth=2)\n",
+        "axes[2].set_title('LatentMamba MSE Loss')\n",
+        "axes[2].set_xlabel('Step')\n",
+        "axes[2].set_ylabel('MSE')\n",
+        "axes[2].grid(True, alpha=0.3)\n",
+        "\n",
+        "plt.tight_layout()\n",
+        "plt.savefig('/content/training_curves.png', dpi=150)\n",
+        "plt.show()\n",
+        "print('📊 Training curves saved to /content/training_curves.png')"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Generation\n",
+        "\n",
+        "Generate MIDI music using the trained model!"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ============================================================\n",
+        "# 9. Generate Music!\n",
+        "# ============================================================\n",
+        "from model import MuseMorphic\n",
+        "from tokenizer import REMIPlusTokenizer, notes_to_midi_file\n",
+        "\n",
+        "# Assemble full model\n",
+        "full_model = MuseMorphic(config).to(device)\n",
+        "full_model.phrase_vae = vae\n",
+        "full_model.latent_mamba = mamba\n",
+        "full_model.eval()\n",
+        "\n",
+        "# Unfreeze VAE for generation (was frozen for Stage 2)\n",
+        "# (No gradient computation needed for generation anyway)\n",
+        "\n",
+        "# Generation settings\n",
+        "N_PHRASES = 16    # @param {type: 'integer'}\n",
+        "TEMPERATURE = 0.7 # @param {type: 'number'}\n",
+        "\n",
+        "print(f'Generating {N_PHRASES} phrases at temperature {TEMPERATURE}...')\n",
+        "\n",
+        "with torch.no_grad():\n",
+        "    # Generate latent sequence\n",
+        "    z_generated = mamba.generate(\n",
+        "        n_phrases=N_PHRASES,\n",
+        "        temperature=TEMPERATURE,\n",
+        "        batch_size=1,\n",
+        "    )\n",
+        "    print(f'Generated latent shape: {z_generated.shape}')\n",
+        "    \n",
+        "    # Decode each phrase latent to tokens\n",
+        "    all_tokens = []\n",
+        "    for t in range(z_generated.shape[1]):\n",
+        "        z = z_generated[:, t]  # (1, latent_dim)\n",
+        "        \n",
+        "        # Autoregressive decode\n",
+        "        generated_ids = [config.bos_token_id]\n",
+        "        max_decode_len = 128\n",
+        "        \n",
+        "        for _ in range(max_decode_len):\n",
+        "            input_tensor = torch.tensor([generated_ids], dtype=torch.long, device=device)\n",
+        "            with torch.autocast(device_type=device.type, dtype=amp_dtype):\n",
+        "                logits = vae.decode(z, input_tensor)\n",
+        "            \n",
+        "            next_logits = logits[0, -1] / max(TEMPERATURE, 0.1)\n",
+        "            probs = F.softmax(next_logits, dim=-1)\n",
+        "            next_token = torch.multinomial(probs, 1).item()\n",
+        "            generated_ids.append(next_token)\n",
+        "            \n",
+        "            if next_token == config.eos_token_id:\n",
+        "                break\n",
+        "        \n",
+        "        phrase_tokens = tokenizer.decode(generated_ids)\n",
+        "        all_tokens.extend(phrase_tokens)\n",
+        "\n",
+        "print(f'\\nGenerated {len(all_tokens)} REMI+ tokens')\n",
+        "print(f'Sample tokens: {all_tokens[:20]}')\n",
+        "\n",
+        "# Convert to MIDI notes\n",
+        "notes = tokenizer.tokens_to_midi_notes(all_tokens)\n",
+        "print(f'Extracted {len(notes)} notes')\n",
+        "\n",
+        "if notes:\n",
+        "    # Write MIDI file\n",
+        "    output_midi = '/content/generated_music.mid'\n",
+        "    success = notes_to_midi_file(notes, output_midi)\n",
+        "    if success:\n",
+        "        print(f'\\n🎵 MIDI file saved to: {output_midi}')\n",
+        "        print(f'  Notes: {len(notes)}')\n",
+        "        if notes:\n",
+        "            total_duration = max(n[\"start\"] + n[\"duration\"] for n in notes)\n",
+        "            print(f'  Duration: ~{total_duration/480:.1f} beats')\n",
+        "            pitches = [n[\"pitch\"] for n in notes]\n",
+        "            print(f'  Pitch range: {min(pitches)}-{max(pitches)}')\n",
+        "    else:\n",
+        "        print('⚠️ MIDI writing failed. Install midiutil: pip install midiutil')\n",
+        "else:\n",
+        "    print('⚠️ No notes generated. Try training for more epochs or adjusting temperature.')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ============================================================\n",
+        "# 10. Save Model\n",
+        "# ============================================================\n",
+        "import os\n",
+        "from dataclasses import asdict\n",
+        "\n",
+        "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
+        "\n",
+        "# Save model\n",
+        "save_path = os.path.join(OUTPUT_DIR, 'musemorphic_model.pt')\n",
+        "torch.save({\n",
+        "    'vae_state_dict': vae.state_dict(),\n",
+        "    'mamba_state_dict': mamba.state_dict(),\n",
+        "    'config': asdict(config),\n",
+        "    'training_history': {\n",
+        "        'vae': history,\n",
+        "        'mamba': mamba_history,\n",
+        "    }\n",
+        "}, save_path)\n",
+        "print(f'✅ Model saved to {save_path}')\n",
+        "print(f'  File size: {os.path.getsize(save_path)/1e6:.1f} MB')\n",
+        "\n",
+        "# Save tokenizer\n",
+        "tokenizer.save(os.path.join(OUTPUT_DIR, 'tokenizer'))\n",
+        "print(f'✅ Tokenizer saved')\n",
+        "\n",
+        "# Optional: Push to HF Hub\n",
+        "if PUSH_TO_HUB and HUB_MODEL_ID:\n",
+        "    from huggingface_hub import HfApi\n",
+        "    api = HfApi()\n",
+        "    api.upload_folder(\n",
+        "        folder_path=OUTPUT_DIR,\n",
+        "        repo_id=HUB_MODEL_ID,\n",
+        "        repo_type='model',\n",
+        "    )\n",
+        "    print(f'✅ Pushed to https://huggingface.co/{HUB_MODEL_ID}')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ============================================================\n",
+        "# 11. Listen to Generated MIDI (in Colab)\n",
+        "# ============================================================\n",
+        "try:\n",
+        "    from IPython.display import Audio, display\n",
+        "    import pretty_midi\n",
+        "    \n",
+        "    # Synthesize MIDI to audio using FluidSynth\n",
+        "    pm = pretty_midi.PrettyMIDI('/content/generated_music.mid')\n",
+        "    audio = pm.fluidsynth(fs=22050)\n",
+        "    \n",
+        "    print('🎧 Listen to generated music:')\n",
+        "    display(Audio(audio, rate=22050))\n",
+        "except Exception as e:\n",
+        "    print(f'Audio playback not available: {e}')\n",
+        "    print('Download the MIDI file and play it in any MIDI player.')\n",
+        "    print('File: /content/generated_music.mid')"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "---\n",
+        "\n",
+        "## Architecture Summary\n",
+        "\n",
+        "### Novel Contributions\n",
+        "\n",
+        "1. **First SSM-based latent music generator**: Mamba operating on compressed phrase latents\n",
+        "2. **FME with log-frequency encoding**: Physics-aware embeddings respecting harmonic series\n",
+        "3. **Multi-attribute control via latent conditioning**: Tempo, key, density, style\n",
+        "4. **Guaranteed training stability stack**: σReparam + ZClip + Pre-LN + BF16 + label smoothing\n",
+        "5. **Three-stage PhraseVAE curriculum**: Prevents posterior collapse\n",
+        "6. **Sub-1GB inference**: Phrase-level Mamba recurrence with fixed-size state\n",
+        "\n",
+        "### Key References\n",
+        "\n",
+        "- Gu & Dao (2023). **Mamba**: Linear-Time Sequence Modeling with Selective State Spaces\n",
+        "- **MIDI-RWKV** (2025). Personalizable Long-Context Symbolic Music Infilling\n",
+        "- **PhraseVAE** (2024). Phrase-level latent diffusion for music\n",
+        "- **FME** (2022). Domain-Knowledge-Inspired Music Embedding\n",
+        "- **σReparam** (2023). Stabilizing Transformer Training\n",
+        "- **ZClip** (2025). Adaptive Spike Mitigation for LLM Pre-Training\n",
+        "- **REMI** (2020). Pop Music Transformer\n",
+        "\n",
+        "📄 Full architecture document: [https://huggingface.co/asdf98/MuseMorphic](https://huggingface.co/asdf98/MuseMorphic)"
+      ]
+    }
+  ]
+}