Inference.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1fc06181",
   "metadata": {},
   "source": [
    "# Inference Notebook\n",
    "- This notebook serves to import trained models and generate new samples.\n",
    "- You should only need to edits the [Checkpoint & Configs](#Checkpoint-\\&-Configs) and [Define Sample Parameters](#Define-Sample-Parameters) cells.\n",
    "- Currently this notebook only offers unconditional generation, but I plan to include more features in the future.\n",
    "- Have fun creating new sounds!\n",
    "\n",
    "*(NOTE: Don't do \"run all\" in Jupyter. For some reason it doesn't output anything when I used this option. It may work in other environments, but just a heads up!)*"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7b67d821",
   "metadata": {},
   "source": [
    "#### Imports\n",
    "Import necessary libraries to run the notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dc381cc",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Imports\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torchaudio\n",
    "from torch import nn\n",
    "import pytorch_lightning as pl\n",
    "from ema_pytorch import EMA\n",
    "import IPython.display as ipd\n",
    "import yaml\n",
    "from audio_diffusion_pytorch import DiffusionModel, UNetV0, VDiffusion, VSampler\n",
    "from diffusion import sampling, utils"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "49d620cb",
   "metadata": {},
   "source": [
    "### Checkpoint & Configs\n",
    "- Replace these paths with the path to your model's checkpoint and configs.\n",
    "- Pre-trained models are availlable to download on Hugging Face.\n",
    "\n",
    "|Model|Link|\n",
    "|---|---|\n",
    "|Kicks|[crlandsc/tiny-audio-diffusion-kicks](https://huggingface.co/crlandsc/tiny-audio-diffusion-kicks)|\n",
    "|Snares|[crlandsc/tiny-audio-diffusion-snares](https://huggingface.co/crlandsc/tiny-audio-diffusion-snares)|\n",
    "|Hi-hats|[crlandsc/tiny-audio-diffusion-hihats](https://huggingface.co/crlandsc/tiny-audio-diffusion-hihats)|\n",
    "|Percussion (all drum types)|[crlandsc/tiny-audio-diffusion-percussion](https://huggingface.co/crlandsc/tiny-audio-diffusion-percussion)|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5037ead6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load model checkpoint\n",
    "ckpt_path = \"./saved_models/kicks/kicks_v7.ckpt\" # path to model checkpoint\n",
    "config_path = \"./saved_models/kicks/config.yaml\" # path to model config"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f2f61804",
   "metadata": {},
   "source": [
    "### Functions & Models\n",
    "- Functions and models definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46eec999",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load configs\n",
    "with open(config_path, 'r') as file:\n",
    "    config = yaml.safe_load(file)\n",
    "pl_configs = config['model']\n",
    "model_configs = config['model']['model']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4797122",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_mel_spectrogram(sample):\n",
    "    transform = torchaudio.transforms.MelSpectrogram(\n",
    "        sample_rate=sr,\n",
    "        n_fft=1024,\n",
    "        hop_length=512,\n",
    "        n_mels=80,\n",
    "        center=True,\n",
    "        norm=\"slaney\",\n",
    "    )\n",
    "\n",
    "    spectrogram = transform(torch.mean(sample, dim=0)) # downmix and cal spectrogram\n",
    "    spectrogram = torchaudio.functional.amplitude_to_DB(spectrogram, 1.0, 1e-10, 80.0)\n",
    "\n",
    "    # Plot the Mel spectrogram\n",
    "    fig = plt.figure(figsize=(7, 4))\n",
    "    plt.imshow(spectrogram, aspect='auto', origin='lower')\n",
    "    plt.colorbar(format='%+2.0f dB')\n",
    "    plt.xlabel('Frame')\n",
    "    plt.ylabel('Mel Bin')\n",
    "    plt.title('Mel Spectrogram')\n",
    "    plt.tight_layout()\n",
    "    \n",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e246c0e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define PyTorch Lightning model\n",
    "class Model(pl.LightningModule):\n",
    "    def __init__(\n",
    "        self,\n",
    "        lr: float,\n",
    "        lr_beta1: float,\n",
    "        lr_beta2: float,\n",
    "        lr_eps: float,\n",
    "        lr_weight_decay: float,\n",
    "        ema_beta: float,\n",
    "        ema_power: float,\n",
    "        model: nn.Module,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.lr = lr\n",
    "        self.lr_beta1 = lr_beta1\n",
    "        self.lr_beta2 = lr_beta2\n",
    "        self.lr_eps = lr_eps\n",
    "        self.lr_weight_decay = lr_weight_decay\n",
    "        self.model = model\n",
    "        self.model_ema = EMA(self.model, beta=ema_beta, power=ema_power)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5b2aecab",
   "metadata": {},
   "source": [
    "### Instantiate model\n",
    "*NOTE: This model setup needs to exactly match the model that was trained*\n",
    "\n",
    "- This cell instantiates the model (no weights) using the config.yaml file. This is structure is critical to make sure that the model weights can be loaded in correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d626c6e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate model (must match model that was trained)\n",
    "\n",
    "# Diffusion model\n",
    "model = DiffusionModel(\n",
    "    net_t=UNetV0, # The model type used for diffusion (U-Net V0 in this case)\n",
    "    in_channels=model_configs['in_channels'], # U-Net: number of input/output (audio) channels\n",
    "    channels=model_configs['channels'], # U-Net: channels at each layer\n",
    "    factors=model_configs['factors'], # U-Net: downsampling and upsampling factors at each layer\n",
    "    items=model_configs['items'], # U-Net: number of repeating items at each layer\n",
    "    attentions=model_configs['attentions'], # U-Net: attention enabled/disabled at each layer\n",
    "    attention_heads=model_configs['attention_heads'], # U-Net: number of attention heads per attention item\n",
    "    attention_features=model_configs['attention_features'], # U-Net: number of attention features per attention item\n",
    "    diffusion_t=VDiffusion, # The diffusion method used\n",
    "    sampler_t=VSampler # The diffusion sampler used\n",
    ")\n",
    "\n",
    "# pl model\n",
    "model = Model(\n",
    "    lr=pl_configs['lr'],\n",
    "    lr_beta1=pl_configs['lr_beta1'],\n",
    "    lr_beta2=pl_configs['lr_beta2'],\n",
    "    lr_eps=pl_configs['lr_eps'],\n",
    "    lr_weight_decay=pl_configs['lr_weight_decay'],\n",
    "    ema_beta=pl_configs['ema_beta'],\n",
    "    ema_power=pl_configs['ema_power'],\n",
    "    model=model\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c2d2f702",
   "metadata": {},
   "source": [
    "### Check if GPU available\n",
    "- This checks to see if a CUDe capable GPU is available to utilize. If so, the model is assigned to the GPU. If not, the model simply remains on the CPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ce84487",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assign to GPU\n",
    "if torch.cuda.is_available():\n",
    "    model = model.to('cuda')\n",
    "    print(f\"Device: {model.device}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "825d96d9",
   "metadata": {},
   "source": [
    "### Load model\n",
    "- This cell loads the checkpoint weights into the model. It should return `\"<All keys matched successfully>\"` if successfully loaded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "845993e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load model checkpoint\n",
    "checkpoint = torch.load(ckpt_path, map_location='cpu')['state_dict']\n",
    "model.load_state_dict(checkpoint) # should output \"<All keys matched successfully>\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "aaaa996f",
   "metadata": {},
   "source": [
    "## Unconditional Sample Generation\n",
    "Generate new sounds from noise with no conditioning."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c93dc280",
   "metadata": {},
   "source": [
    "#### Define Sample Parameters\n",
    "- sample_length: how long to make the output (measured in samples). Recommended $2^{15}=32768$ (~0.75 sec), as that is what the model was trained on.\n",
    "- sr (sample rate): sampling rate to output. Recommended industry standard 44.1kHz (44100Hz).\n",
    "- num_samples: number of new samples that will be generated.\n",
    "- num_steps: number of diffusion steps - tradeoff inference speed for sample quality (10-100 is a good range).\n",
    "    - 10+ steps - quick generation, alright samples but noticeable high-freq hiss.\n",
    "    - 50+ steps - moderate generation speed, good tradeoff for speed and qualiy (less high-freq hiss).\n",
    "    - 100+ steps - slow generation speed, high quality samples.\n",
    "\n",
    "Have fun playing around with these parameters! Note that sometimes the model outputs some wild things. This is likely due to the small size of the model as well as the limited training data. Larger models and/or larger and more diverse datasets would improve this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f0438d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define diffusion paramters\n",
    "sample_length = 2**15 # 32768 samples @ 44100 = .75 sec\n",
    "sr = 44100\n",
    "num_samples = 3 # number of samples to generate\n",
    "num_steps = 50 # number of diffusion steps, tradeoff inference speed for sample quality (10-100 is a good range)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4952b429",
   "metadata": {},
   "source": [
    "#### Generate samples\n",
    "Run the following cell to generate samples based on previously defined parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf88849c",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    all_samples = torch.zeros(2, 0)\n",
    "    for i in range(num_samples):\n",
    "        noise = torch.randn((1, 2, sample_length), device=model.device) # [batch_size, in_channels, length]\n",
    "        generated_sample = model.model_ema.ema_model.sample(noise, num_steps=num_steps).squeeze(0).cpu() # Suggested num_steps 10-100\n",
    "\n",
    "        print(f\"Generated Sample {i+1}\")\n",
    "        display(ipd.Audio(generated_sample, rate=sr))\n",
    "        \n",
    "        # concatenate all samples:\n",
    "        all_samples = torch.concat((all_samples, generated_sample), dim=1)\n",
    "        \n",
    "        fig = plot_mel_spectrogram(generated_sample)\n",
    "        plt.title(f\"Mel Spectrogram (Sample {i+1})\")\n",
    "        plt.show()\n",
    "        \n",
    "        torch.cuda.empty_cache()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c18bda2c",
   "metadata": {},
   "source": [
    "#### Combine all samples\n",
    "- Option to combine all samples into a single sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8003140",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional concatenate all samples\n",
    "print(f\"All Samples\")\n",
    "display(ipd.Audio(all_samples, rate=sr))\n",
    "fig = plot_mel_spectrogram(all_samples)\n",
    "plt.title(f\"Mel Spectrogram)\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9b5cc6ca",
   "metadata": {},
   "source": [
    "## Conditional \"Style-Transfer\" Generation\n",
    "Generate new sounds conditioned on input audio.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f122e399",
   "metadata": {},
   "source": [
    "#### Define Sample Parameters\n",
    "- audio_file_path: Path to audio file for conditioning the model.\n",
    "- sample_with_noise: Option to output the conditioning sample with noise added to listen, or suppress it.\n",
    "- trim_sample: Option to trim/pad sample if it is too long/short.\n",
    "- sample_length: how long to make the output (measured in samples). Recommended $2^{15}=32768$ (~0.75 sec), as that is what the model was trained on.\n",
    "- sr (sample rate): sampling rate to output. Recommended industry standard 44.1kHz (44100Hz).\n",
    "- num_samples: number of new samples that will be generated.\n",
    "- noise_level: The amount of noise to be added to the input sample.\n",
    "- num_steps: number of diffusion steps - tradeoff inference speed for sample quality.\n",
    "  - The number of steps for conditional diffusion varies more compared to unconditional diffusion. For example, if you input a transient sound (like a snare hit) and want to transfer it to the `kicks` model, then you may not want to add any noise and keep the steps below 10 for an interesting sound. But, if you want to transfer something like a guitar to the percussion model, you may want to add some more noise and increase the number of steps.\n",
    "\n",
    "*NOTE:* The less noise that is added to a sample, the less varied the outputs will be. For example, if you ad 0 noise to a sample and generate it 3 times, it will produce the exact same thing 3 times (because the input remains consistent). As you increase the noise added, increasing the variation of the inputs, the outputs will vary more widely as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38a62339",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define diffusion paramters\n",
    "audio_file_path = \"samples/snare1.wav\"\n",
    "\n",
    "# Listen to noised sample\n",
    "sample_with_noise = False # True to listen to sample + noise, false to not output\n",
    "\n",
    "# If sample too long\n",
    "trim_sample = False # True - if sample too long / False does not trim\n",
    "sample_length = 2**15 # NA\n",
    "\n",
    "sr = 44100 # Sampling rate\n",
    "num_samples = 1 # number of samples to generate\n",
    "noise_level = 0 # between 0 and 1\n",
    "num_steps = 6 # number of diffusion steps, tradeoff inference speed for sample quality (10-100 is a good range)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "15298646",
   "metadata": {},
   "source": [
    "#### Generate samples\n",
    "Run the following cell to generate samples based on previously defined parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bf5203d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate samples\n",
    "with torch.no_grad():\n",
    "\n",
    "    # load audio sample\n",
    "    audio_sample = torchaudio.load(audio_file_path)[0].unsqueeze(0).to(model.device) # unsqueeze for correct tensor shape\n",
    "\n",
    "    # Trim audio\n",
    "    if trim_sample:\n",
    "        og_shape = audio_sample.shape\n",
    "        if sample_length < og_shape[2]:\n",
    "            audio_sample = audio_sample[:,:,:sample_length]\n",
    "        elif sample_length > og_shape[2]:\n",
    "            # Pad tensor with zeros to match sample length\n",
    "            audio_sample = torch.concat((audio_sample, torch.zeros(og_shape[0], og_shape[1], sample_length - og_shape[2]).to(model.device)), dim=2)\n",
    "\n",
    "\n",
    "    original_audio = audio_sample.squeeze(0).squeeze(0).cpu()\n",
    "\n",
    "    # Display original audio sample\n",
    "    print(f\"Original Sample\")\n",
    "    display(ipd.Audio(original_audio, rate=sr))\n",
    "\n",
    "    # Plot original audio\n",
    "    fig = plot_mel_spectrogram(original_audio)\n",
    "    plt.title(f\"Mel Spectrogram (Original Sample)\")\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "    # Display original audio sample + noise\n",
    "    if sample_with_noise:\n",
    "        noise = torch.randn_like(audio_sample, device=model.device) * noise_level # combine input signal and noise\n",
    "        noised_sample = (audio_sample + noise).squeeze(0).cpu() # normalize?\n",
    "        print(f\"Original Noised Sample\")\n",
    "        display(ipd.Audio(noised_sample, rate=sr))\n",
    "\n",
    "        # Plot original audio + noise\n",
    "        fig = plot_mel_spectrogram(noised_sample)\n",
    "        plt.title(f\"Mel Spectrogram (Noised Sample)\")\n",
    "        plt.show()\n",
    "\n",
    "\n",
    "    all_samples = torch.zeros(2, 0)\n",
    "    for i in range(num_samples):\n",
    "        noise = torch.randn_like(audio_sample, device=model.device) * noise_level # combine input signal and noise\n",
    "        audio = audio_sample + noise # normalize?\n",
    "        generated_sample = model.model_ema.ema_model.sample(audio, num_steps=num_steps).squeeze(0).cpu()\n",
    "\n",
    "        print(f\"Generated Sample {i+1}\")\n",
    "        display(ipd.Audio(generated_sample, rate=sr))\n",
    "        \n",
    "        # concatenate all samples:\n",
    "        all_samples = torch.concat((all_samples, generated_sample), dim=1)\n",
    "        \n",
    "        fig = plot_mel_spectrogram(generated_sample)\n",
    "        plt.title(f\"Mel Spectrogram (Sample {i+1})\")\n",
    "        plt.show()\n",
    "        \n",
    "        torch.cuda.empty_cache()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4762cc6b",
   "metadata": {},
   "source": [
    "#### Combine all samples\n",
    "- Option to combine all samples into a single sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5501ce6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional concatenate all samples\n",
    "print(f\"All Samples\")\n",
    "display(ipd.Audio(all_samples, rate=sr))\n",
    "fig = plot_mel_spectrogram(all_samples)\n",
    "plt.title(f\"Mel Spectrogram)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ef83b4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Add normalization?\n",
    "# TODO: Add other smapling methods (currently only DDIM)\n",
    "# TODO: clean cell (make functions)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tiny-audio-diffusion (Python 3.10)",
   "language": "python",
   "name": "tiny-audio-diffusion"
  },
  "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}