{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def append_zero(x):\n",
    "    return torch.cat([x, x.new_zeros([1])])\n",
    "\n",
    "\n",
    "def get_sigmas_polyexponential(n, sigma_min, sigma_max, rho=1.0, device=\"cpu\"):\n",
    "    \"\"\"Constructs an polynomial in log sigma noise schedule.\"\"\"\n",
    "    ramp = torch.linspace(1, 0, n, device=device) ** rho\n",
    "    sigmas = torch.exp(ramp * (math.log(sigma_max) - math.log(sigma_min)) + math.log(sigma_min))\n",
    "    return append_zero(sigmas)\n",
    "\n",
    "def get_sigmas_cosine(n, sigma_min, sigma_max, device=\"cpu\"):\n",
    "    \"\"\"Constructs a cosine noise schedule.\"\"\"\n",
    "    ramp = torch.cos(torch.linspace(0, math.pi / 2, n)) ** 2\n",
    "    sigmas = torch.exp(ramp * (math.log(sigma_max) - math.log(sigma_min)) + math.log(sigma_min))\n",
    "    return append_zero(sigmas)\n",
    "\n",
    "def get_sigmas_sigmoid(n, sigma_min, sigma_max, device=\"cpu\"):\n",
    "    \"\"\"Constructs a sigmoid noise schedule.\"\"\"\n",
    "    ramp = 1 / (1 + torch.exp(10 * (torch.linspace(0, 1, n) - 0.5)))\n",
    "    sigmas = torch.exp(ramp * (math.log(sigma_max) - math.log(sigma_min)) + math.log(sigma_min))\n",
    "    return append_zero(sigmas)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "steps = 16\n",
    "sigma_min = 0.5\n",
    "sigma_max = 50.0\n",
    "rho = 1.0\n",
    "\n",
    "configs = [\n",
    "  {\"steps\": steps, \"sigma_min\": 0.5, \"sigma_max\": 50.0, \"rho\": 1.0},\n",
    " # {\"steps\": steps, \"sigma_min\": 0.01, \"sigma_max\": 50.0, \"rho\": 1.0},\n",
    " # {\"steps\": steps, \"sigma_min\": 0.1, \"sigma_max\": 50.0, \"rho\": 1.0},\n",
    " # {\"steps\": steps, \"sigma_min\": 0.5, \"sigma_max\": 50.0, \"rho\": 1.0},\n",
    " # {\"steps\": steps, \"sigma_min\": 0.5, \"sigma_max\": 50.0, \"rho\": 0.5},\n",
    "\n",
    "]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for config in configs:\n",
    "    sigmas = get_sigmas_polyexponential(config[\"steps\"], config[\"sigma_min\"], config[\"sigma_max\"], config[\"rho\"], device=\"cuda\")\n",
    "    plt.plot(sigmas.cpu().numpy(), label=f\"steps={config['steps']}, sigma_min={config['sigma_min']}, sigma_max={config['sigma_max']}, rho={config['rho']}\")\n",
    "    sigmas = get_sigmas_cosine(config[\"steps\"], config[\"sigma_min\"], config[\"sigma_max\"], device=\"cuda\")\n",
    "    plt.plot(sigmas.cpu().numpy(), label=f\"steps={config['steps']}, sigma_min={config['sigma_min']}, sigma_max={config['sigma_max']}\")\n",
    "    sigmas = get_sigmas_sigmoid(config[\"steps\"], config[\"sigma_min\"], config[\"sigma_max\"], device=\"cuda\")\n",
    "    plt.plot(sigmas.cpu().numpy(), label=f\"steps={config['steps']}, sigma_min={config['sigma_min']}, sigma_max={config['sigma_max']}\")\n",
    "plt.legend()\n",
    "plt.show()  "
   ]
  },
  {
   "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
}