Upload 2085 files

HW4/dataset.py

@@ -0,0 +1,213 @@
+"""
+dataset.py — NSynth spectrogram dataset
+
+Loads 0.5-second chunks of NSynth audio and converts them to complex spectrograms.
+Each item is ((2, FREQ_BINS, TIME_FRAMES), pitch) where pitch is a MIDI integer 0-127.
+
+Filename format: {family}_{source}_{id}-{pitch:03d}-{velocity:03d}.wav
+  e.g. keyboard_electronic_098-100-075.wav  →  pitch = 100
+
+Normalization
+-------------
+We use power-law magnitude compression on the raw STFT:
+
+    X_norm = β · |X|^α · exp(j·∠X)      (default α=0.5, β=1.0)
+
+This is sqrt-magnitude compression — a standard technique in audio generation
+that compresses the wide dynamic range without needing pre-computed dataset stats.
+The inverse is exact: given X_norm, recover X via |X_norm/β|^(1/α) · exp(j·∠X_norm).
+
+With n_fft=256, hop=128, and 0.5s at 16kHz:
+    FREQ_BINS   = 129   (n_fft // 2 + 1)
+    TIME_FRAMES = 63    (chunk_samples // hop_length + 1)
+"""
+
+import random
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+from torch.utils.data import Dataset
+from tqdm import tqdm
+
+# ── Audio / STFT constants ─────────────────────────────────────────────────────
+SR             = 16_000
+CHUNK_DURATION = 0.5
+CHUNK_SAMPLES  = int(SR * CHUNK_DURATION)  # 8 000
+N_FFT          = 256
+HOP_LENGTH     = 128
+FREQ_BINS      = N_FFT // 2 + 1            # 129
+TIME_FRAMES    = CHUNK_SAMPLES // HOP_LENGTH + 1  # 63
+
+# ── Power-law compression parameters ──────────────────────────────────────────
+ALPHA_RESCALE = 0.5   # magnitude exponent (0.5 = sqrt compression)
+BETA_RESCALE  = 1.0   # scale factor (1.0 = no rescaling)
+
+
+# ── Core normalization helpers ─────────────────────────────────────────────────
+
+def normalize_complex_powerlaw(
+    stft_complex: torch.Tensor,
+    alpha: float = ALPHA_RESCALE,
+    beta: float  = BETA_RESCALE,
+) -> torch.Tensor:
+    """
+    Power-law magnitude compression, preserving phase.
+
+    Transforms each STFT bin:  X → β · |X|^α · exp(j·∠X)
+
+    Parameters
+    ----------
+    stft_complex : complex tensor (..., FREQ_BINS, TIME_FRAMES)
+    """
+    mag   = stft_complex.abs().clamp(min=1e-9)
+    phase = torch.angle(stft_complex)
+    mag_c = beta * mag.pow(alpha)
+    return torch.complex(mag_c * torch.cos(phase), mag_c * torch.sin(phase))
+
+
+def denormalize_complex_powerlaw(
+    stft_norm: torch.Tensor,
+    alpha: float = ALPHA_RESCALE,
+    beta: float  = BETA_RESCALE,
+) -> torch.Tensor:
+    """Invert normalize_complex_powerlaw exactly: X_norm → original STFT complex."""
+    mag_c = stft_norm.abs().clamp(min=1e-9)
+    phase = torch.angle(stft_norm)
+    mag   = (mag_c / beta).pow(1.0 / alpha)
+    return torch.complex(mag * torch.cos(phase), mag * torch.sin(phase))
+
+
+# ── Public pipeline ────────────────────────────────────────────────────────────
+
+def get_pitch(path: Path) -> int:
+    """Parse the MIDI pitch (0-127) from an NSynth filename.
+
+    e.g. keyboard_electronic_098-100-075.wav  →  100
+    """
+    return int(path.stem.split("-")[-2])
+
+
+def wav_to_spec(path: Path) -> torch.Tensor:
+    """
+    Load one WAV file, grab a random 0.5-second chunk, compute its STFT,
+    and apply power-law magnitude compression.
+
+    No pre-computed dataset statistics are needed.
+
+    Returns
+    -------
+    spec : (2, FREQ_BINS, TIME_FRAMES)  float32
+        Channel 0 = real part, Channel 1 = imaginary part of the compressed STFT.
+    """
+    audio, sr = torchaudio.load(str(path))
+    if sr != SR:
+        audio = torchaudio.functional.resample(audio, sr, SR)
+    audio = audio[0]  # mono → (N,)
+
+    # Always take the first CHUNK_SAMPLES: NSynth notes decay into silence
+    # quickly, so a random crop would mostly yield silence and bias training.
+    chunk = audio[:CHUNK_SAMPLES]
+    if len(chunk) < CHUNK_SAMPLES:
+        chunk = F.pad(chunk, (0, CHUNK_SAMPLES - len(chunk)))
+
+    # STFT + power-law compression
+    window = torch.hann_window(N_FFT)
+    stft = torch.stft(
+        chunk, n_fft=N_FFT, hop_length=HOP_LENGTH,
+        window=window, return_complex=True, center=True,
+    )  # complex (FREQ_BINS, T)
+    stft = normalize_complex_powerlaw(stft)
+
+    # Split complex → two real channels
+    spec = torch.stack([stft.real, stft.imag], dim=0)  # (2, F, T)
+
+    # Ensure exact time dimension (numerical STFT edge cases)
+    if spec.shape[2] > TIME_FRAMES:
+        spec = spec[:, :, :TIME_FRAMES]
+    elif spec.shape[2] < TIME_FRAMES:
+        spec = F.pad(spec, (0, TIME_FRAMES - spec.shape[2]))
+
+    return spec  # (2, 129, 63)
+
+
+def spec_to_audio(spec: torch.Tensor) -> torch.Tensor:
+    """
+    Convert a power-law-compressed (2, FREQ_BINS, TIME_FRAMES) spectrogram
+    back to a waveform via the inverse power law and ISTFT.
+
+    Returns
+    -------
+    audio : (CHUNK_SAMPLES,)  float32
+    """
+    stft_norm = torch.complex(spec[0], spec[1])
+    stft = denormalize_complex_powerlaw(stft_norm)
+    window = torch.hann_window(N_FFT, device=spec.device)
+    audio = torch.istft(
+        stft, n_fft=N_FFT, hop_length=HOP_LENGTH,
+        window=window, center=True, length=CHUNK_SAMPLES,
+    )
+    return audio  # (CHUNK_SAMPLES,)
+
+
+# ── Dataset ────────────────────────────────────────────────────────────────────
+
+class NSynthSpecDataset(Dataset):
+    """
+    Dataset of (spectrogram, pitch) pairs from NSynth WAV files.
+
+    Spectrograms use power-law magnitude compression (no pre-computed stats needed).
+
+    Parameters
+    ----------
+    audio_dir         : path to a directory full of .wav files
+    max_files         : cap the number of files (useful for fast experiments)
+    instrument_filter : if set, only keep files whose stem starts with this prefix.
+                        Examples:
+                          "keyboard_synthetic"      → all keyboard synth samples
+                          "keyboard_synthetic_099"  → one specific instrument
+                          "guitar_acoustic"         → all acoustic guitars
+    cache             : if True (default), cache spectrograms in RAM after first load.
+                        Epoch 1 reads from disk; all subsequent epochs hit RAM.
+                        At ~65 KB/spec, 5 000 files ≈ 325 MB — fits easily in Colab.
+    """
+
+    def __init__(
+        self,
+        audio_dir: str,
+        max_files: int = None,
+        instrument_filter: str = None,
+        cache: bool = True,
+    ):
+        files = sorted(Path(audio_dir).glob("*.wav"))
+
+        if instrument_filter:
+            files = [f for f in files if f.stem.startswith(instrument_filter)]
+            print(f"Instrument filter '{instrument_filter}': {len(files)} files match")
+
+        if max_files and len(files) > max_files:
+            rng = random.Random(42)
+            rng.shuffle(files)
+            files = files[:max_files]
+
+        self.files = files
+        self._cache: list = [None] * len(files) if cache else None
+        print(f"NSynthSpecDataset: {len(self.files)} files  (RAM cache: {'on' if cache else 'off'})")
+
+    def __len__(self):
+        return len(self.files)
+
+    def __getitem__(self, idx):
+        if self._cache is not None and self._cache[idx] is not None:
+            return self._cache[idx]
+
+        path  = self.files[idx]
+        spec  = wav_to_spec(path)                                    # (2, F, T)
+        pitch = torch.tensor(get_pitch(path), dtype=torch.long)     # scalar
+        item  = (spec, pitch)
+
+        if self._cache is not None:
+            self._cache[idx] = item
+
+        return item

HW4/homework4_stub.ipynb

@@ -0,0 +1,1166 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "h6IdxQC3ivWO"
+   },
+   "source": [
+    "# Assignment: Flow Matching for Audio Generation\n",
+    "\n",
+    "In this assignment you will implement the **inference and training** code for a pitch-conditioned\n",
+    "flow matching model trained on the [NSynth](https://magenta.tensorflow.org/datasets/nsynth)\n",
+    "dataset. A pretrained model (`pretrained_keyboard.pt`) is provided so you can hear results\n",
+    "immediately and focus on understanding the algorithms.\n",
+    "\n",
+    "**Before you begin:** Runtime -> Change runtime type -> **T4 GPU** (required).\n",
+    "\n",
+    "---\n",
+    "\n",
+    "## What is given to you\n",
+    "| File | Contents |\n",
+    "|---|---|\n",
+    "| `dataset.py` | Spectrogram extraction, normalization, `NSynthSpecDataset` |\n",
+    "| `model.py` | Three model architectures and `build_model_from_config` |\n",
+    "| `pretrained_keyboard.pt` | Pretrained UNet-based flow model (~125k params, keyboard sounds) |\n",
+    "\n",
+    "---\n",
+    "\n",
+    "## Parts and points\n",
+    "| Part | Task | Points |\n",
+    "|---|---|---|\n",
+    "| 1 | Euler sampling | 2 |\n",
+    "| 2a | Naive velocity scaling (warmup) | 1 |\n",
+    "| 2b | Classifier-Free Guidance (CFG) | 2 |\n",
+    "| 3a | Heun's method | 2 |\n",
+    "| 3b | RK4 | 1 |\n",
+    "| 4a | Timestep Sampling | 0.5 |\n",
+    "| 4b | Flow Loss | 1.5 |\n",
+    "| **Total** | | **10** |\n",
+    "| 4c (bonus) | Fine-tuning on a new instrument | +0.5 |\n",
+    "| 5 (bonus) | Beat the baseline (open-ended) | +0.5 |\n",
+    "\n",
+    "---\n",
+    "\n",
+    "## Submission\n",
+    "\n",
+    "When you are done, **convert this notebook to a Python file** for submission:\n",
+    "\n",
+    "```\n",
+    "File -> Download -> Download .py\n",
+    "```\n",
+    "\n",
+    "or from the command line:\n",
+    "\n",
+    "```bash\n",
+    "jupyter nbconvert --to script assignment.ipynb\n",
+    "```\n",
+    "\n",
+    "**Submit the following:**\n",
+    "1. `assignment.py` — your converted notebook\n",
+    "2. (Bonus) `submission_q4.npz` and `model_ft_q4.pt` - Part 4 generated samples and fine-tuned checkpoint\n",
+    "3. (Bonus) `submission_q5.npz` and `model_q5.pt` if attempting Part 5"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "7mQDjgM8ivWQ"
+   },
+   "source": [
+    "## Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "EKOv3SOAivWR"
+   },
+   "outputs": [],
+   "source": [
+    "# Install missing deps (torch/torchaudio are pre-installed on Colab)\n",
+    "!pip install -q librosa tqdm\n",
+    "\n",
+    "import os, sys, json\n",
+    "from pathlib import Path\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.nn.functional as F\n",
+    "import torchaudio\n",
+    "from torch.utils.data import DataLoader\n",
+    "import IPython.display as ipd\n",
+    "\n",
+    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
+    "print(f'Device: {device}')\n",
+    "if torch.cuda.is_available():\n",
+    "    print(f'GPU: {torch.cuda.get_device_name(0)}')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "2329c5be"
+   },
+   "outputs": [],
+   "source": [
+    "# dataset.py, model.py, and pretrained_keyboard.pt are provided alongside this notebook.\n",
+    "# Make sure they are in the same directory as this file before running.\n",
+    "import os, sys\n",
+    "sys.path.insert(0, os.path.dirname(os.path.abspath('__file__')) if '__file__' in dir() else '.')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "d4GKlQBBivWS"
+   },
+   "outputs": [],
+   "source": [
+    "# Load dataset and model utilities\n",
+    "from dataset import NSynthSpecDataset, wav_to_spec, spec_to_audio, FREQ_BINS, TIME_FRAMES, SR\n",
+    "from model  import load_flow_model, save_flow_model, build_model_from_config, FlowModelWrapper, NULL_PITCH\n",
+    "\n",
+    "# Load the pretrained model (wrapped for standard diffusion convention: t=1 noise, t=0 data)\n",
+    "CKPT_PATH = 'pretrained_keyboard.pt'\n",
+    "model, ckpt = load_flow_model(CKPT_PATH, device=device)\n",
+    "\n",
+    "print(f'Model: {ckpt[\"n_params\"]:,} parameters  ({ckpt[\"n_params\"]/1e3:.0f}k)')\n",
+    "print(f'Input shape : (batch, 2, {FREQ_BINS}, {TIME_FRAMES})')\n",
+    "print(f'  2 channels = real + imaginary parts of a 0.5-second complex STFT')\n",
+    "print(f'  freq_bins  = {FREQ_BINS} = n_fft/2 + 1  (n_fft=256)')\n",
+    "print(f'  time_frames= {TIME_FRAMES}')\n",
+    "print(f'NULL_PITCH  : {NULL_PITCH}  (the \"no conditioning\" token used for CFG)')\n",
+    "print(f'Sample rate : {SR} Hz')\n",
+    "print(f'Time convention: t=1 is noise, t=0 is data')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "BP3I8fNeivWS"
+   },
+   "source": [
+    "---\n",
+    "## Background: Rectified Flow in one page\n",
+    "\n",
+    "In class, we studied diffusion models, which have roots in stochastic differential equations (SDEs). In recent years, the popular and highly related *flow-matching* formulation has become increasingly popular, which has its roots in deterministic *Ordinary* differential equations (ODEs). When choosing flow matching with a gaussian distribution, and in our specific choice of optimal transport paths, the math becomes very simple! For those wondering, despite some common confusion on diffusion vs. flow matching, they are fundamentally [the same thing](https://diffusionflow.github.io/).\n",
+    "\n",
+    "**Training objective** (Flow Matching with Optimal Transport paths):\n",
+    "\n",
+    "Given a data sample $x_0$ (spectrogram), noise $\\epsilon \\sim \\mathcal{N}(0,I)$, and pitch label $p$:\n",
+    "$$x_t = (1-t)\\,x_0 + t\\,\\epsilon  \\qquad t \\sim U[0,1]$$\n",
+    "At $t=0$ this is pure data; at $t=1$ it is pure noise. The velocity along the straight path is:\n",
+    "$$v^* = \\epsilon - x_0 \\quad\\text{(points from data toward noise)}$$\n",
+    "$$\\mathcal{L} = \\mathbb{E}_{t,x_0,x_1}\\bigl[\\|v_\\theta(x_t, t, p) - v^*\\|^2\\bigr]$$\n",
+    "\n",
+    "The model $v_\\theta(x_t, t, p)$ learns to predict this velocity field. From this field, we can then use any classical ODE solver.\n",
+    "\n",
+    "**Key model interface:**\n",
+    "```python\n",
+    "v = model(x, t, pitch)\n",
+    "# x     : (B, 2, FREQ_BINS, TIME_FRAMES)  — current noisy spectrogram\n",
+    "# t     : (B,)  float in [0, 1]           — current time (1=noise, 0=data)\n",
+    "# pitch : (B,)  int  in [0, 127]          — MIDI pitch (or NULL_PITCH=128 for uncond)\n",
+    "# v     : same shape as x                 — predicted velocity (data → noise direction, i.e. ε - x₀)\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "HlhykRBxivWS"
+   },
+   "source": [
+    "---\n",
+    "## Part 1 — Euler Sampling  `[2 pts]`\n",
+    "\n",
+    "Implement the basic Euler ODE solver. Starting from Gaussian noise at $t=1$, take `n_steps` equal\n",
+    "steps of size $\\Delta t = 1/n$ toward the data distribution at $t=0$:\n",
+    "$$x_{t-\\Delta t} = x_t - v_\\theta(x_t,\\, t,\\, p)\\,\\Delta t, \\qquad t = 1,\\; 1{-}\\Delta t,\\; 1{-}2\\Delta t,\\; \\ldots$$\n",
+    "\n",
+    "**Hints:**\n",
+    "- All tensors are already on the right device — don't call `.to()` inside the loop.\n",
+    "- Create the time tensor `t_batch = torch.full((B,), t_val, device=x.device)` before each model call.\n",
+    "- The first step starts at $t = 1$ and the last step starts at $t = 1/n$, arriving at $t = 0$."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "EBU5XvvYivWS"
+   },
+   "outputs": [],
+   "source": [
+    "def euler_sample(model, x1, pitches, n_steps=50):\n",
+    "    \"\"\"\n",
+    "    Euler ODE integration from noise (t=1) to data (t=0).\n",
+    "\n",
+    "    Parameters\n",
+    "    ----------\n",
+    "    model   : flow model (eval mode, on device)\n",
+    "    x1      : (B, 2, FREQ_BINS, TIME_FRAMES)  initial Gaussian noise\n",
+    "    pitches : (B,) MIDI pitches, dtype=torch.long\n",
+    "    n_steps : number of Euler steps\n",
+    "\n",
+    "    Returns\n",
+    "    -------\n",
+    "    x : (B, 2, FREQ_BINS, TIME_FRAMES)  generated spectrograms at t=0\n",
+    "    \"\"\"\n",
+    "    ### YOUR CODE HERE ###\n",
+    "    dt = ...\n",
+    "    x = x1.clone()\n",
+    "    B = ...\n",
+    "    with torch.no_grad():\n",
+    "        for i in range(n_steps):\n",
+    "            t_batch = ...\n",
+    "            v = ...\n",
+    "            x = ...\n",
+    "    return x\n",
+    "    raise NotImplementedError"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "MW9o9BeHivWS"
+   },
+   "outputs": [],
+   "source": [
+    "# === Sanity Check 1 — do not modify ===\n",
+    "torch.manual_seed(0)\n",
+    "x1_ag = torch.randn(4, 2, FREQ_BINS, TIME_FRAMES, device=device)\n",
+    "p_ag  = torch.tensor([60, 62, 64, 67], dtype=torch.long, device=device)\n",
+    "\n",
+    "out1 = euler_sample(model, x1_ag.clone(), p_ag, n_steps=20)\n",
+    "\n",
+    "assert out1.shape == (4, 2, FREQ_BINS, TIME_FRAMES), \\\n",
+    "    f'Wrong output shape: {out1.shape}'\n",
+    "assert not torch.allclose(out1, x1_ag, atol=1e-3), \\\n",
+    "    'Output equals input — did you implement the integration loop?'\n",
+    "assert out1.isfinite().all(), 'Output contains NaN or Inf'\n",
+    "assert out1.std() > 0.05, f'Output std={out1.std():.4f} is suspiciously low'\n",
+    "print(f'\\u2713 euler_sample | shape={out1.shape}, mean={out1.mean():.4f}, std={out1.std():.4f}')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "HKCQomYdivWT"
+   },
+   "outputs": [],
+   "source": [
+    "# Listen: C major scale (C4 D4 E4 F4 G4 A4 B4 C5)\n",
+    "NOTE_NAMES = ['C','C#','D','D#','E','F','F#','G','G#','A','A#','B']\n",
+    "\n",
+    "torch.manual_seed(7)\n",
+    "demo_pitches = torch.tensor([60,62,64,65,67,69,71,72], dtype=torch.long, device=device)\n",
+    "demo_noise   = torch.randn(len(demo_pitches), 2, FREQ_BINS, TIME_FRAMES, device=device)\n",
+    "\n",
+    "samples_euler = euler_sample(model, demo_noise.clone(), demo_pitches, n_steps=50)\n",
+    "\n",
+    "print('Euler samples (no CFG):  notice how pitch identity may be weak at scale=1\\n')\n",
+    "for spec, pitch in zip(samples_euler, demo_pitches.tolist()):\n",
+    "    audio = spec_to_audio(spec.cpu())\n",
+    "    audio = audio / (audio.abs().max() + 1e-8)\n",
+    "    name  = NOTE_NAMES[pitch % 12] + str(pitch // 12 - 1)\n",
+    "    print(f'  {name} (MIDI {pitch})')\n",
+    "    display(ipd.Audio(audio.numpy(), rate=SR))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "M91nx0aDivWT"
+   },
+   "source": [
+    "---\n",
+    "## Part 2a — Naive Velocity Scaling  `[1 pts]`\n",
+    "\n",
+    "Before implementing proper CFG, let's explore a **wrong but instructive** idea:\n",
+    "just multiply the velocity by a scalar `scale` at each step.\n",
+    "\n",
+    "$$x_{t-\\Delta t} = x_t - \\underbrace{v_\\theta(x_t, t, p) \\cdot \\texttt{scale}}_{\\text{scaled velocity}} \\cdot \\Delta t$$\n",
+    "\n",
+    "When `scale=1.0` this should be identical to `euler_sample`.  \n",
+    "When `scale>1` the integration \"goes faster\" — think about what happens when you overshoot.\n",
+    "\n",
+    "Try `scale=2.0` and `scale=0.5` and listen to the results. What do you notice?"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "2pBp-lN5ivWT"
+   },
+   "outputs": [],
+   "source": [
+    "def naive_scale_sample(model, x1, pitches, n_steps=50, scale=1.0):\n",
+    "    \"\"\"\n",
+    "    Euler sampling (t=1 → t=0) with velocity multiplied by `scale`.\n",
+    "\n",
+    "    Parameters: same as euler_sample, plus\n",
+    "    scale : float — multiply every velocity prediction by this factor\n",
+    "\n",
+    "    Returns: (B, 2, FREQ_BINS, TIME_FRAMES)\n",
+    "    \"\"\"\n",
+    "    ### YOUR CODE HERE ###\n",
+    "    raise NotImplementedError"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "VrWWQFd4ivWT"
+   },
+   "outputs": [],
+   "source": [
+    "# === Sanity Check 2a — do not modify ===\n",
+    "torch.manual_seed(0)\n",
+    "x1_ag = torch.randn(4, 2, FREQ_BINS, TIME_FRAMES, device=device)\n",
+    "p_ag  = torch.tensor([60, 62, 64, 67], dtype=torch.long, device=device)\n",
+    "\n",
+    "out_s1 = naive_scale_sample(model, x1_ag.clone(), p_ag, n_steps=20, scale=1.0)\n",
+    "out_e  = euler_sample(      model, x1_ag.clone(), p_ag, n_steps=20)\n",
+    "out_s2 = naive_scale_sample(model, x1_ag.clone(), p_ag, n_steps=20, scale=2.0)\n",
+    "\n",
+    "assert torch.allclose(out_s1, out_e, atol=1e-5), \\\n",
+    "    'naive_scale_sample(scale=1.0) must match euler_sample exactly'\n",
+    "assert not torch.allclose(out_s2, out_e, atol=1e-3), \\\n",
+    "    'naive_scale_sample(scale=2.0) should differ from scale=1.0'\n",
+    "print('\\u2713 naive_scale_sample | scale=1.0 matches Euler, scale=2.0 differs')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "kHnV8jDcivWT"
+   },
+   "outputs": [],
+   "source": [
+    "# Compare scale=1.0 vs scale=2.0 — what does scaling velocity actually do?\n",
+    "torch.manual_seed(42)\n",
+    "test_pitch  = torch.full((1,), 60, dtype=torch.long, device=device)\n",
+    "test_noise  = torch.randn(1, 2, FREQ_BINS, TIME_FRAMES, device=device)\n",
+    "\n",
+    "for scale in [0.5, 1.0, 2.0, 4.0]:\n",
+    "    s = naive_scale_sample(model, test_noise.clone(), test_pitch, n_steps=50, scale=scale)\n",
+    "    audio = spec_to_audio(s[0].cpu())\n",
+    "    audio = audio / (audio.abs().max() + 1e-8)\n",
+    "    print(f'scale={scale}')\n",
+    "    display(ipd.Audio(audio.numpy(), rate=SR))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "ij2kU-7CivWT"
+   },
+   "source": [
+    "---\n",
+    "## Part 2b — Classifier-Free Guidance (CFG)  `[2 pts]`\n",
+    "\n",
+    "The correct way to amplify conditioning is to run the model **twice** per step:\n",
+    "once with the actual pitch and once with `NULL_PITCH` (the \"no conditioning\" token),\n",
+    "then combine:\n",
+    "\n",
+    "$$v_{\\text{cond}}   = v_\\theta(x_t,\\, t,\\, p)\\\\ v_{\\text{uncond}} = v_\\theta(x_t,\\, t,\\, \\texttt{NULL\\_PITCH})\\\\ v = v_{\\text{uncond}} + s \\cdot (v_{\\text{cond}} - v_{\\text{uncond}})$$\n",
+    "\n",
+    "Notice: at $s=1$, this equals $v_{\\text{cond}}$ (standard Euler). At $s=0$ it equals $v_{\\text{uncond}}$.\n",
+    "Values $s>1$ extrapolate **beyond** the conditional estimate, sharpening pitch adherence.\n",
+    "\n",
+    "The Euler update with CFG is then: $x_{t-\\Delta t} = x_t - v \\cdot \\Delta t$\n",
+    "\n",
+    "**Why is this better than naive scaling?**  \n",
+    "Naive scaling changes only the *magnitude* of the velocity (affects \"how far\" you step), which can lead to severe oversaturation artifacts.  \n",
+    "CFG changes the *direction and magnitude*, as it steers toward regions that are more associated\n",
+    "with the conditioned pitch, without just overshooting.\n",
+    "\n",
+    "**Hints:**\n",
+    "- `torch.full_like(pitches, NULL_PITCH)` creates the unconditional pitch batch.\n",
+    "- When `guidance_scale == 1.0`, skip the second model call (just return `v_cond`)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "bvkYNRhJivWT"
+   },
+   "outputs": [],
+   "source": [
+    "def cfg_sample(model, x1, pitches, n_steps=50, guidance_scale=1.0):\n",
+    "    \"\"\"\n",
+    "    Euler sampling (t=1 → t=0) with Classifier-Free Guidance.\n",
+    "\n",
+    "    Parameters\n",
+    "    ----------\n",
+    "    model          : flow model\n",
+    "    x1             : (B, 2, FREQ_BINS, TIME_FRAMES)  initial noise\n",
+    "    pitches        : (B,) MIDI pitches, dtype=torch.long\n",
+    "    n_steps        : Euler integration steps\n",
+    "    guidance_scale : float; 1.0 = no guidance, 3-7 = strong\n",
+    "\n",
+    "    Returns\n",
+    "    -------\n",
+    "    (B, 2, FREQ_BINS, TIME_FRAMES)\n",
+    "    \"\"\"\n",
+    "    ### YOUR CODE HERE ###\n",
+    "    raise NotImplementedError"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "iWz5TYuEivWU"
+   },
+   "outputs": [],
+   "source": [
+    "# === Sanity Check 2b — do not modify ===\n",
+    "torch.manual_seed(0)\n",
+    "x1_ag = torch.randn(4, 2, FREQ_BINS, TIME_FRAMES, device=device)\n",
+    "p_ag  = torch.tensor([60, 62, 64, 67], dtype=torch.long, device=device)\n",
+    "\n",
+    "out_cfg1   = cfg_sample(model, x1_ag.clone(), p_ag, n_steps=20, guidance_scale=1.0)\n",
+    "out_euler  = euler_sample(model, x1_ag.clone(), p_ag, n_steps=20)\n",
+    "out_cfg6   = cfg_sample(model, x1_ag.clone(), p_ag, n_steps=20, guidance_scale=6.0)\n",
+    "out_naive2 = naive_scale_sample(model, x1_ag.clone(), p_ag, n_steps=20, scale=2.0)\n",
+    "\n",
+    "assert torch.allclose(out_cfg1, out_euler, atol=1e-5), \\\n",
+    "    'cfg_sample(guidance_scale=1.0) must equal euler_sample'\n",
+    "assert not torch.allclose(out_cfg6, out_euler, atol=1e-3), \\\n",
+    "    'cfg_sample(guidance_scale=6.0) should differ from scale=1.0'\n",
+    "assert not torch.allclose(out_cfg6, out_naive2, atol=1e-3), \\\n",
+    "    'CFG (scale=6) must differ from naive scaling (scale=2) — they use different formulas'\n",
+    "print('\\u2713 cfg_sample | gs=1.0 matches Euler, gs=6.0 differs, CFG != naive scaling')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "QxzyMDWKivWU"
+   },
+   "outputs": [],
+   "source": [
+    "# Listen: same noise, vary guidance scale 1 → 3 → 6 → 10\n",
+    "torch.manual_seed(42)\n",
+    "test_pitch = torch.full((1,), 60, dtype=torch.long, device=device)\n",
+    "test_noise = torch.randn(1, 2, FREQ_BINS, TIME_FRAMES, device=device)\n",
+    "\n",
+    "for gs in [1.0, 3.0, 6.0, 10.0]:\n",
+    "    s = cfg_sample(model, test_noise.clone(), test_pitch, n_steps=50, guidance_scale=gs)\n",
+    "    audio = spec_to_audio(s[0].cpu())\n",
+    "    audio = audio / (audio.abs().max() + 1e-8)\n",
+    "    print(f'guidance_scale={gs}  — pitch adherence increases, diversity decreases at high scales')\n",
+    "    display(ipd.Audio(audio.numpy(), rate=SR))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "xpBDeRyTivWU"
+   },
+   "source": [
+    "---\n",
+    "## Part 3a — Heun's Method  `[2 pts]`\n",
+    "\n",
+    "Euler's method has **first-order** local truncation error $O(\\Delta t^2)$.  \n",
+    "**Heun's method** (improved Euler / explicit trapezoidal rule) achieves **second-order** accuracy\n",
+    "$O(\\Delta t^3)$ by using two velocity evaluations per step:\n",
+    "\n",
+    "$$k_1 = v_\\theta(x_t,\\, t)$$\n",
+    "$$\\hat{x} = x_t - k_1 \\cdot \\Delta t \\quad\\text{(Euler predictor)}$$\n",
+    "$$k_2 = v_\\theta(\\hat{x},\\, t - \\Delta t)$$\n",
+    "$$x_{t-\\Delta t} = x_t - \\tfrac{1}{2}(k_1 + k_2) \\cdot \\Delta t \\quad\\text{(corrected update)}$$\n",
+    "\n",
+    "At the **same NFE budget** (number of function evaluations), Heun with $n/2$ steps\n",
+    "uses the same compute as Euler with $n$ steps, but with much lower discretization error.\n",
+    "\n",
+    "Your implementation should also work with CFG from the previous part."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "Q0Q1T6lCivWU"
+   },
+   "outputs": [],
+   "source": [
+    "def heun_sample(model, x1, pitches, n_steps=50, guidance_scale=1.0):\n",
+    "    \"\"\"\n",
+    "    Heun's method (2nd-order Runge-Kutta) from t=1 to t=0, with optional CFG.\n",
+    "\n",
+    "    Parameters: same as cfg_sample.\n",
+    "    Returns: (B, 2, FREQ_BINS, TIME_FRAMES)\n",
+    "    \"\"\"\n",
+    "    ### YOUR CODE HERE ###\n",
+    "    raise NotImplementedError"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "n_tRQIE5ivWU"
+   },
+   "outputs": [],
+   "source": [
+    "# === Sanity Check 3a — do not modify ===\n",
+    "torch.manual_seed(0)\n",
+    "x1_ag = torch.randn(4, 2, FREQ_BINS, TIME_FRAMES, device=device)\n",
+    "p_ag  = torch.tensor([60, 62, 64, 67], dtype=torch.long, device=device)\n",
+    "\n",
+    "out_euler = euler_sample(model, x1_ag.clone(), p_ag, n_steps=25)\n",
+    "out_heun  = heun_sample( model, x1_ag.clone(), p_ag, n_steps=25, guidance_scale=1.0)\n",
+    "\n",
+    "assert out_heun.shape == (4, 2, FREQ_BINS, TIME_FRAMES), \\\n",
+    "    f'Wrong shape: {out_heun.shape}'\n",
+    "assert out_heun.isfinite().all(), 'Heun output contains NaN/Inf'\n",
+    "assert not torch.allclose(out_heun, out_euler, atol=1e-3), \\\n",
+    "    'heun_sample must differ from euler_sample'\n",
+    "\n",
+    "# CFG version: gs=1 should match no-CFG Heun\n",
+    "out_heun_gs1 = heun_sample(model, x1_ag.clone(), p_ag, n_steps=25, guidance_scale=1.0)\n",
+    "assert torch.allclose(out_heun, out_heun_gs1, atol=1e-5), \\\n",
+    "    'heun_sample results should be deterministic given the same inputs'\n",
+    "\n",
+    "out_heun_gs6 = heun_sample(model, x1_ag.clone(), p_ag, n_steps=25, guidance_scale=6.0)\n",
+    "assert not torch.allclose(out_heun_gs6, out_heun, atol=1e-3), \\\n",
+    "    'heun_sample with guidance_scale=6 should differ from guidance_scale=1'\n",
+    "\n",
+    "l2 = (out_heun - out_euler).norm().item()\n",
+    "print(f'\\u2713 heun_sample | shape OK, differs from Euler (L2={l2:.3f}), CFG works')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "0o0u3rmuivWU"
+   },
+   "source": [
+    "---\n",
+    "## Part 3b — RK4  `[1 pts]`\n",
+    "\n",
+    "The **classic 4th-order Runge-Kutta** method uses 4 velocity evaluations per step.\n",
+    "Since we integrate **backward** from $t$ to $t - \\Delta t$:\n",
+    "\n",
+    "$$k_1 = v_\\theta(x_t,\\;t)$$\n",
+    "$$k_2 = v_\\theta(x_t - k_1\\tfrac{\\Delta t}{2},\\;t - \\tfrac{\\Delta t}{2})$$\n",
+    "$$k_3 = v_\\theta(x_t - k_2\\tfrac{\\Delta t}{2},\\;t - \\tfrac{\\Delta t}{2})$$\n",
+    "$$k_4 = v_\\theta(x_t - k_3\\Delta t,\\;t - \\Delta t)$$\n",
+    "$$x_{t-\\Delta t} = x_t - \\tfrac{\\Delta t}{6}(k_1 + 2k_2 + 2k_3 + k_4)$$\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "Q_0dDlLFivWU"
+   },
+   "outputs": [],
+   "source": [
+    "def rk4_sample(model, x1, pitches, n_steps=25, guidance_scale=1.0):\n",
+    "    \"\"\"\n",
+    "    4th-order Runge-Kutta from t=1 to t=0, with optional CFG.\n",
+    "\n",
+    "    Parameters: same as heun_sample.\n",
+    "    Returns: (B, 2, FREQ_BINS, TIME_FRAMES)\n",
+    "    \"\"\"\n",
+    "    ### YOUR CODE HERE ###\n",
+    "    raise NotImplementedError"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "esbLX00ZivWU"
+   },
+   "outputs": [],
+   "source": [
+    "# === Sanity Check 3b — do not modify ===\n",
+    "torch.manual_seed(0)\n",
+    "x1_ag = torch.randn(4, 2, FREQ_BINS, TIME_FRAMES, device=device)\n",
+    "p_ag  = torch.tensor([60, 62, 64, 67], dtype=torch.long, device=device)\n",
+    "\n",
+    "out_rk4 = rk4_sample(model, x1_ag.clone(), p_ag, n_steps=12, guidance_scale=1.0)\n",
+    "assert out_rk4.shape == (4, 2, FREQ_BINS, TIME_FRAMES)\n",
+    "assert out_rk4.isfinite().all()\n",
+    "out_heun = heun_sample(model, x1_ag.clone(), p_ag, n_steps=25, guidance_scale=1.0)\n",
+    "assert not torch.allclose(out_rk4, out_heun, atol=1e-3), 'RK4 should differ from Heun'\n",
+    "print('\\u2713 rk4_sample: shape OK, differs from Heun')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "TDCogL8CivWU"
+   },
+   "outputs": [],
+   "source": [
+    "# Compare solvers at equal NFE budget (50 model evaluations each)\n",
+    "# Euler: 50 steps x 1 eval/step = 50 NFE\n",
+    "# Heun:  25 steps x 2 eval/step = 50 NFE\n",
+    "# RK4:   12 steps x 4 eval/step = 48 NFE\n",
+    "\n",
+    "torch.manual_seed(7)\n",
+    "pitch1 = torch.full((1,), 60, dtype=torch.long, device=device)\n",
+    "noise1 = torch.randn(1, 2, FREQ_BINS, TIME_FRAMES, device=device)\n",
+    "\n",
+    "GS = 6.0  # guidance scale for all\n",
+    "\n",
+    "comparisons = [\n",
+    "    ('Euler  (50 steps, gs=6)', lambda: cfg_sample(  model, noise1.clone(), pitch1, n_steps=50,  guidance_scale=GS)),\n",
+    "    ('Heun   (25 steps, gs=6)', lambda: heun_sample( model, noise1.clone(), pitch1, n_steps=25,  guidance_scale=GS)),\n",
+    "    ('RK4    (12 steps, gs=6)', lambda: rk4_sample(  model, noise1.clone(), pitch1, n_steps=12,  guidance_scale=GS)),\n",
+    "]\n",
+    "\n",
+    "for label, fn in comparisons:\n",
+    "    s = fn()\n",
+    "    audio = spec_to_audio(s[0].cpu())\n",
+    "    audio = audio / (audio.abs().max() + 1e-8)\n",
+    "    print(label)\n",
+    "    display(ipd.Audio(audio.numpy(), rate=SR))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "HPzW6nlIivWU"
+   },
+   "source": [
+    "---\n",
+    "## Part 4a — Timestep Sampling  `[0.5 pts]`\n",
+    "\n",
+    "During flow matching training, we sample a random time $t \\in [0,1]$ for each example,\n",
+    "controlling *how much* noise is mixed in. Two strategies:\n",
+    "\n",
+    "- **Uniform:** $t \\sim U[0,1]$ — equal weight across all noise levels\n",
+    "- **Logit-normal:** $t = \\sigma(z),\\; z \\sim \\mathcal{N}(0,1)$ — concentrates weight near\n",
+    "  $t = 0.5$ (the hardest denoising regime, used in Stable Diffusion 3 / Flux)\n",
+    "\n",
+    "Implement `sample_timesteps` to support both modes."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "IyLjcASeivWU"
+   },
+   "outputs": [],
+   "source": [
+    "def sample_timesteps(B: int, device, t_sample: str = 'logit_normal') -> torch.Tensor:\n",
+    "    \"\"\"\n",
+    "    Sample B timestep values in [0, 1].\n",
+    "\n",
+    "    Parameters\n",
+    "    ----------\n",
+    "    B        : batch size\n",
+    "    device   : torch device (e.g. 'cuda' or x.device)\n",
+    "    t_sample : 'uniform' or 'logit_normal'\n",
+    "\n",
+    "    Returns\n",
+    "    -------\n",
+    "    t : (B,) tensor of floats in [0, 1]\n",
+    "    \"\"\"\n",
+    "    ### YOUR CODE HERE ###\n",
+    "    raise NotImplementedError"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "x4H3WQ6givWU"
+   },
+   "outputs": [],
+   "source": [
+    "# === Sanity Check 4a — do not modify ===\n",
+    "torch.manual_seed(0)\n",
+    "t_unif  = sample_timesteps(1000, device, 'uniform')\n",
+    "t_logit = sample_timesteps(1000, device, 'logit_normal')\n",
+    "\n",
+    "assert t_unif.shape  == (1000,), f'Wrong shape: {t_unif.shape}'\n",
+    "assert (t_unif  >= 0).all() and (t_unif  <= 1).all(), 'Uniform t must be in [0, 1]'\n",
+    "assert t_logit.shape == (1000,), f'Wrong shape: {t_logit.shape}'\n",
+    "assert (t_logit >= 0).all() and (t_logit <= 1).all(), 'Logit-normal t must be in [0, 1]'\n",
+    "\n",
+    "# logit-normal should be more concentrated near 0.5 than uniform\n",
+    "assert (t_logit - 0.5).abs().mean() < (t_unif - 0.5).abs().mean(), \\\n",
+    "    'Logit-normal should concentrate more near t=0.5 than uniform'\n",
+    "\n",
+    "print(f'\\u2713 sample_timesteps | '\n",
+    "      f'uniform mean|t-0.5|={(t_unif-0.5).abs().mean():.3f}, '\n",
+    "      f'logit-normal mean|t-0.5|={(t_logit-0.5).abs().mean():.3f}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "-ZxC2ubKivWU"
+   },
+   "source": [
+    "---\n",
+    "## Part 4b — Flow Loss  `[1.5 pts]`\n",
+    "\n",
+    "Now implement the core loss function. Given a batch of clean data $x_0$, pitch labels $p$,\n",
+    "and pre-sampled timesteps $t$, your function should:\n",
+    "\n",
+    "1. Sample $\\epsilon \\sim \\mathcal{N}(0,I)$ (noise, same shape as $x_0$)\n",
+    "2. Interpolate: $x_t = (1-t)\\,x_0 + t\\,\\epsilon$\n",
+    "3. Target velocity: $v^* = \\epsilon - x_0$\n",
+    "4. **CFG dropout:** replace $p$ with `NULL_PITCH` with probability `p_uncond`\n",
+    "5. Predict: $\\hat{v} = v_\\theta(x_t, t, p)$\n",
+    "6. Return $\\mathcal{L} = \\text{MSE}(\\hat{v},\\, v^*)$\n",
+    "\n",
+    "**Hints:**\n",
+    "- Broadcast `t` from shape `(B,)` to `(B,2,F,T)` via `t[:,None,None,None]`\n",
+    "- Return `F.mse_loss(v_pred, target)` — **do not call `.item()`** here; the training loop calls `.backward()` on your return value\n",
+    "- CFG dropout: `mask = torch.rand(B, device=x_data.device) < p_uncond`; use `pitch.clone()` to avoid modifying the original"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "fjKraYfKivWU"
+   },
+   "outputs": [],
+   "source": [
+    "def flow_loss(model, x_data, pitch, t, p_uncond=0.1):\n",
+    "    \"\"\"\n",
+    "    Compute the flow matching loss for one batch.\n",
+    "\n",
+    "    Parameters\n",
+    "    ----------\n",
+    "    model   : flow model (should be in train mode when called)\n",
+    "    x_data  : (B, 2, FREQ_BINS, TIME_FRAMES) — clean data batch (x_0), on device\n",
+    "    pitch   : (B,) — MIDI pitch labels, dtype=torch.long, on device\n",
+    "    t       : (B,) — sampled timestep values in [0, 1], on device\n",
+    "    p_uncond: CFG dropout probability (replace pitch with NULL_PITCH)\n",
+    "\n",
+    "    Returns\n",
+    "    -------\n",
+    "    loss : scalar tensor (differentiable — do NOT call .item())\n",
+    "    \"\"\"\n",
+    "    ### YOUR CODE HERE ###\n",
+    "    raise NotImplementedError"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "wJzxwq1aivWU"
+   },
+   "outputs": [],
+   "source": [
+    "# === Sanity Check 4b — do not modify ===\n",
+    "import copy\n",
+    "torch.manual_seed(42)\n",
+    "\n",
+    "x_data_ag = torch.randn(4, 2, FREQ_BINS, TIME_FRAMES, device=device) * 0.5\n",
+    "p_ag      = torch.randint(0, 128, (4,), dtype=torch.long, device=device)\n",
+    "t_ag      = sample_timesteps(4, device, 'logit_normal')\n",
+    "\n",
+    "model_ag = copy.deepcopy(model)\n",
+    "opt_ag   = torch.optim.AdamW(model_ag.parameters(), lr=1e-4)\n",
+    "\n",
+    "model_ag.train()\n",
+    "loss = flow_loss(model_ag, x_data_ag, p_ag, t_ag, p_uncond=0.0)\n",
+    "\n",
+    "assert loss.shape == (), \\\n",
+    "    f'flow_loss must return a scalar tensor, got shape {loss.shape}. Do not call .item()!'\n",
+    "assert loss.grad_fn is not None, \\\n",
+    "    'flow_loss must return a differentiable tensor for .backward() to work'\n",
+    "assert 0 < loss.item() < 10, f'Loss {loss.item():.4f} outside expected range (0, 10)'\n",
+    "\n",
+    "# Verify a full training step works end-to-end\n",
+    "params_before = [p.data.clone() for p in model_ag.parameters()]\n",
+    "opt_ag.zero_grad(set_to_none=True)\n",
+    "loss.backward()\n",
+    "torch.nn.utils.clip_grad_norm_(model_ag.parameters(), 1.0)\n",
+    "opt_ag.step()\n",
+    "model_ag.eval()\n",
+    "\n",
+    "changed = any(not torch.allclose(pb, pa)\n",
+    "              for pb, pa in zip(params_before, model_ag.parameters()))\n",
+    "assert changed, 'Model parameters did not change — check the gradient path in flow_loss'\n",
+    "\n",
+    "print(f'\\u2713 flow_loss | loss={loss.item():.4f}, training step works')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "aVjheN_wivWU"
+   },
+   "source": [
+    "---\n",
+    "## Part 4c — Fine-tuning on a new instrument  `[+0.5 pts bonus]`\n",
+    "### Download fine-tuning data\n",
+    "\n",
+    "We use `nsynth-valid` (~1.4 GB), which contains all instrument families.  \n",
+    "Use `instrument_filter` to select your target (default: `'guitar'`).  \n",
+    "Available families: `bass, brass, flute, guitar, keyboard, mallet, organ, reed, string, synth_lead, vocal`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "4fRu3Ba6ivWV"
+   },
+   "outputs": [],
+   "source": [
+    "DATA_ROOT = '/content/nsynth'\n",
+    "VALID_DIR = f'{DATA_ROOT}/nsynth-valid/audio'\n",
+    "\n",
+    "if not os.path.exists(VALID_DIR):\n",
+    "    print('Downloading nsynth-valid (~1.4 GB)...')\n",
+    "    !mkdir -p {DATA_ROOT}\n",
+    "    !wget -q http://download.magenta.tensorflow.org/datasets/nsynth/nsynth-valid.jsonwav.tar.gz \\\n",
+    "           -O /tmp/nsynth-valid.tar.gz\n",
+    "    !tar -xf /tmp/nsynth-valid.tar.gz -C {DATA_ROOT}\n",
+    "    !rm /tmp/nsynth-valid.tar.gz\n",
+    "    print('Done.')\n",
+    "else:\n",
+    "    print('nsynth-valid already present.')\n",
+    "\n",
+    "import glob\n",
+    "all_valid = glob.glob(f'{VALID_DIR}/*.wav')\n",
+    "print(f'Total valid files: {len(all_valid):,}')\n",
+    "for family in ['guitar', 'bass', 'flute', 'brass', 'reed', 'keyboard']:\n",
+    "    n = len([f for f in all_valid if os.path.basename(f).startswith(family)])\n",
+    "    print(f'  {family:12s}: {n:4d} files')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "cjUJ6TItivWV"
+   },
+   "outputs": [],
+   "source": [
+    "# ── YOUR CHOICES ──────────────────────────────────────────────────────────────\n",
+    "TARGET_INSTRUMENT = 'guitar'    # Change to any family listed above\n",
+    "FT_MAX_FILES      = 2000        # Number of files to use\n",
+    "FT_EPOCHS         = 300         # Training epochs\n",
+    "FT_LR             = 1e-3        # Learning rate\n",
+    "FT_BATCH_SIZE     = 64\n",
+    "# ─────────────────────────────────────────────────────────────────────────────\n",
+    "\n",
+    "import copy, time\n",
+    "\n",
+    "model_ft, ckpt_ft = load_flow_model(CKPT_PATH, device=device)\n",
+    "\n",
+    "dataset_ft = NSynthSpecDataset(VALID_DIR, instrument_filter=TARGET_INSTRUMENT,\n",
+    "                               max_files=FT_MAX_FILES)\n",
+    "loader_ft  = DataLoader(dataset_ft, batch_size=FT_BATCH_SIZE, shuffle=True,\n",
+    "                        drop_last=True, num_workers=2)\n",
+    "optimizer_ft = torch.optim.AdamW(model_ft.parameters(), lr=FT_LR, weight_decay=1e-4)\n",
+    "\n",
+    "print(f'Fine-tuning on {len(dataset_ft)} {TARGET_INSTRUMENT} files, '\n",
+    "      f'{len(loader_ft)} batches/epoch, {FT_EPOCHS} epochs')\n",
+    "\n",
+    "# ── Training loop — infrastructure is given; your sample_timesteps + flow_loss do the work ──\n",
+    "t0 = time.time()\n",
+    "for epoch in range(1, FT_EPOCHS + 1):\n",
+    "    model_ft.train()\n",
+    "    epoch_loss = []\n",
+    "    for x_data, pitch in loader_ft:\n",
+    "        x_data = x_data.to(device, non_blocking=True)\n",
+    "        pitch  = pitch.to(device, non_blocking=True)\n",
+    "\n",
+    "        t = sample_timesteps(x_data.shape[0], x_data.device)   # ← your function\n",
+    "\n",
+    "        optimizer_ft.zero_grad(set_to_none=True)\n",
+    "        loss = flow_loss(model_ft, x_data, pitch, t, p_uncond=0.1)  # ← your function\n",
+    "        loss.backward()\n",
+    "        torch.nn.utils.clip_grad_norm_(model_ft.parameters(), 1.0)\n",
+    "        optimizer_ft.step()\n",
+    "\n",
+    "        epoch_loss.append(loss.item())\n",
+    "\n",
+    "    if epoch % max(1, FT_EPOCHS // 10) == 0 or epoch == 1:\n",
+    "        print(f'Epoch {epoch:4d}/{FT_EPOCHS}  loss={np.mean(epoch_loss):.4f}  '\n",
+    "              f'elapsed={time.time()-t0:.0f}s')\n",
+    "\n",
+    "model_ft.eval()\n",
+    "print(f'\\nFine-tuning done in {(time.time()-t0)/60:.1f} min')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "X1DXJufSivWV"
+   },
+   "source": [
+    "### Generate & submit 100 samples\n",
+    "\n",
+    "Choose your best sampler and guidance scale.\n",
+    "**The submission must include the starting noise for each sample** so we can verify reproducibility."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "Oe-zPG3NivWV"
+   },
+   "outputs": [],
+   "source": [
+    "# ── YOUR CHOICES ──────────────────────────────────────────────────────────────\n",
+    "Q4_SAMPLER        = 'heun'   # 'euler' | 'cfg' | 'heun' | 'rk4'\n",
+    "Q4_GUIDANCE_SCALE = 6.0\n",
+    "Q4_N_STEPS        = 50\n",
+    "# ─────────────────────────────────────────────────────────────────────────────\n",
+    "\n",
+    "N_SUB = 100\n",
+    "# Spread pitches across 3 octaves (C3–B5)\n",
+    "q4_pitches = torch.tensor(\n",
+    "    [(48 + i % 36) for i in range(N_SUB)], dtype=torch.long, device=device)\n",
+    "\n",
+    "torch.manual_seed(0)\n",
+    "q4_noises = torch.randn(N_SUB, 2, FREQ_BINS, TIME_FRAMES, device=device)\n",
+    "\n",
+    "model_ft.eval()\n",
+    "q4_samples = []\n",
+    "BATCH = 16\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    for i in range(0, N_SUB, BATCH):\n",
+    "        x0b = q4_noises[i:i+BATCH]\n",
+    "        pb  = q4_pitches[i:i+BATCH]\n",
+    "        if Q4_SAMPLER == 'euler':\n",
+    "            out = euler_sample(model_ft, x0b.clone(), pb, n_steps=Q4_N_STEPS)\n",
+    "        elif Q4_SAMPLER == 'heun':\n",
+    "            out = heun_sample( model_ft, x0b.clone(), pb, n_steps=Q4_N_STEPS,\n",
+    "                               guidance_scale=Q4_GUIDANCE_SCALE)\n",
+    "        elif Q4_SAMPLER == 'rk4':\n",
+    "            out = rk4_sample(  model_ft, x0b.clone(), pb, n_steps=Q4_N_STEPS,\n",
+    "                               guidance_scale=Q4_GUIDANCE_SCALE)\n",
+    "        else:  # cfg\n",
+    "            out = cfg_sample(  model_ft, x0b.clone(), pb, n_steps=Q4_N_STEPS,\n",
+    "                               guidance_scale=Q4_GUIDANCE_SCALE)\n",
+    "        q4_samples.append(out.cpu())\n",
+    "\n",
+    "q4_samples = torch.cat(q4_samples)  # (100, 2, F, T)\n",
+    "print(f'Generated {len(q4_samples)} samples')\n",
+    "\n",
+    "# Listen to a few\n",
+    "for i in range(3):\n",
+    "    audio = spec_to_audio(q4_samples[i])\n",
+    "    audio = audio / (audio.abs().max() + 1e-8)\n",
+    "    print(f'Sample {i}, pitch={q4_pitches[i].item()}')\n",
+    "    display(ipd.Audio(audio.numpy(), rate=SR))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "SMoXWnVjivWV"
+   },
+   "outputs": [],
+   "source": [
+    "# Save submission for Q4\n",
+    "os.makedirs('/content', exist_ok=True)\n",
+    "\n",
+    "np.savez_compressed(\n",
+    "    '/content/submission_q4.npz',\n",
+    "    samples        = q4_samples.numpy().astype(np.float32),\n",
+    "    noises         = q4_noises.cpu().numpy().astype(np.float32),\n",
+    "    pitches        = q4_pitches.cpu().numpy().astype(np.int64),\n",
+    "    guidance_scale = np.array(Q4_GUIDANCE_SCALE, dtype=np.float32),\n",
+    "    n_steps        = np.array(Q4_N_STEPS,        dtype=np.int32),\n",
+    "    sampler        = np.array(Q4_SAMPLER,         dtype=object),\n",
+    ")\n",
+    "print('Saved: /content/submission_q4.npz')\n",
+    "print()\n",
+    "print('Submit the following files:')\n",
+    "print('  1. /content/submission_q4.npz')\n",
+    "print('  2. /content/model_ft_q4.pt  (your fine-tuned checkpoint)')\n",
+    "print('  3. This notebook (.ipynb)')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "q7D3qcb0ivWV"
+   },
+   "source": [
+    "---\n",
+    "## Part 5 — Beat the Baseline  `[+0.5 pts bonus]`\n",
+    "\n",
+    "**Baseline** (pretrained keyboard model, Heun 25 steps, guidance=6):  \n",
+    "- FD@6 ~ 354  \n",
+    "- Pitch class accuracy ~ 79%\n",
+    "\n",
+    "Achieve a **lower FD** or **higher pitch accuracy** using any approach:\n",
+    "\n",
+    "| Idea | Notes |\n",
+    "|---|---|\n",
+    "| More fine-tuning epochs | Longer training on same data |\n",
+    "| Different target instrument | Easier task = better FD |\n",
+    "| Better inference | Heun/RK4 vs. Euler; higher guidance scale |\n",
+    "| Mixed training | Fine-tune on multiple families |\n",
+    "| Train from scratch | Use `train_step` with a fresh model |\n",
+    "| Model architecture | Load a larger model config |\n",
+    "\n",
+    "Describe your approach in the markdown cell below, then generate and save 100 samples."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "zDj6OQj8ivWV"
+   },
+   "source": [
+    "### My approach for Part 5\n",
+    "\n",
+    "*(Replace this with your description: what did you try, what worked, and why?)*"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "Klx0o90GivWg"
+   },
+   "outputs": [],
+   "source": [
+    "# ── YOUR CODE HERE — no structure imposed ─────────────────────────────────────\n",
+    "# Train a better model (model_q5), then generate q5_samples and q5_noises below.\n",
+    "\n",
+    "# Example starting point — fine-tune longer on keyboard (the easiest task):\n",
+    "# ckpt_q5  = torch.load(CKPT_PATH, map_location=device, weights_only=False)\n",
+    "# model_q5 = build_model_from_config(ckpt_q5['config']).to(device)\n",
+    "# model_q5.load_state_dict(ckpt_q5['model_state'])\n",
+    "# ... your training loop ..."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "2Fy1RkuSivWg"
+   },
+   "outputs": [],
+   "source": [
+    "# Generate 100 samples from your best Q5 model\n",
+    "# ── YOUR CHOICES ──────────────────────────────────────────────────────────────\n",
+    "Q5_SAMPLER        = 'heun'\n",
+    "Q5_GUIDANCE_SCALE = 6.0\n",
+    "Q5_N_STEPS        = 50\n",
+    "# Use model_q5 (defined in cell above)\n",
+    "# ─────────────────────────────────────────────────────────────────────────────\n",
+    "\n",
+    "q5_pitches = torch.tensor(\n",
+    "    [(48 + i % 36) for i in range(N_SUB)], dtype=torch.long, device=device)\n",
+    "\n",
+    "torch.manual_seed(0)\n",
+    "q5_noises = torch.randn(N_SUB, 2, FREQ_BINS, TIME_FRAMES, device=device)\n",
+    "\n",
+    "# model_q5.eval()\n",
+    "# q5_samples = ... (same pattern as the Q4 generation cell)\n",
+    "\n",
+    "# Once generated, run the save cell below:\n",
+    "print('TODO: generate q5_samples (100, 2, FREQ_BINS, TIME_FRAMES) and q5_noises')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "T6LGCxGIivWg"
+   },
+   "outputs": [],
+   "source": [
+    "# Save submission for Q5\n",
+    "#os.makedirs('/content', exist_ok=True)\n",
+    "\n",
+    "Q5_CKPT_PATH = '/content/model_q5.pt'\n",
+    "# torch.save({'model_state': model_q5.state_dict(), 'config': ..., 'n_params': ...}, Q5_CKPT_PATH)\n",
+    "\n",
+    "np.savez_compressed(\n",
+    "    '/content/submission_q5.npz',\n",
+    "    samples        = q5_samples.numpy().astype(np.float32),   # fill in q5_samples\n",
+    "    noises         = q5_noises.cpu().numpy().astype(np.float32),\n",
+    "    pitches        = q5_pitches.cpu().numpy().astype(np.int64),\n",
+    "    guidance_scale = np.array(Q5_GUIDANCE_SCALE, dtype=np.float32),\n",
+    "    n_steps        = np.array(Q5_N_STEPS,        dtype=np.int32),\n",
+    "    sampler        = np.array(Q5_SAMPLER,         dtype=object),\n",
+    ")\n",
+    "print('Saved: /content/submission_q5.npz')\n",
+    "print()\n",
+    "print('Submit:')\n",
+    "print('  1. /content/submission_q5.npz')\n",
+    "print('  2. /content/model_q5.pt')\n",
+    "print('  3. This notebook (.ipynb)')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "ER4aIO5pivWg"
+   },
+   "source": [
+    "---\n",
+    "## Submission checklist\n",
+    "\n",
+    "Before submitting, run **Kernel -> Restart and Run All** to confirm everything executes\n",
+    "from scratch without errors. Then convert your notebook to a `.py` file:\n",
+    "\n",
+    "```bash\n",
+    "jupyter nbconvert --to script assignment.ipynb\n",
+    "```\n",
+    "\n",
+    "| Item | File | Required for |\n",
+    "|---|---|---|\n",
+    "| Converted notebook | `assignment.py` | All parts |\n",
+    "| Q4 samples + noises (bonus) | `submission_q4.npz` | Part 4 |\n",
+    "| Q4 checkpoint (bonus) | `model_ft_q4.pt` | Part 4 |\n",
+    "| Q5 samples + noises (bonus) | `submission_q5.npz` | Part 5 |\n",
+    "| Q5 checkpoint (bonus) | `model_q5.pt` | Part 5 |"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "2fDHwQCfivWg"
+   },
+   "source": []
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

HW4/model.py

@@ -0,0 +1,420 @@
+"""
+model.py — TinyFlowNet, UNet2DFlowNet, DiTFlowNet (all pitch-conditioned, CFG-ready)
+
+Time convention (via FlowModelWrapper)
+--------------------------------------
+Students interact with the model through ``FlowModelWrapper``, which uses
+the standard diffusion convention:
+
+    t = 1  →  pure noise          t = 0  →  clean data
+
+    model(x, t, pitch) → velocity pointing from noise toward data
+
+Generation integrates from t=1 to t=0:   x_{t−Δt} = x_t − v·Δt
+
+The raw network architectures below use the opposite internal convention
+(t=0 noise, t=1 data, v = data − noise).  The wrapper handles the mapping.
+
+Classifier-Free Guidance (CFG)
+-------------------------------
+Pitch index 128 is reserved as the null / unconditional token.
+
+Model overview
+--------------
+  TinyFlowNet   (~88k params)  : flat stack of ResBlocks, no downsampling
+  UNet2DFlowNet (~213k params) : 2-level encoder-decoder with skip connections
+  DiTFlowNet    (~221k params) : patch-based Diffusion Transformer (adaLN-Zero)
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+NULL_PITCH = 128   # reserved index for unconditional (CFG null token)
+
+
+# ── Shared building blocks ─────────────────────────────────────────────────────
+
+class SinusoidalEmbedding(nn.Module):
+    """
+    Fixed sinusoidal embedding of a scalar time value t ∈ [0, 1].
+    No learnable parameters — the MLP after it does the heavy lifting.
+    """
+
+    def __init__(self, dim: int):
+        super().__init__()
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(10_000) * torch.arange(half).float() / max(half - 1, 1)
+        )
+        self.register_buffer("freqs", freqs)
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        # t : (B,)
+        emb = t[:, None] * self.freqs[None]          # (B, half)
+        return torch.cat([emb.sin(), emb.cos()], -1)  # (B, dim)
+
+
+def _make_t_emb(t_dim: int) -> nn.Sequential:
+    """Shared time-embedding MLP: sinusoidal → 2-layer MLP → (B, t_dim)."""
+    return nn.Sequential(
+        SinusoidalEmbedding(t_dim),
+        nn.Linear(t_dim, t_dim * 2),
+        nn.SiLU(),
+        nn.Linear(t_dim * 2, t_dim),
+    )
+
+
+class ResBlock(nn.Module):
+    """
+    Pre-norm residual conv block with combined time+pitch conditioning.
+
+      x  ──►  GroupNorm ──► Conv ──► SiLU  ──► + cond_shift ──► GroupNorm ──► Conv ──► SiLU  ──► + x
+    """
+
+    def __init__(self, channels: int, t_dim: int, groups: int = 8):
+        super().__init__()
+        self.norm1  = nn.GroupNorm(groups, channels)
+        self.conv1  = nn.Conv2d(channels, channels, 3, padding=1)
+        self.t_proj = nn.Linear(t_dim, channels)     # conditioning → additive shift
+        self.norm2  = nn.GroupNorm(groups, channels)
+        self.conv2  = nn.Conv2d(channels, channels, 3, padding=1)
+        self.act    = nn.SiLU()
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        # cond: (B, t_dim) — combined time + pitch embedding
+        h  = self.act(self.conv1(self.norm1(x)))
+        h  = h + self.t_proj(self.act(cond))[:, :, None, None]  # broadcast over (F,T)
+        h  = self.act(self.conv2(self.norm2(h)))
+        return x + h
+
+
+# ── TinyFlowNet ────────────────────────────────────────────────────────────────
+
+class TinyFlowNet(nn.Module):
+    """
+    Predicts the vector field v_θ(x_t, t, pitch) for flow matching.
+
+    A flat stack of ResBlocks — no spatial downsampling.
+    Simple and fast; good baseline.
+
+    Default config (hidden=32, n_blocks=4, t_dim=32) ≈ 88k parameters.
+    """
+
+    def __init__(self, hidden: int = 32, n_blocks: int = 4, t_dim: int = 32):
+        super().__init__()
+        groups = min(8, hidden)
+
+        self.t_emb     = _make_t_emb(t_dim)
+        self.pitch_emb = nn.Embedding(NULL_PITCH + 1, t_dim)
+
+        self.input_proj = nn.Conv2d(2, hidden, 3, padding=1)
+        self.blocks = nn.ModuleList(
+            [ResBlock(hidden, t_dim, groups=groups) for _ in range(n_blocks)]
+        )
+        self.output_proj = nn.Sequential(
+            nn.GroupNorm(groups, hidden),
+            nn.SiLU(),
+            nn.Conv2d(hidden, 2, 1),
+        )
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor, pitch: torch.Tensor) -> torch.Tensor:
+        cond = self.t_emb(t) + self.pitch_emb(pitch)  # (B, t_dim)
+        h = self.input_proj(x)
+        for block in self.blocks:
+            h = block(h, cond)
+        return self.output_proj(h)
+
+
+# ── UNet2DFlowNet ──────────────────────────────────────────────────────────────
+
+class UNet2DFlowNet(nn.Module):
+    """
+    2D UNet vector-field network for flow matching on spectrograms.
+
+    Two spatial downsampling levels with skip connections:
+
+      Encoder:  [2→C] → ResBlock(C) ─── ↓ → ResBlock(C) ─── ↓ → ResBlock(2C)
+                                skip1 ↗              skip2 ↗
+      Decoder:  ↑+skip2 → merge(3C→C) → ResBlock(C) → ↑+skip1 → merge(2C→C) → ResBlock(C) → [C→2]
+
+    Bilinear upsampling (size read from skip tensor) handles odd input dimensions
+    (129 × 63) without size-mismatch issues.
+
+    Default config (hidden=32, t_dim=32): ~213k parameters.
+    """
+
+    def __init__(self, hidden: int = 32, t_dim: int = 32):
+        super().__init__()
+        C = hidden
+        g = min(8, C)
+
+        self.t_emb     = _make_t_emb(t_dim)
+        self.pitch_emb = nn.Embedding(NULL_PITCH + 1, t_dim)
+
+        # Encoder
+        self.input_proj = nn.Conv2d(2, C, 3, padding=1)
+        self.enc1       = ResBlock(C,     t_dim, groups=g)
+        self.down1      = nn.AvgPool2d(2)
+        self.chan_up1   = nn.Conv2d(C, C, 1)           # identity channel change (C→C)
+        self.enc2       = ResBlock(C,     t_dim, groups=g)
+        self.down2      = nn.AvgPool2d(2)
+        self.chan_up2   = nn.Conv2d(C, C * 2, 1)       # C → 2C at bottleneck
+        self.bottleneck = ResBlock(C * 2, t_dim, groups=min(8, C * 2))
+
+        # Decoder — merge convs reduce concatenated channels before ResBlock
+        self.merge1 = nn.Conv2d(C * 2 + C, C, 3, padding=1)   # cat(2C, C) → C
+        self.dec1   = ResBlock(C, t_dim, groups=g)
+        self.merge2 = nn.Conv2d(C + C, C, 3, padding=1)        # cat(C, C) → C
+        self.dec2   = ResBlock(C, t_dim, groups=g)
+
+        self.output_proj = nn.Sequential(
+            nn.GroupNorm(g, C),
+            nn.SiLU(),
+            nn.Conv2d(C, 2, 1),
+        )
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor, pitch: torch.Tensor) -> torch.Tensor:
+        cond = self.t_emb(t) + self.pitch_emb(pitch)  # (B, t_dim)
+
+        # Encoder
+        h  = self.input_proj(x)                         # (B, C, F, T)
+        s1 = self.enc1(h, cond)                         # (B, C, F, T)   — skip1
+        h  = self.chan_up1(self.down1(s1))               # (B, C, F//2, T//2)
+        s2 = self.enc2(h, cond)                         # (B, C, F//2, T//2) — skip2
+        h  = self.chan_up2(self.down2(s2))               # (B, 2C, F//4, T//4)
+        h  = self.bottleneck(h, cond)                    # (B, 2C, F//4, T//4)
+
+        # Decoder — upsample to match skip spatial size, cat, reduce, ResBlock
+        h = F.interpolate(h, size=s2.shape[2:], mode='bilinear', align_corners=False)
+        h = self.merge1(torch.cat([h, s2], dim=1))       # (B, C, F//2, T//2)
+        h = self.dec1(h, cond)
+
+        h = F.interpolate(h, size=s1.shape[2:], mode='bilinear', align_corners=False)
+        h = self.merge2(torch.cat([h, s1], dim=1))       # (B, C, F, T)
+        h = self.dec2(h, cond)
+
+        return self.output_proj(h)                        # (B, 2, F, T)
+
+
+# ── DiTFlowNet ─────────────────────────────────────────────────────────────────
+
+class DiTBlock(nn.Module):
+    """
+    Diffusion Transformer block with adaLN-Zero conditioning.
+
+    Given a conditioning vector cond ∈ R^{t_dim}, a learned MLP produces six
+    per-sample parameters (scale1, shift1, gate1, scale2, shift2, gate2) that
+    modulate the attention and FFN sublayers independently.
+
+    The final linear in the adaLN MLP is zero-initialized so each block starts
+    as a near-identity residual connection (gate=0, scale≈1, shift≈0).
+    """
+
+    def __init__(self, d_model: int, n_heads: int, t_dim: int, ffn_mult: int = 4):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(d_model, elementwise_affine=False)
+        self.norm2 = nn.LayerNorm(d_model, elementwise_affine=False)
+        self.attn  = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
+        self.ffn   = nn.Sequential(
+            nn.Linear(d_model, ffn_mult * d_model),
+            nn.GELU(),
+            nn.Linear(ffn_mult * d_model, d_model),
+        )
+        # adaLN-Zero: (B, t_dim) → 6 × (B, d_model) for scale/shift/gate × 2 sublayers
+        self.adaLN_mlp = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(t_dim, 6 * d_model),
+        )
+        nn.init.zeros_(self.adaLN_mlp[-1].weight)
+        nn.init.zeros_(self.adaLN_mlp[-1].bias)
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        # cond: (B, t_dim);  x: (B, n_tokens, d_model)
+        g1, b1, a1, g2, b2, a2 = self.adaLN_mlp(cond).chunk(6, dim=-1)
+
+        # Attention sub-block
+        h = (1 + g1[:, None]) * self.norm1(x) + b1[:, None]
+        h, _ = self.attn(h, h, h)
+        x = x + a1[:, None] * h
+
+        # FFN sub-block
+        h = (1 + g2[:, None]) * self.norm2(x) + b2[:, None]
+        h = self.ffn(h)
+        x = x + a2[:, None] * h
+
+        return x
+
+
+class DiTFlowNet(nn.Module):
+    """
+    Diffusion Transformer vector-field network for flow matching on spectrograms.
+
+    Patchifies the (2, freq_bins, time_frames) input into tokens, applies N
+    transformer blocks with adaLN-Zero conditioning, then unpatches back.
+
+    Input padding: the spectrogram is zero-padded to the nearest multiple of
+    patch_size in each spatial dimension before patchification and cropped back
+    to the original size at output.
+
+    Default config (d_model=64, n_layers=3, patch_size=8): ~221k parameters.
+    For (2, 129, 63): pads to (2, 136, 64) → 17×8 = 136 tokens, patch_dim=128.
+
+    Parameters
+    ----------
+    freq_bins   : input frequency dimension (e.g. 129)
+    time_frames : input time dimension (e.g. 63)
+    d_model     : transformer hidden dimension
+    n_layers    : number of DiT blocks
+    n_heads     : attention heads (must divide d_model)
+    t_dim       : conditioning embedding dimension
+    patch_size  : spatial patch size applied to both freq and time axes
+    """
+
+    def __init__(
+        self,
+        freq_bins: int  = 129,
+        time_frames: int = 63,
+        d_model: int    = 64,
+        n_layers: int   = 3,
+        n_heads: int    = 4,
+        t_dim: int      = 32,
+        patch_size: int = 8,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        p = patch_size
+        patch_dim = 2 * p * p   # 2 channels × p × p pixels per patch
+
+        # Number of tokens for the fixed-size spectrograms
+        nf = math.ceil(freq_bins  / p)
+        nt = math.ceil(time_frames / p)
+        n_tokens = nf * nt
+
+        self.t_emb     = _make_t_emb(t_dim)
+        self.pitch_emb = nn.Embedding(NULL_PITCH + 1, t_dim)
+
+        self.patch_embed = nn.Linear(patch_dim, d_model)
+        self.pos_embed   = nn.Parameter(torch.zeros(1, n_tokens, d_model))
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+
+        self.blocks = nn.ModuleList([
+            DiTBlock(d_model, n_heads, t_dim) for _ in range(n_layers)
+        ])
+        self.norm         = nn.LayerNorm(d_model)
+        self.unpatch_proj = nn.Linear(d_model, patch_dim, bias=False)
+
+    def _patchify(self, x: torch.Tensor) -> tuple:
+        """(B, 2, freq, time) → (B, nf*nt, 2*p*p)"""
+        B, C, freq, time = x.shape
+        p = self.patch_size
+        pad_f = (-freq) % p
+        pad_t = (-time) % p
+        if pad_f or pad_t:
+            x = F.pad(x, (0, pad_t, 0, pad_f))
+        _, _, Fp, Tp = x.shape
+        nf, nt = Fp // p, Tp // p
+        # (B, C, nf, p, nt, p) → (B, nf, nt, C, p, p) → (B, nf*nt, C*p*p)
+        x = x.reshape(B, C, nf, p, nt, p)
+        x = x.permute(0, 2, 4, 1, 3, 5).reshape(B, nf * nt, C * p * p)
+        return x, (freq, time, nf, nt)
+
+    def _unpatchify(self, x: torch.Tensor, freq_orig: int, time_orig: int,
+                    nf: int, nt: int) -> torch.Tensor:
+        """(B, nf*nt, 2*p*p) → (B, 2, freq_orig, time_orig)"""
+        B = x.shape[0]
+        p = self.patch_size
+        x = x.reshape(B, nf, nt, 2, p, p)
+        x = x.permute(0, 3, 1, 4, 2, 5).reshape(B, 2, nf * p, nt * p)
+        return x[:, :, :freq_orig, :time_orig]
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor, pitch: torch.Tensor) -> torch.Tensor:
+        cond = self.t_emb(t) + self.pitch_emb(pitch)          # (B, t_dim)
+        tokens, (freq_orig, time_orig, nf, nt) = self._patchify(x)
+        tokens = self.patch_embed(tokens) + self.pos_embed     # (B, n_tokens, d_model)
+        for block in self.blocks:
+            tokens = block(tokens, cond)
+        tokens = self.unpatch_proj(self.norm(tokens))          # (B, n_tokens, patch_dim)
+        return self._unpatchify(tokens, freq_orig, time_orig, nf, nt)
+
+
+# ── Flow model wrapper (diffusion convention) ────────────────────────────────
+
+class FlowModelWrapper(nn.Module):
+    """Wraps a raw flow model to use the standard diffusion time convention:
+
+        t = 1  →  pure noise
+        t = 0  →  clean data
+
+    The wrapped model's ``forward(x, t, pitch)`` returns the velocity field
+    pointing from noise toward data, so that generation integrates from t=1
+    down to t=0 via  x_{t-Δt} = x_t − v·Δt.
+
+    Internally the raw network was trained with the opposite convention
+    (t=0 = noise, t=1 = data, velocity = data − noise), so the wrapper
+    simply flips time and negates the output.  Gradients flow through
+    correctly, so fine-tuning works as expected.
+    """
+
+    def __init__(self, inner: nn.Module):
+        super().__init__()
+        self.inner = inner
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor,
+                pitch: torch.Tensor) -> torch.Tensor:
+        return -self.inner(x, 1.0 - t, pitch)
+
+
+# ── Utilities ──────────────────────────────────────────────────────────────────
+
+def count_params(model: nn.Module) -> int:
+    inner = model.inner if isinstance(model, FlowModelWrapper) else model
+    return sum(p.numel() for p in inner.parameters())
+
+
+def build_model_from_config(cfg: dict) -> nn.Module:
+    """Reconstruct the correct model class from a saved checkpoint config dict."""
+    model_type = cfg.get("model_type", "tiny")
+    if model_type == "tiny":
+        return TinyFlowNet(
+            hidden=cfg["hidden"], n_blocks=cfg["n_blocks"], t_dim=cfg["t_dim"]
+        )
+    elif model_type == "unet":
+        return UNet2DFlowNet(hidden=cfg["hidden"], t_dim=cfg["t_dim"])
+    elif model_type == "dit":
+        return DiTFlowNet(
+            freq_bins=cfg["freq_bins"], time_frames=cfg["time_frames"],
+            d_model=cfg["d_model"], n_layers=cfg["n_layers"],
+            n_heads=cfg["n_heads"], t_dim=cfg["t_dim"],
+            patch_size=cfg["patch_size"],
+        )
+    else:
+        raise ValueError(f"Unknown model_type: {model_type!r}")
+
+
+def load_flow_model(ckpt_path: str, device: str = "cpu"):
+    """Load a checkpoint and return ``(wrapped_model, ckpt_dict)``.
+
+    The returned model uses the standard diffusion convention
+    (t=1 noise, t=0 data).
+    """
+    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+    raw  = build_model_from_config(ckpt["config"]).to(device)
+    raw.load_state_dict(ckpt["model_state"])
+    model = FlowModelWrapper(raw)
+    model.eval()
+    return model, ckpt
+
+
+def save_flow_model(model: nn.Module, path: str, config: dict,
+                    n_params: int, **extra):
+    """Save a model checkpoint (unwraps ``FlowModelWrapper`` automatically)."""
+    inner = model.inner if isinstance(model, FlowModelWrapper) else model
+    torch.save({
+        "model_state": inner.state_dict(),
+        "config":      config,
+        "n_params":    n_params,
+        **extra,
+    }, path)

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-021-025.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8878d4a2bc99e55cc4caedbc0d9c087f06841982fbd90b5f679538fd07541242
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-021-050.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c71e53b0fee3f5f223b6cbc42e35d2dce096b9e7db24718736df699f00bcb06f
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-021-075.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5c87d4471c7d92ab80df60bf9fa55df5e7d7af8a4e4449ff5ee926725f0e8ab5
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-021-100.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1cee6413154b1cd7a607305810027720245062380d762e4638568eecb9d17a59
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-022-025.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ff6a602f7e10ce9a36243c7833a40100f2429faa24407c270fa2e52a4b600431
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-022-100.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:545ac61d631271b9fd595ea04618584b642a5ec1925af105f6bbbad995e8f59a
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-023-050.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67a013fa74ce17091acaf7b7d448fd98c265401c00378518140f21b7b610ee8a
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-023-075.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d0ffb9cfc439c6c303f896c766236ef141a4812546e7e1d6010e26326b5f1f6e
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-023-100.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6bf07c752febb19443cf2366c2522a0f1b4c48eb806c778549f89ba8497ba5f0
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-025-050.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:242fec67dad566f77a3cf53c7ebcb5f6ca1dab3cc2b7103e549bbf7c0ffb4b58
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-025-075.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bc7c56c199166cd5196100fc8c3da5f13989bf4427119dae077aee4c0d208714
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-025-100.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b547c20c61a33cc76a4e7466de2a5a9f946c62c0bba442b7e48f1aa2c5a92b2
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-025-127.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd755b0c82c13b9231b676146ec14db493b8a6f56ba193ff0d4ed6716c17dc69
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-026-025.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:465dfa073255429136ce6def173a94d94bcde117683531453224da42a62f962f
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-026-050.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:96cba7b434484d657005640f8f9697c82580a1d9358748fa26da773c59172fe7
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-026-075.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e6ba26b58b42e128232911ca75194aeef8a0f9cc6b1cc0712a8d14e9ffd49b8
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-026-100.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e43baedcf1f9e5769fdf89b4315865554ce6b706c0e723609c15ba5df207932e
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-026-127.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ac5d11098ee7a53d01fe30fc360dd76cdd2dadaf31873680752dda4dfd19fad
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-027-025.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1963553d755aa98e0d26e778f2f2553441010cf5f617458bdefafc999027afad
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-027-050.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e23596b51eb46b6d408ad2d2f7862239f6c7a03a1645adb4493f2f5e1ae4507
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-027-075.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c90f36d35cd28b7b6d9ff60dc1e7def761136651d01d529b761cc1d78515943e
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-027-100.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b105ee0272403fc44c83a850716d0eb83314646f73c625ff63472c5690ce67f
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-028-050.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:867828a89d3e61028ba1ec063a5ce84d0d8191a4568af24e2ca85215907bb6e9
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-028-075.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:beb43834ca548705d85063b3c48a7e5239abe2ec7f0f215d999710058c6ae988
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-028-100.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f042c649fbb4476b98656b8590d418dbd9fecf8b98ae0533f40354593105fba6
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-028-127.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:569775184581da8bb6f1605c17c254cfc01747ef5ed99a01aba674d2ed1fbe5b
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-029-025.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:70fb4ae6b04cd023d3a927c8b326433b0d5d249e7fe8e3bceb5996d441a5ef1d
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-029-050.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:44aec925810c4393aebb3517b374d2a6884f952e01fbe50809daad736bfc6e7b
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-029-075.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e6787498992e0aadba4894c8cca5789cd0459300c8d974f0797a717caf24228b
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-029-127.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:01987bdb4b1bb7d760dc0903bb9df307504ba122ec64062e7b8f0cd071fe465e
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-030-025.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3cf42c523c34dc361ec67183e0e356d7e3b2c96ae17eea3962508dbec73de533
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-030-075.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c75597b0552c8f15670aef064eea0fc9111763c38875edec5bcbe3ae1f4954bb
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-030-100.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a0e2c619af19310d070414a2a053ccfff8e23b50bdd3c1f19e54c381a678b9a9
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-030-127.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c9143cd6f9d5bf5a526859e75a29f5a95320deb8e4eb656941f0539c6db0a47
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-031-025.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b80f106368887833481473009d4cfbe06edf9805d60a7d14c35944a048f55006
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-031-050.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:06c5452658d1a43e3328ccb8b7a4964936262089f843dec9048cf66a980f5b14
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-031-075.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2da41595e426fd863b0cb4082957b7297bd6fbfde46945eb79d998e4f0891816
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-031-127.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2062daf4f1d4e69fd2a69fe5e7627b0266e77779e596810d46e843d4540f5e6d
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-032-025.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:36653c96644a42f5f63ba147a42c3218e5d3fb64023c5b3684bba92655c420a9
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-032-050.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d2a03506783b8513dda10375a77c35345f92536744af34913e963a7d9233f86
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-032-075.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fff2f2262a8afe02134bee1cdd636196be92d957523e47b918bce43e3d3aefc6
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-032-100.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0def44b1c20130dc273900e44468cd3c70be5816b37377e78fc43d3726fc915c
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-032-127.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d1cc63a3b8901f5718d64b7fa2c3069bf40c12cf41d9e508f6f9c3acd31ab7a6
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-033-025.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2591be421a1a4d52febe4dd7e003cf8070c760d4f5b1f32e51ee9cce7d72e2e3
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-033-050.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d1bd74503cf2a9a924b9275a675231e017395a8f06689555b07685c7f80a0893
+size 128044

HW4/nsynth/nsynth-valid/audio/guitar_acoustic_010-033-075.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:69dcb55df9c4c287b6a8c5651fdfbabcd5f0bf7e71caf9a09b741da97c6c7098
+size 128044