{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import IPython\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "import math\n",
    "\n",
    "def mdct(x):\n",
    "    \"\"\"\n",
    "    Performs Modified Discrete Cosine Transform (MDCT)\n",
    "    x: input tensor of shape (..., N) where N is frame size\n",
    "    returns: MDCT coefficients of shape (..., N/2)\n",
    "    \"\"\"\n",
    "    N = x.shape[-1]\n",
    "    n = torch.arange(N, device=x.device)\n",
    "    k = torch.arange(N // 2, device=x.device)\n",
    "    \n",
    "    # Corrected MDCT matrix\n",
    "    # The phase term is (2n + 1 + N/2) * (2k + 1) * π/(2N)\n",
    "    arg = (math.pi / (2 * N)) * ((2 * n + 1 + N//2).view(-1, 1)) * ((2 * k + 1).view(1, -1))\n",
    "    mdct_matrix = torch.cos(arg) * (2.0 / N)**0.5\n",
    "    \n",
    "    return torch.matmul(x, mdct_matrix)\n",
    "\n",
    "def imdct(X):\n",
    "    \"\"\"\n",
    "    Performs Inverse Modified Discrete Cosine Transform (IMDCT)\n",
    "    X: input tensor of shape (..., N/2) where N is frame size\n",
    "    returns: IMDCT output of shape (..., N)\n",
    "    \"\"\"\n",
    "    half_N = X.shape[-1]\n",
    "    N = half_N * 2\n",
    "    n = torch.arange(N, device=X.device)\n",
    "    k = torch.arange(half_N, device=X.device)\n",
    "    \n",
    "    # Corrected IMDCT matrix with proper phase term\n",
    "    arg = (math.pi / (2 * N)) * ((2 * n + 1 + N//2).view(-1, 1)) * ((2 * k + 1).view(1, -1))\n",
    "    imdct_matrix = torch.cos(arg) * (2.0 / N)**0.5\n",
    "    \n",
    "    return torch.matmul(X, imdct_matrix.T) * 2.0\n",
    "\n",
    "def audio_to_mdct_frames(audio, frame_size=1920, midside=False):\n",
    "    \"\"\"\n",
    "    Convert stereo audio to overlapped MDCT coefficients with optional mid-side processing.\n",
    "    \n",
    "    Args:\n",
    "        audio: Tensor of shape (batch, 2, samples)\n",
    "        frame_size: Size of each frame (default 1920)\n",
    "        midside: If True, process as mid/side channels instead of left/right (default False)\n",
    "    Returns:\n",
    "        Tensor of MDCT coefficients (batch, 2, n_frames, frame_size//2)\n",
    "    \"\"\"\n",
    "    batch, channels, samples = audio.shape\n",
    "    hop_size = frame_size // 2\n",
    "    \n",
    "    # Convert to mid-side if requested\n",
    "    if midside:\n",
    "        mid = (audio[:, 0, :] + audio[:, 1, :]) / 2.0\n",
    "        side = (audio[:, 0, :] - audio[:, 1, :]) / 2.0\n",
    "        audio = torch.stack([mid, side], dim=1)\n",
    "    \n",
    "    # Calculate number of full frames\n",
    "    n_frames = (samples - frame_size) // hop_size + 1\n",
    "    \n",
    "    # Create overlapping frames\n",
    "    frames = []\n",
    "    for i in range(n_frames):\n",
    "        start = i * hop_size\n",
    "        frame = audio[:, :, start:start + frame_size]\n",
    "        if frame.shape[-1] == frame_size:\n",
    "            frames.append(frame)\n",
    "    \n",
    "    frames = torch.stack(frames, dim=2)\n",
    "    \n",
    "    # Apply sine window\n",
    "    window = torch.sin(torch.pi / frame_size * (torch.arange(frame_size, device=audio.device) + 0.5))\n",
    "    frames = frames * window.view(1, 1, 1, -1)\n",
    "    \n",
    "    # Reshape and apply MDCT\n",
    "    shape = frames.shape\n",
    "    frames_reshaped = frames.reshape(-1, frame_size)\n",
    "    mdct_coeffs = mdct(frames_reshaped)\n",
    "    \n",
    "    # Reshape back\n",
    "    mdct_coeffs = mdct_coeffs.reshape(shape[0], shape[1], shape[2], -1)\n",
    "    \n",
    "    return mdct_coeffs\n",
    "\n",
    "def mdct_frames_to_audio_slow(mdct_coeffs, frame_size=1920, midside=False):\n",
    "    \"\"\"\n",
    "    Reconstruct audio from MDCT coefficients with optional mid-side processing.\n",
    "    \n",
    "    Args:\n",
    "        mdct_coeffs: Tensor of shape (batch, 2, n_frames, frame_size//2)\n",
    "        frame_size: Size of each frame (default 1920)\n",
    "        midside: If True, assume coefficients are mid/side encoded (default False)\n",
    "    Returns:\n",
    "        Reconstructed audio tensor of shape (batch, 2, samples)\n",
    "    \"\"\"\n",
    "    batch, channels, n_frames, half_frame_size = mdct_coeffs.shape\n",
    "    hop_size = frame_size // 2\n",
    "    \n",
    "    # Reshape and apply IMDCT\n",
    "    shape = mdct_coeffs.shape\n",
    "    coeffs_reshaped = mdct_coeffs.reshape(-1, half_frame_size)\n",
    "    frames = imdct(coeffs_reshaped)\n",
    "    frames = frames.reshape(shape[0], shape[1], shape[2], -1)\n",
    "    \n",
    "    # Apply synthesis window\n",
    "    window = torch.sin(torch.pi / frame_size * (torch.arange(frame_size, device=mdct_coeffs.device) + 0.5))\n",
    "    frames = frames * window.view(1, 1, 1, -1)\n",
    "    \n",
    "    # Overlap-add reconstruction\n",
    "    total_samples = (n_frames - 1) * hop_size + frame_size\n",
    "    output = torch.zeros(batch, channels, total_samples, device=mdct_coeffs.device)\n",
    "    \n",
    "    for i in range(n_frames):\n",
    "        start = i * hop_size\n",
    "        output[:, :, start:start + frame_size] += frames[:, :, i]\n",
    "    \n",
    "    # Convert back from mid-side to stereo if needed\n",
    "    if midside:\n",
    "        mid = output[:, 0, :]\n",
    "        side = output[:, 1, :]\n",
    "        left = mid + side\n",
    "        right = mid - side\n",
    "        output = torch.stack([left, right], dim=1)\n",
    "    \n",
    "    return output\n",
    "\n",
    "def mdct_frames_to_audio(mdct_coeffs, frame_size=1920, midside=False):\n",
    "    batch, channels, n_frames, half_frame_size = mdct_coeffs.shape\n",
    "    hop_size = frame_size // 2\n",
    "\n",
    "    # Reshape and apply IMDCT\n",
    "    shape = mdct_coeffs.shape\n",
    "    coeffs_reshaped = mdct_coeffs.reshape(-1, half_frame_size)\n",
    "    frames = imdct(coeffs_reshaped)\n",
    "    frames = frames.reshape(shape[0], shape[1], shape[2], -1)\n",
    "\n",
    "    # Apply window\n",
    "    window = torch.sin(\n",
    "        torch.pi\n",
    "        / frame_size\n",
    "        * (torch.arange(frame_size, device=mdct_coeffs.device) + 0.5)\n",
    "    )\n",
    "    frames = frames * window.view(1, 1, 1, -1)\n",
    "\n",
    "    # Calculate total samples and create output shape\n",
    "    total_samples = (n_frames - 1) * hop_size + frame_size\n",
    "\n",
    "    # Reshape frames to prepare for folding\n",
    "    frames = frames.permute(0, 1, 3, 2)  # [batch, channels, frame_size, n_frames]\n",
    "    frames = frames.reshape(batch * channels, frame_size, n_frames)\n",
    "\n",
    "    # Use fold operation to overlap-add frames\n",
    "    output = torch.nn.functional.fold(\n",
    "        frames,\n",
    "        output_size=(1, total_samples),\n",
    "        kernel_size=(1, frame_size),\n",
    "        stride=(1, hop_size),\n",
    "    )\n",
    "\n",
    "    # Reshape output to expected dimensions\n",
    "    output = output.view(batch, channels, total_samples)\n",
    "\n",
    "    if midside:\n",
    "        mid = output[:, 0, :]\n",
    "        side = output[:, 1, :]\n",
    "        left = mid + side\n",
    "        right = mid - side\n",
    "        output = torch.stack([left, right], dim=1)\n",
    "\n",
    "    return output\n",
    "\n",
    "def sparsify_mdct_frames(\n",
    "    frames: torch.Tensor,\n",
    "    energy_threshold: float = 0.01,\n",
    "    min_kept_coeffs: int = 100,\n",
    "    freq_weighting: bool = True,\n",
    "    smooth_window: int = 3\n",
    "):\n",
    "    \"\"\"\n",
    "    Sparsify MDCT frames by zeroing out less perceptually significant coefficients.\n",
    "    Inspired by MP3's psychoacoustic model but simplified.\n",
    "    \n",
    "    Args:\n",
    "        frames: Tensor of shape (batch_size, channels, seq_len, coeffs)\n",
    "        energy_threshold: Threshold for coefficient energy (relative to max)\n",
    "        min_kept_coeffs: Minimum number of coefficients to keep per frame\n",
    "        freq_weighting: Whether to apply frequency-dependent weighting\n",
    "        smooth_window: Size of window for smoothing energy calculation\n",
    "    \n",
    "    Returns:\n",
    "        sparsified_frames: Frames with less significant coefficients zeroed\n",
    "        mask: Boolean mask indicating which coefficients were kept\n",
    "    \"\"\"\n",
    "    bs, chs, seq_len, coeffs = frames.shape\n",
    "    \n",
    "    # Reshape to 2D: (batch * channels * seq_len, coeffs)\n",
    "    x = frames.reshape(-1, coeffs)\n",
    "    \n",
    "    # Calculate energy per coefficient\n",
    "    energy = torch.mean(x.pow(2), dim=0)\n",
    "    \n",
    "    if freq_weighting:\n",
    "        # Apply basic frequency weighting inspired by human hearing\n",
    "        # More sensitive in mid-range frequencies (roughly 2-5 kHz)\n",
    "        freq_weights = torch.ones_like(energy)\n",
    "        mid_range_start = coeffs // 4\n",
    "        mid_range_end = coeffs // 2\n",
    "        freq_weights[mid_range_start:mid_range_end] = 1.5\n",
    "        energy = energy * freq_weights\n",
    "    \n",
    "    if smooth_window > 1:\n",
    "        # Smooth energy across neighboring frequency bins\n",
    "        padding = (smooth_window - 1) // 2\n",
    "        energy = torch.nn.functional.avg_pool1d(\n",
    "            energy.unsqueeze(0).unsqueeze(0),\n",
    "            kernel_size=smooth_window,\n",
    "            padding=padding,\n",
    "            stride=1\n",
    "        ).squeeze()\n",
    "    \n",
    "    max_energy = torch.max(energy)\n",
    "    \n",
    "    # Create mask for significant coefficients\n",
    "    mask = energy > (max_energy * energy_threshold)\n",
    "    \n",
    "    # Ensure we keep at least min_kept_coeffs coefficients\n",
    "    if torch.sum(mask) < min_kept_coeffs:\n",
    "        _, top_indices = torch.topk(energy, min_kept_coeffs)\n",
    "        mask = torch.zeros_like(mask, dtype=torch.bool)\n",
    "        mask[top_indices] = True\n",
    "    \n",
    "    # Apply mask\n",
    "    sparsified = x.clone()\n",
    "    sparsified[:, ~mask] = 0\n",
    "    \n",
    "    # Reshape back to original dimensions\n",
    "    sparsified_frames = sparsified.reshape(bs, chs, seq_len, coeffs)\n",
    "    \n",
    "    return sparsified_frames, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# what we want to do\n",
    "# load a smallish dataset of audio files\n",
    "# for each file, load it, then convert to mdct frames\n",
    "# from this we build a small dataset of mdct frames\n",
    "# we then train a simple PCA model on this dataset\n",
    "# then test reconstruction quality of unseen data\n",
    "\n",
    "manifest_filepath = \"/app/suno/data/audio_2ch_48khz_lg/metas_tr.jsonl\"\n",
    "manifest = read_jsonl(manifest_filepath)\n",
    "print(len(manifest))\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cut down the training set to 100 files\n",
    "manifest_subset = manifest[:100]\n",
    "\n",
    "for meta in manifest_subset:\n",
    "    audio_filepath = meta[\"filepath\"]\n",
    "    audio, sr = torchaudio.load(audio_filepath)\n",
    "    mdct_coeffs = audio_to_mdct_frames(audio.unsqueeze(0))\n",
    "    sparsified_coeffs, mask = sparsify_mdct_frames(mdct_coeffs)\n",
    "    print(mdct_coeffs.shape)\n",
    "    print(sparsified_coeffs.shape)\n",
    "    break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from torch.utils.data import IterableDataset\n",
    "\n",
    "class AudioBatchGenerator(IterableDataset):\n",
    "    def __init__(self, metas: list, batch_size: int, midside=False, sparsify=False):\n",
    "        self.metas = metas\n",
    "        self.batch_size = batch_size\n",
    "        self.midside = midside\n",
    "        self.sparsify = sparsify\n",
    "        self.current_idx = 0\n",
    "        self.total_batches = len(metas) // batch_size + (1 if len(metas) % batch_size else 0)\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.total_batches\n",
    "        \n",
    "    def __iter__(self):\n",
    "        return self\n",
    "        \n",
    "    def __next__(self):\n",
    "        if self.current_idx >= len(self.metas):\n",
    "            self.current_idx = 0\n",
    "            raise StopIteration\n",
    "            \n",
    "        batch_tensors = []\n",
    "        batch_end = min(self.current_idx + self.batch_size, len(self.metas))\n",
    "        \n",
    "        for meta in self.metas[self.current_idx:batch_end]:\n",
    "            waveform, _ = torchaudio.load(meta[\"filepath\"])\n",
    "            waveform = waveform[:, :48000*30]\n",
    "            mdct_frames = audio_to_mdct_frames(waveform.unsqueeze(0), midside=self.midside)  # Shape: [2, frames, 960]\n",
    "\n",
    "            if self.sparsify:\n",
    "                sparsified_coeffs, mask = sparsify_mdct_frames(mdct_frames)\n",
    "                mdct_frames = sparsified_coeffs\n",
    "\n",
    "            mdct_flat = torch.cat((mdct_frames[:, 0, :, :], mdct_frames[:, 1, :, :]), dim=-1)\n",
    "            batch_tensors.append(mdct_flat.squeeze(0))\n",
    "            \n",
    "        self.current_idx = batch_end\n",
    "        return torch.cat(batch_tensors, dim=0)  # Shape: [total_frames, 1920]\n",
    "\n",
    "batch_generator = AudioBatchGenerator(manifest_subset, batch_size=10, midside=True)\n",
    "\n",
    "for batch in batch_generator:\n",
    "    print(batch.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "from sklearn.decomposition import IncrementalPCA\n",
    "\n",
    "n_components = 512\n",
    "batch_size = 8\n",
    "midside = True\n",
    "\n",
    "ipca = IncrementalPCA(n_components=n_components, batch_size=batch_size)\n",
    "generator = AudioBatchGenerator(manifest_subset, batch_size, midside=midside)\n",
    "\n",
    "# Progress tracking\n",
    "total_files = len(manifest_subset)\n",
    "pbar = tqdm(generator, total=len(generator), \n",
    "            desc=f\"Fitting PCA ({total_files} files)\",\n",
    "            unit=\"batch\")\n",
    "\n",
    "frames_processed = 0\n",
    "total_frames = 0\n",
    "\n",
    "for batch in pbar:\n",
    "    batch_np = batch.numpy().astype(np.float32)\n",
    "    ipca.partial_fit(batch_np)\n",
    "    \n",
    "    # Update progress metrics\n",
    "    frames_processed += batch.shape[0]\n",
    "    total_frames += batch.shape[0]\n",
    "    \n",
    "    # Update progress bar postfix\n",
    "    pbar.set_postfix({\n",
    "        'frames': f\"{frames_processed:,}\",\n",
    "        'var_explained': f\"{ipca.explained_variance_ratio_.sum():.3f}\"\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "import torchaudio\n",
    "# test on new audio \n",
    "test_filepath = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n",
    "test_filepath = \"/home/christian/audio/50_genre_songs/Miles Davis - Freddie Freeloader (Official Audio).mp3\"\n",
    "sparsify = True\n",
    "midside = True\n",
    "waveform, sr = torchaudio.load(test_filepath)\n",
    "\n",
    "if sr != 48000:\n",
    "    waveform = torchaudio.transforms.Resample(sr, 48000)(waveform)\n",
    "\n",
    "waveform = waveform[:, :48000*30]\n",
    "\n",
    "mdct_frames = audio_to_mdct_frames(waveform.unsqueeze(0))\n",
    "\n",
    "print(mdct_frames.shape)\n",
    "print(mdct_frames[0, 0, 1000, :])\n",
    "\n",
    "if sparsify:\n",
    "    sparsified_coeffs, mask = sparsify_mdct_frames(mdct_frames)\n",
    "    mdct_frames = sparsified_coeffs\n",
    "\n",
    "print(mdct_frames.shape)\n",
    "mdct_flat = torch.cat((mdct_frames[:, 0, :], mdct_frames[:, 1, :]), dim=-1)\n",
    "print(mdct_flat.shape)\n",
    "#compressed_frames = ipca.transform(mdct_flat.squeeze(0).numpy())\n",
    "#print(compressed_frames.shape)\n",
    "#uncompressed_frames = ipca.inverse_transform(compressed_frames)#\n",
    "#print(uncompressed_frames.shape)\n",
    "\n",
    "#uncompressed_frames = torch.from_numpy(mdct_flat).float()\n",
    "uncompressed_frames = mdct_flat\n",
    "# unfold the frames into the original 2 channel waveform\n",
    "uncompressed_frames = torch.stack((uncompressed_frames[..., :960], uncompressed_frames[..., 960:]), dim=0)\n",
    "print(uncompressed_frames.shape)\n",
    "uncompressed_waveform = mdct_frames_to_audio(uncompressed_frames, midside=midside)\n",
    "print(uncompressed_waveform.shape)\n",
    "\n",
    "IPython.display.display(IPython.display.Audio(uncompressed_waveform.squeeze(0).numpy(), rate=48000))\n",
    "IPython.display.display(IPython.display.Audio(waveform.squeeze(0).numpy(), rate=48000))\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import math\n",
    "import torch\n",
    "\n",
    "def mdct(x):\n",
    "    N = x.shape[-1]\n",
    "    n = torch.arange(N, device=x.device)\n",
    "    k = torch.arange(N // 2, device=x.device)\n",
    "\n",
    "    arg = (\n",
    "        (math.pi / (2 * N))\n",
    "        * ((2 * n + 1 + N // 2).view(-1, 1))\n",
    "        * ((2 * k + 1).view(1, -1))\n",
    "    )\n",
    "    mdct_matrix = torch.cos(arg) * (2.0 / N) ** 0.5\n",
    "\n",
    "    return torch.matmul(x, mdct_matrix)\n",
    "\n",
    "\n",
    "def imdct(X):\n",
    "    half_N = X.shape[-1]\n",
    "    N = half_N * 2\n",
    "    n = torch.arange(N, device=X.device)\n",
    "    k = torch.arange(half_N, device=X.device)\n",
    "\n",
    "    arg = (\n",
    "        (math.pi / (2 * N))\n",
    "        * ((2 * n + 1 + N // 2).view(-1, 1))\n",
    "        * ((2 * k + 1).view(1, -1))\n",
    "    )\n",
    "    imdct_matrix = torch.cos(arg) * (2.0 / N) ** 0.5\n",
    "\n",
    "    return torch.matmul(X, imdct_matrix.T) * 2.0\n",
    "\n",
    "\n",
    "def audio_to_mdct_frames(audio, frame_size=1920, midside=False):\n",
    "    batch, channels, samples = audio.shape\n",
    "    hop_size = frame_size // 2\n",
    "\n",
    "    if midside:\n",
    "        mid = (audio[:, 0, :] + audio[:, 1, :]) / 2.0\n",
    "        side = (audio[:, 0, :] - audio[:, 1, :]) / 2.0\n",
    "        audio = torch.stack([mid, side], dim=1)\n",
    "\n",
    "    n_frames = (samples - frame_size) // hop_size + 1\n",
    "\n",
    "    frames = []\n",
    "    for i in range(n_frames):\n",
    "        start = i * hop_size\n",
    "        frame = audio[:, :, start : start + frame_size]\n",
    "        if frame.shape[-1] == frame_size:\n",
    "            frames.append(frame)\n",
    "\n",
    "    frames = torch.stack(frames, dim=2)\n",
    "\n",
    "    window = torch.sin(\n",
    "        torch.pi / frame_size * (torch.arange(frame_size, device=audio.device) + 0.5)\n",
    "    )\n",
    "    frames = frames * window.view(1, 1, 1, -1)\n",
    "\n",
    "    shape = frames.shape\n",
    "    frames_reshaped = frames.reshape(-1, frame_size)\n",
    "    mdct_coeffs = mdct(frames_reshaped)\n",
    "\n",
    "    return mdct_coeffs.reshape(shape[0], shape[1], shape[2], -1)\n",
    "\n",
    "\n",
    "def mdct_frames_to_audio(mdct_coeffs, frame_size=1920, midside=False):\n",
    "    batch, channels, n_frames, half_frame_size = mdct_coeffs.shape\n",
    "    hop_size = frame_size // 2\n",
    "\n",
    "    shape = mdct_coeffs.shape\n",
    "    coeffs_reshaped = mdct_coeffs.reshape(-1, half_frame_size)\n",
    "    frames = imdct(coeffs_reshaped)\n",
    "    frames = frames.reshape(shape[0], shape[1], shape[2], -1)\n",
    "\n",
    "    window = torch.sin(\n",
    "        torch.pi\n",
    "        / frame_size\n",
    "        * (torch.arange(frame_size, device=mdct_coeffs.device) + 0.5)\n",
    "    )\n",
    "    frames = frames * window.view(1, 1, 1, -1)\n",
    "\n",
    "    total_samples = (n_frames - 1) * hop_size + frame_size\n",
    "    output = torch.zeros(batch, channels, total_samples, device=mdct_coeffs.device)\n",
    "\n",
    "    for i in range(n_frames):\n",
    "        start = i * hop_size\n",
    "        output[:, :, start : start + frame_size] += frames[:, :, i]\n",
    "\n",
    "    if midside:\n",
    "        mid = output[:, 0, :]\n",
    "        side = output[:, 1, :]\n",
    "        left = mid + side\n",
    "        right = mid - side\n",
    "        output = torch.stack([left, right], dim=1)\n",
    "\n",
    "    return output\n",
    "\n",
    "\n",
    "class WaveformMDCTVAE(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        frame_size: int = 1920,\n",
    "        latent_dim: int = 256,\n",
    "        hidden_dims: list = None,\n",
    "        dropout: float = 0.1,\n",
    "        midside: bool = False,\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        self.frame_size = frame_size\n",
    "        self.n_coeffs = frame_size // 2\n",
    "        self.latent_dim = latent_dim\n",
    "        self.midside = midside\n",
    "\n",
    "        if hidden_dims is None:\n",
    "            hidden_dims = [512, 256]\n",
    "\n",
    "        # Encoder layers\n",
    "        modules = []\n",
    "        input_dim = 2 * self.n_coeffs\n",
    "\n",
    "        for h_dim in hidden_dims:\n",
    "            modules.append(\n",
    "                nn.Sequential(\n",
    "                    nn.Linear(input_dim, h_dim),\n",
    "                    nn.LayerNorm(h_dim),\n",
    "                    nn.LeakyReLU(),\n",
    "                    nn.Dropout(dropout),\n",
    "                )\n",
    "            )\n",
    "            input_dim = h_dim\n",
    "\n",
    "        self.encoder = nn.Sequential(*modules)\n",
    "        self.fc_mu = nn.Linear(hidden_dims[-1], latent_dim)\n",
    "        self.fc_var = nn.Linear(hidden_dims[-1], latent_dim)\n",
    "\n",
    "        # Decoder layers\n",
    "        modules = []\n",
    "        hidden_dims.reverse()\n",
    "\n",
    "        self.decoder_input = nn.Sequential(\n",
    "            nn.Linear(latent_dim, hidden_dims[0]),\n",
    "            nn.LayerNorm(hidden_dims[0]),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Dropout(dropout),\n",
    "        )\n",
    "\n",
    "        for i in range(len(hidden_dims) - 1):\n",
    "            modules.append(\n",
    "                nn.Sequential(\n",
    "                    nn.Linear(hidden_dims[i], hidden_dims[i + 1]),\n",
    "                    nn.LayerNorm(hidden_dims[i + 1]),\n",
    "                    nn.LeakyReLU(),\n",
    "                    nn.Dropout(dropout),\n",
    "                )\n",
    "            )\n",
    "\n",
    "        self.decoder = nn.Sequential(*modules)\n",
    "        self.final_layer = nn.Linear(hidden_dims[-1], 2 * self.n_coeffs)\n",
    "\n",
    "    def _encode(self, mdct_frames: torch.Tensor) -> list[torch.Tensor]:\n",
    "        batch_size, _, n_frames, _ = mdct_frames.shape\n",
    "        ch1 = mdct_frames[:, 0, :, :]\n",
    "        ch2 = mdct_frames[:, 1, :, :]\n",
    "\n",
    "        x = torch.cat((ch1, ch2), dim=-1)\n",
    "        result = self.encoder(x)\n",
    "        mu = self.fc_mu(result)\n",
    "        log_var = self.fc_var(result)\n",
    "\n",
    "        return [mu, log_var]\n",
    "\n",
    "    def _decode(self, z: torch.Tensor) -> torch.Tensor:\n",
    "        batch_size, n_frames, _ = z.shape\n",
    "\n",
    "        result = self.decoder_input(z)\n",
    "        result = self.decoder(result)\n",
    "        result = self.final_layer(result)\n",
    "        ch1 = result[..., :self.n_coeffs]\n",
    "        ch2 = result[..., self.n_coeffs:]\n",
    "        result = torch.stack((ch1, ch2), dim=1)\n",
    "\n",
    "        return result\n",
    "\n",
    "    def reparameterize(self, mu: torch.Tensor, log_var: torch.Tensor) -> torch.Tensor:\n",
    "        if self.training:\n",
    "            std = torch.exp(0.5 * log_var)\n",
    "            eps = torch.randn_like(std)\n",
    "            return eps * std + mu\n",
    "        else:\n",
    "            return mu\n",
    "\n",
    "    def forward(\n",
    "        self, waveform: torch.Tensor\n",
    "    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:\n",
    "        \"\"\"\n",
    "        Forward pass handling both waveform conversion and VAE operations.\n",
    "\n",
    "        Args:\n",
    "            waveform: Input audio tensor of shape (batch_size, channels, samples)\n",
    "\n",
    "        Returns:\n",
    "            tuple containing:\n",
    "            - reconstructed waveform\n",
    "            - reconstructed MDCT coefficients\n",
    "            - original MDCT coefficients (for loss computation)\n",
    "            - mu\n",
    "            - log_var\n",
    "        \"\"\"\n",
    "        # Convert input waveform to MDCT frames\n",
    "        mdct_frames = audio_to_mdct_frames(\n",
    "            waveform, frame_size=self.frame_size, midside=self.midside\n",
    "        )  # (batch, channels, n_frames, n_coeffs)\n",
    "\n",
    "        # Encode and decode\n",
    "        mu, log_var = self._encode(mdct_frames)\n",
    "        z = self.reparameterize(mu, log_var)\n",
    "        mdct_recon = self._decode(z)\n",
    "\n",
    "        # Convert back to waveform\n",
    "        waveform_recon = mdct_frames_to_audio(\n",
    "            mdct_recon, frame_size=self.frame_size, midside=self.midside\n",
    "        )\n",
    "\n",
    "        return waveform_recon, mdct_recon, mdct_frames, mu, log_var\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "import torch\n",
    "import torchaudio\n",
    "\n",
    "model = WaveformMDCTVAE(frame_size=1920, latent_dim=128, hidden_dims=[512, 256], dropout=0.1, midside=True)\n",
    "\n",
    "# test on new audio \n",
    "test_filepath = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n",
    "waveform, sr = torchaudio.load(test_filepath)\n",
    "\n",
    "if sr != 48000:\n",
    "    waveform = torchaudio.transforms.Resample(sr, 48000)(waveform)\n",
    "\n",
    "x = waveform[:, :48000*30].unsqueeze(0)\n",
    "\n",
    "with torch.no_grad():\n",
    "    mdct_frames = audio_to_mdct_frames(x, frame_size=1920, midside=True)\n",
    "    mu, log_var = model._encode(mdct_frames)\n",
    "    z = model.reparameterize(mu, log_var)\n",
    "    mdct_recon = model._decode(z)\n",
    "    print(mdct_recon.shape)\n",
    "\n",
    "    waveform_recon = mdct_frames_to_audio(mdct_recon, frame_size=1920, midside=True)\n",
    "\n",
    "    IPython.display.display(IPython.display.Audio(waveform_recon.squeeze(0).numpy(), rate=48000))\n",
    "    IPython.display.display(IPython.display.Audio(x.squeeze(0).numpy(), rate=48000))\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test mdct reconstruction for multi-item batch\n",
    "\n",
    "import IPython\n",
    "import torch\n",
    "import torchaudio\n",
    "\n",
    "model = WaveformMDCTVAE(frame_size=1920, latent_dim=128, hidden_dims=[512, 256], dropout=0.1, midside=True)\n",
    "\n",
    "# test on new audio \n",
    "test_filepath = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n",
    "test_filepath2 = \"/home/christian/audio/reference-audio-wav/02 Take Five.wav\"\n",
    "waveform, sr = torchaudio.load(test_filepath)\n",
    "waveform2, sr2 = torchaudio.load(test_filepath2)\n",
    "\n",
    "if sr != 48000:\n",
    "    waveform = torchaudio.transforms.Resample(sr, 48000)(waveform)\n",
    "\n",
    "if sr2 != 48000:\n",
    "    waveform2 = torchaudio.transforms.Resample(sr2, 48000)(waveform2)\n",
    "\n",
    "x = waveform[:, :48000*30].unsqueeze(0)\n",
    "x2 = waveform2[:, :48000*30].unsqueeze(0)\n",
    "\n",
    "x = torch.cat((x, x2), dim=0)\n",
    "print(x.shape)\n",
    "\n",
    "with torch.no_grad():\n",
    "    mdct_frames = audio_to_mdct_frames(x, frame_size=1920, midside=False)\n",
    "    print(mdct_frames.shape)\n",
    "    waveform_recon = mdct_frames_to_audio(mdct_frames, frame_size=1920, midside=False)\n",
    "    print(waveform_recon.shape)\n",
    "\n",
    "    IPython.display.display(IPython.display.Audio(waveform_recon[0].numpy(), rate=48000))\n",
    "    IPython.display.display(IPython.display.Audio(x[0].numpy(), rate=48000))\n",
    "\n",
    "    IPython.display.display(IPython.display.Audio(waveform_recon[1].numpy(), rate=48000))\n",
    "    IPython.display.display(IPython.display.Audio(x[1].numpy(), rate=48000))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# look at the distribution of the mdct frames\n",
    "from typing import List\n",
    "import numpy as np\n",
    "import itertools\n",
    "from tqdm import tqdm\n",
    "\n",
    "class BufferedAudioDataset(torch.utils.data.Dataset):\n",
    "    def __init__(\n",
    "        self,\n",
    "        filepaths: List[str],\n",
    "        sample_rate: int,\n",
    "        num_workers: int = 1,\n",
    "        chunk_size_s: float = 10.0,\n",
    "        buffer_size: int = 50_000,\n",
    "    ):\n",
    "        self.filepaths = filepaths\n",
    "        self.sample_rate = sample_rate\n",
    "        self.chunk_size_s = chunk_size_s\n",
    "        self.buffer_size = buffer_size\n",
    "        self.chunk_size_samples = int(chunk_size_s * sample_rate)\n",
    "        self.num_workers = num_workers\n",
    "        self.items_since_last_reload = buffer_size  # force a reload\n",
    "        self.buffer = []\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.buffer_size * self.num_workers\n",
    "\n",
    "    def _reload_buffer(self):\n",
    "        self.buffer = []\n",
    "        rand_idxs = torch.randperm(len(self.filepaths))\n",
    "\n",
    "        print(\"Reloading buffer...\")\n",
    "        # max rand_idxs repeat endlessly\n",
    "        rand_idxs = itertools.cycle(rand_idxs)\n",
    "        pbar = tqdm(rand_idxs, total=len(self.filepaths), desc=\"Loading audio buffer\")\n",
    "        for idx in pbar:\n",
    "            if len(self.buffer) >= self.buffer_size:\n",
    "                break\n",
    "\n",
    "            try:\n",
    "                filepath = self.filepaths[idx]\n",
    "                audio, sr = torchaudio.load(filepath)\n",
    "\n",
    "                if sr != self.sample_rate:\n",
    "                    audio = torchaudio.functional.resample(audio, sr, self.sample_rate)\n",
    "\n",
    "                # Pad if needed to ensure consistent chunk size\n",
    "                if audio.shape[-1] < self.chunk_size_samples:\n",
    "                    continue\n",
    "\n",
    "                # Split into chunks\n",
    "                chunks = audio.unfold(\n",
    "                    -1, self.chunk_size_samples, self.chunk_size_samples\n",
    "                )\n",
    "                chunks = chunks.chunk(chunks.shape[1], dim=1)\n",
    "\n",
    "                # Filter chunks by minimum length\n",
    "                valid_chunks = [\n",
    "                    chunk.squeeze(1)\n",
    "                    for chunk in chunks\n",
    "                    if chunk.shape[-1] >= self.chunk_size_samples\n",
    "                ]\n",
    "\n",
    "                # filter out chunks of silence\n",
    "                valid_chunks = [\n",
    "                    chunk for chunk in valid_chunks if (chunk.abs() ** 2).mean() > 0.001\n",
    "                ]\n",
    "\n",
    "                self.buffer.extend(valid_chunks)\n",
    "\n",
    "                pbar.set_postfix({\"buffer_size\": len(self.buffer)})\n",
    "\n",
    "            except Exception as e:\n",
    "                print(f\"Error loading {filepath}: {e}\")\n",
    "                continue\n",
    "        self.items_since_last_reload = 0\n",
    "\n",
    "    def __getitem__(self, _):\n",
    "        if self.items_since_last_reload >= len(self.buffer):\n",
    "            self._reload_buffer()\n",
    "\n",
    "        # get a random preset and apply it to the audio\n",
    "        buffer_idx = np.random.randint(0, len(self.buffer))\n",
    "        audio = self.buffer[buffer_idx]\n",
    "\n",
    "        # ensure nothing is out of range\n",
    "        if audio.abs().max() > 1.0:\n",
    "            audio = audio / audio.abs().max()\n",
    "\n",
    "        # apply random gain reduction\n",
    "        if np.random.uniform() < 0.5:\n",
    "            gain_reduction_db = np.random.uniform(-10, 0)\n",
    "            audio *= 10 ** (gain_reduction_db / 20.0)\n",
    "\n",
    "        # self.items_since_last_reload += 1\n",
    "\n",
    "        return audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "train_filepaths = glob.glob(\n",
    "    os.path.join(\"/app/suno/data/audio_2ch_48khz_lg/val/genius_hq\", \"*.wav\"), recursive=True\n",
    ")\n",
    "\n",
    "dataset = BufferedAudioDataset(train_filepaths, 48000, num_workers=4, chunk_size_s=10.0, buffer_size=1_000)\n",
    "dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=4)\n",
    "\n",
    "examples = []\n",
    "for batch in dataloader:\n",
    "    # measure the mdct frames\n",
    "    mdct_frames = audio_to_mdct_frames(batch, frame_size=1920, midside=False)\n",
    "    ch1 = mdct_frames[:, 0, :, :]\n",
    "    ch2 = mdct_frames[:, 1, :, :]\n",
    "    x = torch.cat((ch1, ch2), dim=-1)\n",
    "    examples.append(x)\n",
    "\n",
    "examples = torch.cat(examples, dim=0)\n",
    "print(examples.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(examples))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "examples = torch.cat(examples[0:250], dim=0)\n",
    "print(examples.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "flat_examples = examples.flatten()\n",
    "flat_examples /=\n",
    "\n",
    "print(min(flat_examples.numpy()))\n",
    "print(max(flat_examples.numpy()))\n",
    "print(np.mean(flat_examples.numpy()))\n",
    "print(np.std(flat_examples.numpy()))\n",
    "\n",
    "plt.hist(flat_examples.numpy(), bins=100)\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Scale factor\n",
    "k = 100  # Example scaling factor, adjust based on desired std deviation in output\n",
    "\n",
    "# Apply tanh transformation\n",
    "transformed_data = np.tanh(k * flat_examples.numpy())\n",
    "\n",
    "# Compute new standard deviation of transformed data\n",
    "new_std = np.std(transformed_data)\n",
    "print(\"New Standard Deviation of Transformed Data:\", new_std)\n",
    "\n",
    "# Inverse transformation using artanh\n",
    "recovered_data = np.arctanh(transformed_data) \n",
    "recovered_data /= k\n",
    "\n",
    "# Compute standard deviation of recovered data\n",
    "recovered_std = np.std(recovered_data)\n",
    "print(\"Standard Deviation of Recovered Data:\", recovered_std)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}