{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import funcy\n",
    "import IPython\n",
    "import numpy as np\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"6\"\n",
    "\n",
    "from suno_utils.utils.text import (    \n",
    "    write_jsonl,\n",
    "    read_jsonl,\n",
    "    write_json,\n",
    "    read_json,\n",
    "    normalize_whitespace,\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "CODEC_FILEPATH = \"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\"\n",
    "\n",
    "from suno_utils.tasks.dac_vae_fixed_25hz import (\n",
    "    preload_models as v2_preload_codec_models,\n",
    "    decode as v2_codec_decode,\n",
    "    encode as v2_codec_encode,\n",
    "    decode_stream_to_full_audio,\n",
    ")\n",
    "\n",
    "_ = v2_preload_codec_models(CODEC_FILEPATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "CODEC_FILEPATH = \"s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth\"\n",
    "\n",
    "from suno_utils.tasks.dac_vae_100hz_peaq import (  # NOTE: works for 25hz as well\n",
    "    preload_models as v1_preload_codec_models,\n",
    "    decode as v1_codec_decode,\n",
    "    encode as v1_codec_encode,\n",
    "    get_embedding_rate,\n",
    "    load_model as load_codec_model,\n",
    ")\n",
    "\n",
    "_ = v1_preload_codec_models(CODEC_FILEPATH)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "def moving_average(latents, kernel_size=5):\n",
    "    \"\"\"\n",
    "    latents: Tensor of shape (C, N)\n",
    "    kernel_size: Size of the smoothing window\n",
    "    \"\"\"\n",
    "    latents = torch.from_numpy(latents.T)\n",
    "    print(latents.shape)\n",
    "    latents = latents.unsqueeze(0)  # Add batch dim: (1, C, N)\n",
    "    latents = F.pad(latents, (kernel_size//2, kernel_size//2), mode='reflect')\n",
    "    kernel = torch.ones(1, 1, kernel_size, device=latents.device) / kernel_size\n",
    "    smoothed = F.conv1d(latents, kernel.expand(latents.size(1), -1, -1), groups=latents.size(1))\n",
    "    return smoothed.squeeze(0).T  # Shape: (C, N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "metadata": {},
   "outputs": [],
   "source": [
    "def diffuse(x, dist, noise_schedule, rng, batch_size, device, diffusion_objective, t_discretize=1):\n",
    "    if dist == \"training\":\n",
    "        if noise_schedule == \"default\" or noise_schedule == \"early\":\n",
    "            # Draw uniformly distributed continuous timesteps\n",
    "            t = rng.draw(batch_size)[:, 0].to(device).to(torch.bfloat16)\n",
    "        elif noise_schedule == \"logit_normal\":\n",
    "            # Draw from a logit-normal distribution\n",
    "            t = torch.sigmoid(torch.randn(batch_size, device=device).to(torch.bfloat16))\n",
    "        elif noise_schedule == \"trunc_logit_normal\":\n",
    "            t = truncated_logistic_normal_rescaled(batch_size).to(device).to(torch.bfloat16)\n",
    "            # Flip the distribution\n",
    "            t = 1 - t\n",
    "        elif noise_schedule == \"log_snr\":\n",
    "            t = sample_timesteps_logsnr(batch_size).to(device).to(torch.bfloat16)\n",
    "        else:\n",
    "            raise ValueError(f\"Invalid noise schedule: {noise_schedule}\")\n",
    "    elif dist == \"inference\":\n",
    "        # For inference, use linearly spaced timesteps\n",
    "        if False:\n",
    "            if t_discretize == 1:\n",
    "                t = torch.ones(batch_size, device=device).to(torch.bfloat16)\n",
    "            else:\n",
    "                t = torch.linspace(0, 1, t_discretize + 1, dtype=torch.bfloat16, device=device)\n",
    "                # Select random timesteps from the discretized schedule\n",
    "                indices = torch.randint(0, t_discretize + 1, (batch_size,), device=device)\n",
    "                t = t[indices]\n",
    "        else:\n",
    "            t = rng.draw(batch_size)[:, 0].to(device).to(torch.bfloat16)\n",
    "    else:\n",
    "        raise ValueError(f\"Invalid diffusion distribution: {dist}\")\n",
    "\n",
    "    # Calculate the noise schedule parameters for those timesteps\n",
    "    if diffusion_objective in [\"v\"]:\n",
    "        alphas, sigmas = get_alphas_sigmas(t)\n",
    "    elif diffusion_objective in [\"rectified_flow\", \"rf_denoiser\"]:\n",
    "        alphas, sigmas = 1 - t, t\n",
    "    else:\n",
    "        raise ValueError(f\"Invalid diffusion objective: {diffusion_objective}\")\n",
    "\n",
    "    x = x.to(device)\n",
    "    x = x.to(t.dtype)\n",
    "    # Combine the ground truth data and the noise\n",
    "    alphas = alphas[:, None, None]\n",
    "    sigmas = sigmas[:, None, None]\n",
    "    eps = torch.randn_like(x)\n",
    "    noised_inputs = x * alphas + eps * sigmas\n",
    "\n",
    "    return noised_inputs, t, sigmas\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_timesteps_logsnr(batch_size, mean_logsnr=-1.2, std_logsnr=2.0):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        batch_size (int): Number of timesteps to sample\n",
    "        mean_logsnr (float): Mean of the logSNR Gaussian distribution\n",
    "        std_logsnr (float): Standard deviation of the logSNR Gaussian distribution\n",
    "\n",
    "    Returns:\n",
    "        torch.Tensor: Tensor of shape (batch_size,) containing timestep values t in [0, 1]\n",
    "    \"\"\"\n",
    "    # Sample logSNR from Gaussian distribution\n",
    "    logsnr = torch.randn(batch_size) * std_logsnr + mean_logsnr\n",
    "\n",
    "    # Convert logSNR to timesteps using the logistic function\n",
    "    # Since logSNR = ln((1-t)/t), we can solve for t:\n",
    "    # t = 1 / (1 + exp(logsnr))\n",
    "    t = torch.sigmoid(-logsnr)\n",
    "\n",
    "    # Clamp values to ensure numerical stability\n",
    "    t = t.clamp(1e-4, 1 - 1e-4)\n",
    "\n",
    "    return t\n",
    "\n",
    "v2_scale_factor = 0.4\n",
    "bs1 = 1\n",
    "device = \"cpu\"\n",
    "\n",
    "#t = sample_timesteps_logsnr(bs1).to(device)\n",
    "t = torch.ones(bs1, device=device) * 0.1\n",
    "\n",
    "alphas, sigmas = 1 - t, t\n",
    "print(alphas, sigmas)\n",
    "\n",
    "x = torch.from_numpy(z_v2[0]) * v2_scale_factor\n",
    "\n",
    "alphas = alphas[:, None, None]\n",
    "sigmas = sigmas[:, None, None]\n",
    "eps = torch.randn_like(x)\n",
    "noised_inputs = x * alphas #+ eps * sigmas\n",
    "\n",
    "print(noised_inputs.shape)\n",
    "\n",
    "\n",
    "# original decoded\n",
    "#v2_decoded = v2_codec_decode(z_v2[0])\n",
    "v2_decoded_smoothed = v2_codec_decode(noised_inputs.squeeze(0).numpy() / v2_scale_factor)\n",
    "\n",
    "# save the decoded audios\n",
    "#v2_decoded.play()\n",
    "v2_decoded_smoothed.play()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import torchaudio\n",
    "import pyloudnorm as pyln\n",
    "from tqdm import tqdm\n",
    "audio_paths = glob.glob(\"/home/christian/audio/genius_ear_100/*.webm\")\n",
    "\n",
    "loudness_normalize = True\n",
    "target_loudness = -16.0\n",
    "\n",
    "audios = []\n",
    "for audio_path in tqdm(audio_paths[5:6]):\n",
    "    audio, sr = torchaudio.load(audio_path)\n",
    "    audio = audio[:,:30*sr]\n",
    "    # loudness normalize the audio\n",
    "    if loudness_normalize:\n",
    "        meter = pyln.Meter(sr) # create BS.1770 meter\n",
    "        loudness = meter.integrated_loudness(audio.permute(1, 0).numpy())\n",
    "        gain_db = target_loudness - loudness\n",
    "        audio *= 10 ** (gain_db / 20.0)\n",
    "        \n",
    "    audios.append(audio)\n",
    "\n",
    "\n",
    "z_v1 = v1_codec_encode(audios)\n",
    "z_v2 = v2_codec_encode(audios)\n",
    "\n",
    "print(z_v1[0].shape, z_v2[0].shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compress(x: torch.Tensor, alpha: float = 1.0) -> torch.Tensor:\n",
    "    return torch.sign(x) * torch.log1p(alpha * torch.abs(x)) / alpha\n",
    "\n",
    "def decompress(y: torch.Tensor, alpha: float = 1.0) -> torch.Tensor:\n",
    "    return torch.sign(y) * (torch.expm1(alpha * torch.abs(y)) / alpha)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# apply scale factors to the z vectors\n",
    "\n",
    "v2_scale_factor = 0.4\n",
    "v1_scale_factor = 1.0\n",
    "\n",
    "#z_v1_scaled = z_v1[0] * v1_scale_factor\n",
    "\n",
    "z_v1_scaled = z_v2[0] * v2_scale_factor\n",
    "z_v2_scaled = compress(torch.from_numpy(z_v2[0]), alpha=1.0) \n",
    "z_v2_descaled = decompress(z_v2_scaled, alpha=1.0)\n",
    "\n",
    "z_v1_scaled_std = z_v1_scaled.std()\n",
    "z_v2_scaled_std = z_v2_scaled.std()\n",
    "z_v2_descaled_std = z_v2_descaled.std()\n",
    "\n",
    "print(z_v1_scaled_std, z_v2_scaled_std, z_v2_descaled_std)\n",
    "\n",
    "\n",
    "# 2d plot with pcolormesh for v1 and v2 on the same plot\n",
    "\n",
    "# lets do a 10 \n",
    "fig, axs = plt.subplots(1, 1, figsize=(8, 4))\n",
    "\n",
    "# Find the common scale for both subplots\n",
    "#vmin = min(z_v12_derivative_scaled.min(), z_v2_scaled.min())\n",
    "#vmax = max(z_v12_derivative_scaled.max(), z_v2_scaled.max())\n",
    "\n",
    "#axs[0].pcolormesh(z_v12_derivative_scaled.T, vmin=vmin, vmax=vmax)\n",
    "#axs[1].pcolormesh(z_v2_scaled.T, vmin=vmin, vmax=vmax)\n",
    "#plt.show()\n",
    "\n",
    "# lets make line plot of the first dim of the two vectors\n",
    "axs.plot(z_v1_scaled[:, 54], label=\"v1\", alpha=0.5)\n",
    "axs.plot(z_v2_scaled[:, 54], label=\"v2\", alpha=0.5)\n",
    "axs.plot(z_v2_descaled[:, 54], label=\"v2_descaled\", alpha=0.5)\n",
    "axs.legend()\n",
    "plt.show()\n",
    "\n",
    "# lets make line plot of the first dim of the two vectors\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, axs = plt.subplots(10, 10, figsize=(20, 20), sharex=True, sharey=True)\n",
    "fig.suptitle(\"Comparison of z_v1_scaled and z_v2_scaled across all dimensions\", fontsize=16)\n",
    "\n",
    "for i in tqdm(range(10)):\n",
    "    for j in range(10):\n",
    "        dim = i * 10 + j\n",
    "        ax = axs[i, j]\n",
    "        #ax.plot(z_v1_scaled[:, dim], label=\"v1\", alpha=0.5)\n",
    "        #ax.plot(z_v2_scaled[:, dim], label=\"v2\", alpha=0.5)\n",
    "        ax.scatter(range(len(z_v1_scaled[:, dim])), z_v1_scaled[:, dim], label=\"v1\", alpha=0.5, s=1)\n",
    "        ax.scatter(range(len(z_v2_scaled[:, dim])), z_v2_scaled[:, dim], label=\"v2\", alpha=0.5, s=1)\n",
    "\n",
    "        # Add dim number as text inside the plot (top-left corner)\n",
    "        ax.text(0.01, 0.95, f\"{dim}\", transform=ax.transAxes,\n",
    "                fontsize=14, verticalalignment='top')\n",
    "\n",
    "        ax.tick_params(axis='both', which='both', labelsize=6)\n",
    "\n",
    "# Add a legend only to the first subplot\n",
    "#axs[0, 0].legend(fontsize=6)\n",
    "\n",
    "plt.tight_layout(rect=[0, 0.03, 1, 0.95])\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# lets listen to the original decoded and the smoothed decoded\n",
    "\n",
    "# original decoded\n",
    "v2_decoded = v2_codec_decode(z_v2_scaled)\n",
    "v2_decoded_smoothed = v2_codec_decode(z_v12_smoothed_scaled)\n",
    "\n",
    "# save the decoded audios\n",
    "v2_decoded.play()\n",
    "v2_decoded_smoothed.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#audio_path = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n",
    "\n",
    "import glob\n",
    "import torchaudio\n",
    "import pyloudnorm as pyln\n",
    "from tqdm import tqdm\n",
    "audio_paths = glob.glob(\"/home/christian/audio/reference-audio-wav/*.wav\")\n",
    "audio_paths = glob.glob(\"/home/christian/audio/genius_ear_100/*.webm\")\n",
    "\n",
    "loudness_normalize = True\n",
    "target_loudness = -16.0\n",
    "\n",
    "audios = []\n",
    "for audio_path in tqdm(audio_paths):\n",
    "    audio, sr = torchaudio.load(audio_path)\n",
    "    audio = audio[:,:30*sr]\n",
    "    # loudness normalize the audio\n",
    "    if loudness_normalize:\n",
    "        meter = pyln.Meter(sr) # create BS.1770 meter\n",
    "        loudness = meter.integrated_loudness(audio.permute(1, 0).numpy())\n",
    "        gain_db = target_loudness - loudness\n",
    "        audio *= 10 ** (gain_db / 20.0)\n",
    "        \n",
    "    audios.append(audio)\n",
    "\n",
    "\n",
    "z_v1 = v1_codec_encode(audios)\n",
    "z_v2 = v2_codec_encode(audios)\n",
    "\n",
    "print(z_v1[0].shape, z_v2[0].shape)\n",
    "\n",
    "# append the z vectors to an array\n",
    "z_v1_all = np.concatenate(z_v1, axis=0)\n",
    "z_v2_all = np.concatenate(z_v2, axis=0)\n",
    "\n",
    "print(z_v1_all.shape, z_v2_all.shape)\n",
    "\n",
    "# apply scale factors to the z vectors\n",
    "\n",
    "v2_scale_factor = 0.4\n",
    "v1_scale_factor = 1.0\n",
    "\n",
    "z_v1 = z_v1_all * v1_scale_factor\n",
    "z_v2 = z_v2_all * v2_scale_factor\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute statistics for the z vectors, which are numpy arrays\n",
    "\n",
    "z_v1_stats = {\n",
    "    \"mean\": z_v1.mean(axis=0),\n",
    "    \"std\": z_v1.std(axis=0),\n",
    "    \"max\": z_v1.max(axis=0),\n",
    "    \"min\": z_v1.min(axis=0),\n",
    "}\n",
    "\n",
    "z_v2_stats = {\n",
    "    \"mean\": z_v2.mean(axis=0),\n",
    "    \"std\": z_v2.std(axis=0),\n",
    "    \"max\": z_v2.max(axis=0),\n",
    "    \"min\": z_v2.min(axis=0),\n",
    "}\n",
    "\n",
    "print(z_v1.std(), z_v2.std())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "metadata": {},
   "outputs": [],
   "source": [
    "loudness_normalize = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# Assuming z_v1_stats and z_v2_stats are already defined\n",
    "\n",
    "# Extract stat names\n",
    "stat_names = ['mean', 'std', 'min', 'max']\n",
    "num_features = len(z_v1_stats['mean'])\n",
    "x = np.arange(num_features)\n",
    "\n",
    "# Create a subplot for each stat\n",
    "fig, axes = plt.subplots(len(stat_names), 1, figsize=(12, 10), sharex=True)\n",
    "\n",
    "for i, stat in enumerate(stat_names):\n",
    "    axes[i].plot(x, z_v1_stats[stat], label='v1 vae', alpha=0.8)\n",
    "    axes[i].plot(x, z_v2_stats[stat], label='v2 vae', alpha=0.8)\n",
    "    axes[i].set_ylabel(stat)\n",
    "    axes[i].legend()\n",
    "    axes[i].grid(True)\n",
    "\n",
    "axes[-1].set_xlabel('Feature Index')\n",
    "plt.suptitle('Comparison of v1 vae vs v2 vae Statistics per Feature Dimension (loudness normalized = {})'.format(loudness_normalize), fontsize=14)\n",
    "plt.tight_layout(rect=[0, 0.03, 1, 0.95])\n",
    "#plt.show()\n",
    "plt.savefig(\"plots/vae_stats.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_diff",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}