{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import auraloss\n",
    "import torch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "melstft = auraloss.freq.MelSTFTLoss(\n",
    "    n_mels=128,\n",
    "    fft_size=2048,\n",
    "    win_length=2048,\n",
    "    hop_size=1024,\n",
    "    sample_rate=24000,\n",
    "    reduction=\"none\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
      "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
      "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.2000, 0.6000, 0.2000,\n",
      "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
      "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
      "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.2000, 0.6000, 0.2000,\n",
      "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
      "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
      "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.2000, 0.6000, 0.2000,\n",
      "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
      "        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
      "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n",
      "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.2000, 0.6000, 0.2000,\n",
      "         0.0000, 0.0000, 0.0000, 0.0000, 0.0000]])\n"
     ]
    }
   ],
   "source": [
    "\n",
    "min_quant_error = 0.0\n",
    "max_quant_error = 8.0\n",
    "num_quant_levels = 32\n",
    "\n",
    "audio_in_a = torch.randn(4, 2, 131072)\n",
    "audio_out_a = audio_in_a.clone() * 0.99\n",
    "\n",
    "audio_in_b = torch.randn(4, 2, 131072)\n",
    "audio_out_b = audio_in_b.clone() * 0.9\n",
    "\n",
    "# compute metrics for audio pairs using the unprocessed inputs as targets\n",
    "melstft_error_a = melstft(\n",
    "    audio_in_a.mean(dim=1, keepdim=True),\n",
    "    audio_out_a.mean(dim=1, keepdim=True),\n",
    ").mean(dim=(1, 2))\n",
    "melstft_error_b = melstft(\n",
    "    audio_in_b.mean(dim=1, keepdim=True),\n",
    "    audio_out_b.mean(dim=1, keepdim=True),\n",
    ").mean(dim=(1, 2))\n",
    "\n",
    "# create labels for computing loss\n",
    "pref_label = (melstft_error_a > melstft_error_b).float()\n",
    "# when True, this means that \"ref\" (second input) is higher quality than first input\n",
    "\n",
    "# create quantification lablels as bins with special smoothing\n",
    "quant_label = torch.abs(melstft_error_a - melstft_error_b).clamp(\n",
    "    min_quant_error,\n",
    "    max_quant_error,\n",
    ")\n",
    "# Compute bin edges and digitize the vector into bins\n",
    "bin_edges = torch.logspace(\n",
    "    torch.log10(torch.tensor(min_quant_error + 1e-6)),  # Add small epsilon to avoid log(0)\n",
    "    torch.log10(torch.tensor(max_quant_error)),\n",
    "    num_quant_levels + 1,\n",
    ").type_as(quant_label)\n",
    "bin_indices = torch.bucketize(quant_label, bin_edges)\n",
    "\n",
    "\n",
    "quant_label = torch.nn.functional.one_hot(\n",
    "    bin_indices, num_classes=num_quant_levels\n",
    ").float()\n",
    "\n",
    "smoothing_value = 0.6\n",
    "smoothing_neighbor_value = 0.2\n",
    "\n",
    "# Create smoothed vectors\n",
    "smoothed_quant_label = quant_label * smoothing_value\n",
    "\n",
    "# Add neighbor values\n",
    "for i, index in enumerate(bin_indices):\n",
    "    if index > 0:\n",
    "        smoothed_quant_label[i, index - 1] += smoothing_neighbor_value\n",
    "    if index < num_quant_levels - 1:\n",
    "        smoothed_quant_label[i, index + 1] += smoothing_neighbor_value\n",
    "\n",
    "print(smoothed_quant_label)"
   ]
  },
  {
   "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
}