added latent diffusion to notebook

notebooks/test_model.ipynb

@@ -177,13 +177,13 @@
    },
    "outputs": [],
    "source": [
-    "start_steps = 500  #@param {type:\"slider\", min:0, max:1000, step:10}\n",
+    "start_step = 500  #@param {type:\"slider\", min:0, max:1000, step:10}\n",
     "track = AudioDiffusion.loop_it(audio, sample_rate, loops=1)\n",
     "for variation in range(12):\n",
     "    image2, (\n",
     "        sample_rate,\n",
     "        audio2) = audio_diffusion.generate_spectrogram_and_audio_from_audio(\n",
-    "            raw_audio=audio, start_step=start_steps)\n",
+    "            raw_audio=audio, start_step=start_step)\n",
     "    display(image2)\n",
     "    display(Audio(audio2, rate=sample_rate))\n",
     "    track = np.concatenate(\n",
@@ -490,10 +490,124 @@
     "display(Audio(audio, rate=sample_rate))"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "5b7081f7",
+   "metadata": {},
+   "source": [
+    "## Latent Audio Diffusion\n",
+    "Instead of de-noising images directly in the pixel space, we can work in the latent space of a pre-trained VAE (Variational AutoEncoder). This is much faster to train and run inference on, although the quality suffers as there are now three stages involved in encoding / decoding: mel spectrogram, VAE and de-noising."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "17610772",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_id = \"teticio/latent-audio-diffusion-ddim-256\"  #@param [\"teticio/latent-audio-diffusion-256\", \"teticio/latent-audio-diffusion-ddim-256\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9e6c73e6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "audio_diffusion = AudioDiffusion(model_id=model_id)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d37a03a9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "seed = 6015487092443227811  #@param {type:\"integer\"}\n",
+    "generator.manual_seed(seed)\n",
+    "image, (sample_rate, audio) = audio_diffusion.generate_spectrogram_and_audio(\n",
+    "    generator=generator)\n",
+    "display(image)\n",
+    "display(Audio(audio, rate=sample_rate))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c0328a56",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "seed2 = 5623685468252603494  #@param {type:\"integer\"}\n",
+    "generator.manual_seed(seed2)\n",
+    "image2, (sample_rate, audio2) = audio_diffusion.generate_spectrogram_and_audio(\n",
+    "    generator=generator)\n",
+    "display(image2)\n",
+    "display(Audio(audio2, rate=sample_rate))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bd1f2b58",
+   "metadata": {},
+   "source": [
+    "### Interpolation in latent space\n",
+    "As the VAE forces a more compact, lower dimensional representation for the spectrograms, interpolation in latent space can lead to meaningful combinations of audios. In combination with the (deterministic) DDIM from the previous section, the model can be used as an encoder / decoder to a lower dimensional space."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "23ff0ee7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "generator.manual_seed(seed)\n",
+    "noise = torch.randn((1, audio_diffusion.pipe.unet.in_channels,\n",
+    "                     audio_diffusion.pipe.unet.sample_size[0],\n",
+    "                     audio_diffusion.pipe.unet.sample_size[1]),\n",
+    "                    generator=generator)\n",
+    "noise.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ff13a2cb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "generator.manual_seed(seed2)\n",
+    "noise2 = torch.randn((1, audio_diffusion.pipe.unet.in_channels,\n",
+    "                      audio_diffusion.pipe.unet.sample_size[0],\n",
+    "                      audio_diffusion.pipe.unet.sample_size[1]),\n",
+    "                     generator=generator)\n",
+    "noise2.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bea26a5e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "alpha = 0.5  #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
+    "_, (sample_rate, audio3) = audio_diffusion.generate_spectrogram_and_audio(\n",
+    "    noise=audio_diffusion.pipe.slerp(noise, noise2, alpha),\n",
+    "    generator=generator)\n",
+    "display(Audio(audio, rate=mel.get_sample_rate()))\n",
+    "display(Audio(audio2, rate=mel.get_sample_rate()))\n",
+    "display(Audio(audio3, rate=sample_rate))"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0b05539f",
+   "id": "60080eed",
    "metadata": {},
    "outputs": [],
    "source": []
@@ -520,7 +634,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.9 (default, Apr 13 2022, 08:48:06) \n[Clang 13.1.6 (clang-1316.0.21.2.5)]"
+   "version": "3.10.6"
   },
   "toc": {
    "base_numbering": 1,