DM4ASC.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "701920bc",
   "metadata": {},
   "source": [
    "# Diffusion Models for Audio Semantic Communication"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae6fb974-42a4-446a-ae48-7645c90f6264",
   "metadata": {},
   "source": [
    "### Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db736330-cc8b-4233-8300-6290e26ffc0b",
   "metadata": {},
   "source": [
    "Run once to set all the needed stuff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e26ea29",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install librosa==0.9.2\n",
    "!pip install huggingface_hub==0.13.3\n",
    "!pip install einops==0.6.1\n",
    "!pip install transformers==4.27.0\n",
    "!pip install progressbar\n",
    "!pip install pandas\n",
    "!pip install matplotlib\n",
    "!pip install torchmetrics\n",
    "!pip install frechet-audio-distance\n",
    "\n",
    "# Install PyTorch version 1.13.1 with CUDA 11.7 support\n",
    "!pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu117\n",
    "\n",
    "%cd diffusers/\n",
    "!pip install -e .\n",
    "    \n",
    "!git clone https://github.com/declare-lab/tango\n",
    "%cd tango\n",
    "%cd 'diffusers'\n",
    "!pip install -e .\n",
    "%cd ../tango/\n",
    "%mkdir weights\n",
    "%cd weights\n",
    "!git clone https://huggingface.co/declare-lab/tango-full-ft-audiocaps"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d7cddc9-a311-4d90-81c4-519a9c8a0a26",
   "metadata": {},
   "source": [
    "## Initialize tango"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0e4efc13-2cd3-494c-bd6a-172ad890cd3c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/mnt/media/christian/DiffInpainting/tango\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ispamm/miniconda3/envs/DiffInp/lib/python3.10/site-packages/IPython/core/magics/osm.py:417: UserWarning: This is now an optional IPython functionality, setting dhist requires you to install the `pickleshare` library.\n",
      "  self.shell.db['dhist'] = compress_dhist(dhist)[-100:]\n",
      "/home/ispamm/miniconda3/envs/DiffInp/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "%cd tango\n",
    "import os\n",
    "\n",
    "import json\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "from huggingface_hub import snapshot_download\n",
    "from models import AudioDiffusion, DDPMScheduler\n",
    "from audioldm.audio.stft import TacotronSTFT\n",
    "from audioldm.variational_autoencoder import AutoencoderKL\n",
    "import IPython\n",
    "import soundfile as sf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dfaa9b67",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Tango:\n",
    "    def __init__(self, path_to_weights='', device=\"cuda:0\"):\n",
    "\n",
    "        if path_to_weights=='':\n",
    "            path_to_weights = './weights/tango-full-ft-audiocaps'\n",
    "\n",
    "        vae_config = json.load(open(\"{}/vae_config.json\".format(path_to_weights)))\n",
    "        stft_config = json.load(open(\"{}/stft_config.json\".format(path_to_weights)))\n",
    "        main_config = json.load(open(\"{}/main_config.json\".format(path_to_weights)))\n",
    "\n",
    "        self.vae = AutoencoderKL(**vae_config).to(device)\n",
    "        self.stft = TacotronSTFT(**stft_config).to(device)\n",
    "        self.model = AudioDiffusion(**main_config).to(device)\n",
    "\n",
    "        vae_weights = torch.load(\"{}/pytorch_model_vae.bin\".format(path_to_weights), map_location=device)\n",
    "        stft_weights = torch.load(\"{}/pytorch_model_stft.bin\".format(path_to_weights), map_location=device)\n",
    "        main_weights = torch.load(\"{}/pytorch_model_main.bin\".format(path_to_weights), map_location=device)\n",
    "\n",
    "        self.vae.load_state_dict(vae_weights)\n",
    "        self.stft.load_state_dict(stft_weights)\n",
    "        self.model.load_state_dict(main_weights)\n",
    "\n",
    "        print (\"Successfully loaded checkpoint from:\", path_to_weights)\n",
    "\n",
    "        self.vae.eval()\n",
    "        self.stft.eval()\n",
    "        self.model.eval()\n",
    "\n",
    "        self.scheduler = DDPMScheduler.from_pretrained(main_config[\"scheduler_name\"], subfolder=\"scheduler\")\n",
    "\n",
    "    def chunks(self, lst, n):\n",
    "        \"\"\" Yield successive n-sized chunks from a list. \"\"\"\n",
    "        for i in range(0, len(lst), n):\n",
    "            yield lst[i:i + n]\n",
    "\n",
    "    def generate(self, prompt, steps=100, guidance=3, samples=1, disable_progress=True):\n",
    "        \"\"\" Generate audio for a single prompt string. \"\"\"\n",
    "        with torch.no_grad():\n",
    "            latents = self.model.inference([prompt], self.scheduler, steps, guidance, samples, disable_progress=disable_progress)\n",
    "            mel = self.vae.decode_first_stage(latents)\n",
    "            wave = self.vae.decode_to_waveform(mel)\n",
    "        return wave[0]\n",
    "\n",
    "    def generate_for_batch(self, prompts, steps=100, guidance=3, samples=1, batch_size=8, disable_progress=True):\n",
    "        \"\"\" Generate audio for a list of prompt strings. \"\"\"\n",
    "        outputs = []\n",
    "        for k in tqdm(range(0, len(prompts), batch_size)):\n",
    "            batch = prompts[k: k+batch_size]\n",
    "            with torch.no_grad():\n",
    "                latents = self.model.inference(batch, self.scheduler, steps, guidance, samples, disable_progress=disable_progress)\n",
    "                mel = self.vae.decode_first_stage(latents)\n",
    "                wave = self.vae.decode_to_waveform(mel)\n",
    "                outputs += [item for item in wave]\n",
    "        if samples == 1:\n",
    "            return outputs\n",
    "        else:\n",
    "            return list(self.chunks(outputs, samples))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f2b7df8e-a49a-46cc-9237-b6a563c79a9e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mnt/media/christian/DiffInpainting/tango/audioldm/audio/stft.py:42: FutureWarning: Pass size=1024 as keyword args. From version 0.10 passing these as positional arguments will result in an error\n",
      "  fft_window = pad_center(fft_window, filter_length)\n",
      "/mnt/media/christian/DiffInpainting/tango/audioldm/audio/stft.py:151: FutureWarning: Pass sr=16000, n_fft=1024, n_mels=64, fmin=0, fmax=8000 as keyword args. From version 0.10 passing these as positional arguments will result in an error\n",
      "  mel_basis = librosa_mel_fn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "UNet initialized randomly.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at google/flan-t5-large were not used when initializing T5EncoderModel: ['decoder.block.9.layer.1.layer_norm.weight', 'decoder.block.20.layer.1.EncDecAttention.k.weight', 'decoder.block.0.layer.2.DenseReluDense.wo.weight', 'decoder.block.21.layer.0.SelfAttention.q.weight', 'decoder.block.5.layer.0.SelfAttention.q.weight', 'decoder.block.11.layer.1.layer_norm.weight', 'decoder.block.14.layer.2.DenseReluDense.wo.weight', 'decoder.block.18.layer.1.EncDecAttention.q.weight', 'decoder.block.23.layer.2.DenseReluDense.wo.weight', 'decoder.block.7.layer.0.SelfAttention.o.weight', 'decoder.block.15.layer.1.layer_norm.weight', 'decoder.block.22.layer.0.SelfAttention.k.weight', 'decoder.block.15.layer.1.EncDecAttention.v.weight', 'decoder.block.8.layer.1.EncDecAttention.q.weight', 'decoder.block.15.layer.1.EncDecAttention.q.weight', 'decoder.block.21.layer.2.layer_norm.weight', 'decoder.block.2.layer.0.SelfAttention.q.weight', 'decoder.block.1.layer.1.layer_norm.weight', 'decoder.block.10.layer.1.EncDecAttention.o.weight', 'decoder.block.4.layer.2.layer_norm.weight', 'decoder.embed_tokens.weight', 'decoder.block.15.layer.1.EncDecAttention.k.weight', 'decoder.block.21.layer.0.SelfAttention.k.weight', 'decoder.final_layer_norm.weight', 'decoder.block.6.layer.1.EncDecAttention.k.weight', 'decoder.block.4.layer.1.EncDecAttention.v.weight', 'decoder.block.12.layer.1.EncDecAttention.k.weight', 'decoder.block.2.layer.2.DenseReluDense.wo.weight', 'decoder.block.4.layer.1.EncDecAttention.q.weight', 'decoder.block.22.layer.0.SelfAttention.q.weight', 'decoder.block.11.layer.1.EncDecAttention.q.weight', 'decoder.block.7.layer.2.layer_norm.weight', 'decoder.block.12.layer.0.SelfAttention.k.weight', 'decoder.block.11.layer.1.EncDecAttention.o.weight', 'decoder.block.20.layer.1.EncDecAttention.o.weight', 'decoder.block.8.layer.0.SelfAttention.v.weight', 'decoder.block.19.layer.1.EncDecAttention.q.weight', 'decoder.block.0.layer.0.SelfAttention.v.weight', 'decoder.block.18.layer.0.SelfAttention.o.weight', 'decoder.block.8.layer.1.layer_norm.weight', 'decoder.block.5.layer.1.layer_norm.weight', 'decoder.block.1.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.15.layer.0.layer_norm.weight', 'decoder.block.11.layer.0.SelfAttention.o.weight', 'decoder.block.9.layer.1.EncDecAttention.k.weight', 'decoder.block.8.layer.2.layer_norm.weight', 'decoder.block.16.layer.2.layer_norm.weight', 'decoder.block.6.layer.0.SelfAttention.k.weight', 'decoder.block.14.layer.1.EncDecAttention.o.weight', 'decoder.block.22.layer.1.EncDecAttention.q.weight', 'decoder.block.14.layer.2.layer_norm.weight', 'decoder.block.4.layer.2.DenseReluDense.wo.weight', 'decoder.block.9.layer.1.EncDecAttention.q.weight', 'decoder.block.2.layer.0.layer_norm.weight', 'decoder.block.20.layer.0.SelfAttention.o.weight', 'decoder.block.9.layer.0.SelfAttention.o.weight', 'decoder.block.3.layer.1.layer_norm.weight', 'decoder.block.16.layer.1.EncDecAttention.v.weight', 'decoder.block.13.layer.2.layer_norm.weight', 'decoder.block.21.layer.1.EncDecAttention.k.weight', 'decoder.block.14.layer.1.EncDecAttention.q.weight', 'decoder.block.16.layer.1.layer_norm.weight', 'decoder.block.21.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.16.layer.2.DenseReluDense.wo.weight', 'decoder.block.16.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.13.layer.0.SelfAttention.q.weight', 'decoder.block.18.layer.1.EncDecAttention.o.weight', 'decoder.block.22.layer.0.SelfAttention.v.weight', 'decoder.block.6.layer.0.SelfAttention.o.weight', 'decoder.block.12.layer.0.SelfAttention.o.weight', 'decoder.block.3.layer.0.SelfAttention.o.weight', 'decoder.block.21.layer.2.DenseReluDense.wo.weight', 'decoder.block.10.layer.1.EncDecAttention.q.weight', 'decoder.block.10.layer.0.SelfAttention.o.weight', 'decoder.block.7.layer.1.EncDecAttention.o.weight', 'decoder.block.0.layer.0.SelfAttention.k.weight', 'decoder.block.22.layer.0.layer_norm.weight', 'decoder.block.11.layer.1.EncDecAttention.k.weight', 'decoder.block.17.layer.1.EncDecAttention.q.weight', 'decoder.block.0.layer.1.layer_norm.weight', 'decoder.block.15.layer.0.SelfAttention.k.weight', 'decoder.block.6.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.10.layer.2.layer_norm.weight', 'decoder.block.7.layer.2.DenseReluDense.wo.weight', 'decoder.block.17.layer.0.SelfAttention.v.weight', 'decoder.block.20.layer.0.layer_norm.weight', 'decoder.block.14.layer.0.SelfAttention.k.weight', 'decoder.block.2.layer.1.EncDecAttention.k.weight', 'decoder.block.19.layer.0.layer_norm.weight', 'decoder.block.15.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.6.layer.0.SelfAttention.v.weight', 'decoder.block.7.layer.0.SelfAttention.q.weight', 'decoder.block.22.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.2.layer.0.SelfAttention.v.weight', 'decoder.block.7.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.12.layer.0.layer_norm.weight', 'decoder.block.0.layer.1.EncDecAttention.k.weight', 'decoder.block.18.layer.2.DenseReluDense.wo.weight', 'decoder.block.23.layer.0.SelfAttention.v.weight', 'decoder.block.11.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.15.layer.2.DenseReluDense.wo.weight', 'decoder.block.7.layer.1.EncDecAttention.k.weight', 'decoder.block.16.layer.0.SelfAttention.v.weight', 'decoder.block.5.layer.1.EncDecAttention.v.weight', 'decoder.block.9.layer.2.layer_norm.weight', 'decoder.block.15.layer.1.EncDecAttention.o.weight', 'decoder.block.23.layer.1.layer_norm.weight', 'decoder.block.0.layer.1.EncDecAttention.o.weight', 'decoder.block.14.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.1.layer.0.layer_norm.weight', 'decoder.block.6.layer.2.DenseReluDense.wo.weight', 'decoder.block.12.layer.1.EncDecAttention.q.weight', 'decoder.block.17.layer.2.DenseReluDense.wo.weight', 'decoder.block.5.layer.0.SelfAttention.k.weight', 'decoder.block.14.layer.1.EncDecAttention.k.weight', 'decoder.block.19.layer.0.SelfAttention.o.weight', 'decoder.block.13.layer.1.layer_norm.weight', 'decoder.block.22.layer.1.EncDecAttention.k.weight', 'decoder.block.0.layer.0.SelfAttention.o.weight', 'decoder.block.1.layer.2.DenseReluDense.wo.weight', 'decoder.block.4.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.19.layer.1.layer_norm.weight', 'decoder.block.17.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.19.layer.0.SelfAttention.k.weight', 'decoder.block.23.layer.0.SelfAttention.k.weight', 'decoder.block.0.layer.1.EncDecAttention.v.weight', 'decoder.block.17.layer.0.SelfAttention.k.weight', 'decoder.block.5.layer.1.EncDecAttention.o.weight', 'decoder.block.11.layer.0.SelfAttention.q.weight', 'decoder.block.14.layer.0.SelfAttention.v.weight', 'decoder.block.4.layer.0.SelfAttention.v.weight', 'decoder.block.3.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.23.layer.0.layer_norm.weight', 'decoder.block.22.layer.2.layer_norm.weight', 'decoder.block.4.layer.1.layer_norm.weight', 'decoder.block.13.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.6.layer.1.EncDecAttention.q.weight', 'decoder.block.0.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.13.layer.0.SelfAttention.v.weight', 'decoder.block.5.layer.0.layer_norm.weight', 'decoder.block.12.layer.0.SelfAttention.q.weight', 'decoder.block.19.layer.0.SelfAttention.q.weight', 'decoder.block.20.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.7.layer.0.SelfAttention.k.weight', 'decoder.block.5.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.13.layer.1.EncDecAttention.k.weight', 'decoder.block.15.layer.0.SelfAttention.q.weight', 'decoder.block.22.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.17.layer.1.EncDecAttention.v.weight', 'decoder.block.12.layer.1.EncDecAttention.v.weight', 'decoder.block.12.layer.1.EncDecAttention.o.weight', 'decoder.block.22.layer.1.layer_norm.weight', 'decoder.block.17.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.7.layer.1.EncDecAttention.q.weight', 'decoder.block.1.layer.0.SelfAttention.q.weight', 'decoder.block.20.layer.2.layer_norm.weight', 'decoder.block.6.layer.1.layer_norm.weight', 'decoder.block.15.layer.0.SelfAttention.o.weight', 'decoder.block.8.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.21.layer.0.SelfAttention.o.weight', 'decoder.block.4.layer.1.EncDecAttention.k.weight', 'decoder.block.11.layer.0.layer_norm.weight', 'decoder.block.23.layer.1.EncDecAttention.k.weight', 'decoder.block.12.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.19.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.17.layer.1.EncDecAttention.k.weight', 'decoder.block.23.layer.0.SelfAttention.q.weight', 'decoder.block.15.layer.0.SelfAttention.v.weight', 'decoder.block.11.layer.0.SelfAttention.k.weight', 'decoder.block.3.layer.0.SelfAttention.k.weight', 'decoder.block.5.layer.0.SelfAttention.o.weight', 'decoder.block.10.layer.0.SelfAttention.v.weight', 'decoder.block.14.layer.0.SelfAttention.o.weight', 'decoder.block.20.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.4.layer.1.EncDecAttention.o.weight', 'decoder.block.20.layer.0.SelfAttention.k.weight', 'decoder.block.19.layer.2.DenseReluDense.wo.weight', 'decoder.block.16.layer.0.SelfAttention.o.weight', 'decoder.block.11.layer.0.SelfAttention.v.weight', 'decoder.block.14.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.10.layer.0.SelfAttention.k.weight', 'decoder.block.12.layer.0.SelfAttention.v.weight', 'decoder.block.21.layer.0.SelfAttention.v.weight', 'decoder.block.22.layer.0.SelfAttention.o.weight', 'decoder.block.2.layer.0.SelfAttention.k.weight', 'decoder.block.12.layer.2.DenseReluDense.wo.weight', 'decoder.block.4.layer.0.layer_norm.weight', 'decoder.block.12.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.23.layer.1.EncDecAttention.q.weight', 'decoder.block.23.layer.2.layer_norm.weight', 'decoder.block.17.layer.1.EncDecAttention.o.weight', 'decoder.block.8.layer.1.EncDecAttention.v.weight', 'decoder.block.2.layer.1.EncDecAttention.q.weight', 'decoder.block.5.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.3.layer.2.DenseReluDense.wo.weight', 'decoder.block.18.layer.0.SelfAttention.v.weight', 'decoder.block.5.layer.1.EncDecAttention.k.weight', 'decoder.block.23.layer.1.EncDecAttention.v.weight', 'decoder.block.0.layer.2.layer_norm.weight', 'decoder.block.21.layer.1.layer_norm.weight', 'decoder.block.9.layer.2.DenseReluDense.wo.weight', 'decoder.block.1.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.12.layer.1.layer_norm.weight', 'decoder.block.23.layer.1.EncDecAttention.o.weight', 'decoder.block.14.layer.0.layer_norm.weight', 'decoder.block.13.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.21.layer.1.EncDecAttention.v.weight', 'decoder.block.8.layer.0.SelfAttention.k.weight', 'decoder.block.9.layer.0.SelfAttention.v.weight', 'decoder.block.8.layer.1.EncDecAttention.o.weight', 'decoder.block.16.layer.0.layer_norm.weight', 'decoder.block.10.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.19.layer.2.layer_norm.weight', 'decoder.block.3.layer.1.EncDecAttention.v.weight', 'decoder.block.4.layer.0.SelfAttention.q.weight', 'decoder.block.0.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.17.layer.2.layer_norm.weight', 'decoder.block.13.layer.0.SelfAttention.k.weight', 'decoder.block.18.layer.1.EncDecAttention.v.weight', 'decoder.block.3.layer.0.layer_norm.weight', 'decoder.block.17.layer.0.SelfAttention.q.weight', 'decoder.block.0.layer.1.EncDecAttention.q.weight', 'decoder.block.15.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.7.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.2.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.13.layer.2.DenseReluDense.wo.weight', 'decoder.block.18.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.13.layer.1.EncDecAttention.v.weight', 'decoder.block.6.layer.0.layer_norm.weight', 'decoder.block.13.layer.0.SelfAttention.o.weight', 'decoder.block.1.layer.1.EncDecAttention.o.weight', 'decoder.block.18.layer.0.SelfAttention.k.weight', 'decoder.block.9.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.19.layer.1.EncDecAttention.o.weight', 'decoder.block.10.layer.1.layer_norm.weight', 'decoder.block.13.layer.0.layer_norm.weight', 'decoder.block.8.layer.0.SelfAttention.q.weight', 'decoder.block.22.layer.1.EncDecAttention.v.weight', 'decoder.block.5.layer.1.EncDecAttention.q.weight', 'decoder.block.11.layer.1.EncDecAttention.v.weight', 'decoder.block.7.layer.0.SelfAttention.v.weight', 'decoder.block.19.layer.0.SelfAttention.v.weight', 'decoder.block.9.layer.0.layer_norm.weight', 'decoder.block.16.layer.1.EncDecAttention.k.weight', 'decoder.block.21.layer.1.EncDecAttention.q.weight', 'decoder.block.18.layer.0.SelfAttention.q.weight', 'decoder.block.0.layer.0.SelfAttention.q.weight', 'decoder.block.10.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.16.layer.1.EncDecAttention.q.weight', 'decoder.block.4.layer.0.SelfAttention.k.weight', 'decoder.block.20.layer.2.DenseReluDense.wo.weight', 'decoder.block.3.layer.1.EncDecAttention.k.weight', 'decoder.block.19.layer.1.EncDecAttention.v.weight', 'decoder.block.9.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.2.layer.1.layer_norm.weight', 'decoder.block.22.layer.1.EncDecAttention.o.weight', 'decoder.block.8.layer.2.DenseReluDense.wo.weight', 'decoder.block.21.layer.0.layer_norm.weight', 'decoder.block.3.layer.2.layer_norm.weight', 'decoder.block.2.layer.1.EncDecAttention.v.weight', 'decoder.block.20.layer.0.SelfAttention.q.weight', 'decoder.block.21.layer.1.EncDecAttention.o.weight', 'decoder.block.6.layer.1.EncDecAttention.o.weight', 'decoder.block.16.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.9.layer.0.SelfAttention.q.weight', 'decoder.block.4.layer.0.SelfAttention.o.weight', 'decoder.block.11.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.6.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.5.layer.2.layer_norm.weight', 'decoder.block.18.layer.0.layer_norm.weight', 'decoder.block.10.layer.1.EncDecAttention.v.weight', 'decoder.block.13.layer.1.EncDecAttention.o.weight', 'decoder.block.5.layer.2.DenseReluDense.wo.weight', 'decoder.block.8.layer.1.EncDecAttention.k.weight', 'decoder.block.16.layer.0.SelfAttention.k.weight', 'decoder.block.2.layer.2.layer_norm.weight', 'decoder.block.7.layer.0.layer_norm.weight', 'decoder.block.15.layer.2.layer_norm.weight', 'decoder.block.6.layer.2.layer_norm.weight', 'decoder.block.8.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.5.layer.0.SelfAttention.v.weight', 'decoder.block.18.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.6.layer.1.EncDecAttention.v.weight', 'decoder.block.1.layer.0.SelfAttention.o.weight', 'decoder.block.21.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.8.layer.0.layer_norm.weight', 'decoder.block.20.layer.1.EncDecAttention.q.weight', 'decoder.block.10.layer.1.EncDecAttention.k.weight', 'decoder.block.10.layer.0.layer_norm.weight', 'decoder.block.18.layer.1.layer_norm.weight', 'decoder.block.20.layer.1.layer_norm.weight', 'decoder.block.20.layer.1.EncDecAttention.v.weight', 'decoder.block.23.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.10.layer.0.SelfAttention.q.weight', 'decoder.block.17.layer.1.layer_norm.weight', 'decoder.block.3.layer.1.EncDecAttention.o.weight', 'decoder.block.9.layer.0.SelfAttention.k.weight', 'decoder.block.19.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.0.layer.0.layer_norm.weight', 'decoder.block.3.layer.1.EncDecAttention.q.weight', 'decoder.block.16.layer.1.EncDecAttention.o.weight', 'decoder.block.9.layer.1.EncDecAttention.v.weight', 'decoder.block.3.layer.0.SelfAttention.v.weight', 'decoder.block.11.layer.2.layer_norm.weight', 'decoder.block.1.layer.1.EncDecAttention.k.weight', 'decoder.block.2.layer.0.SelfAttention.o.weight', 'decoder.block.2.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.3.layer.2.DenseReluDense.wi_0.weight', 'decoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight', 'lm_head.weight', 'decoder.block.1.layer.0.SelfAttention.k.weight', 'decoder.block.10.layer.2.DenseReluDense.wo.weight', 'decoder.block.8.layer.0.SelfAttention.o.weight', 'decoder.block.17.layer.0.layer_norm.weight', 'decoder.block.20.layer.0.SelfAttention.v.weight', 'decoder.block.14.layer.0.SelfAttention.q.weight', 'decoder.block.16.layer.0.SelfAttention.q.weight', 'decoder.block.19.layer.1.EncDecAttention.k.weight', 'decoder.block.3.layer.0.SelfAttention.q.weight', 'decoder.block.23.layer.0.SelfAttention.o.weight', 'decoder.block.14.layer.1.EncDecAttention.v.weight', 'decoder.block.13.layer.1.EncDecAttention.q.weight', 'decoder.block.22.layer.2.DenseReluDense.wo.weight', 'decoder.block.4.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.2.layer.1.EncDecAttention.o.weight', 'decoder.block.1.layer.1.EncDecAttention.q.weight', 'decoder.block.1.layer.1.EncDecAttention.v.weight', 'decoder.block.12.layer.2.layer_norm.weight', 'decoder.block.23.layer.2.DenseReluDense.wi_1.weight', 'decoder.block.7.layer.1.layer_norm.weight', 'decoder.block.18.layer.1.EncDecAttention.k.weight', 'decoder.block.1.layer.2.layer_norm.weight', 'decoder.block.1.layer.0.SelfAttention.v.weight', 'decoder.block.7.layer.1.EncDecAttention.v.weight', 'decoder.block.6.layer.0.SelfAttention.q.weight', 'decoder.block.17.layer.0.SelfAttention.o.weight', 'decoder.block.18.layer.2.layer_norm.weight', 'decoder.block.14.layer.1.layer_norm.weight', 'decoder.block.11.layer.2.DenseReluDense.wo.weight', 'decoder.block.9.layer.1.EncDecAttention.o.weight']\n",
      "- This IS expected if you are initializing T5EncoderModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing T5EncoderModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Successfully loaded checkpoint from: ./weights/tango-full-ft-audiocaps\n"
     ]
    }
   ],
   "source": [
    "tango = Tango()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5d4e570",
   "metadata": {},
   "source": [
    "## DNNM+ & Repaint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "febbce01-c74c-435b-aaac-f76d30e87747",
   "metadata": {},
   "source": [
    "Change the following parameters to test DDNM+ on the denoising or impainting task.\n",
    "\n",
    "You should first download the AudioCaps dataset and change the paths below accordingly. Please note that different seeds can impact the quality of final results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "49fab913-9a2e-46e5-bda1-b4845a048a72",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/mnt/media/christian/DiffInpainting/versions\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ispamm/miniconda3/envs/DiffInp/lib/python3.10/site-packages/IPython/core/magics/osm.py:417: UserWarning: This is now an optional IPython functionality, setting dhist requires you to install the `pickleshare` library.\n",
      "  self.shell.db['dhist'] = compress_dhist(dhist)[-100:]\n"
     ]
    }
   ],
   "source": [
    "%cd ../versions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "97ff368f-e774-40f4-b0b6-adf2adbce95e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = 'ddnm+'                         # 'ddnm+' or 'repaint'\n",
    "snr = 20 # 17.5, 20, 30                 # Specify the PSNR level of the noisy audio\n",
    "noisy_prompts=True                      # Whether to apply noise to text prompts or not\n",
    "prompt_psnr = snr                       # Here you can control the amount of noise to apply to the prompt embedding\n",
    "time_mask_percentage = (0.45, 0.55)     # IMPAINTING? If not, set (0.,0.)\n",
    "mask_type = 'time'                      # type of data you want to apply the mask to\n",
    "print_report = False                    # True to print final report\n",
    "num_samples = 1                         # Number of audio files with which to test the method\n",
    "save_output_audio = True                # Save the output audio file\n",
    "audio_description = 'ddnmp_snr20'       # filename of the saved files\n",
    "save_noisy_audio = False                # Whether to save noise latents converted to audio or not\n",
    "caption_filter = None                   # Search for specific words in the captions (list of words)\n",
    "\n",
    "name_experiment = 'DDNM+ inpaint SNR20 test'\n",
    "output_path = '../output'\n",
    "dataset_path = '../AudioCaps'\n",
    "dataset_subset = 'AudioCaps_Val'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cc384970",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DDNM+\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████| 1000/1000 [07:13<00:00,  2.31it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 1200x300 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import shutil\n",
    "import utils\n",
    "\n",
    "\n",
    "output_path_generated = os.path.join(output_path, f'{name_experiment}/generated')\n",
    "output_path_noisy  = os.path.join(output_path, f'{name_experiment}/noisy')\n",
    "output_path_clean = os.path.join(output_path, f'{name_experiment}/clean')\n",
    "output_path_spectrograms = os.path.join(output_path, f'{name_experiment}/spectrograms')\n",
    "\n",
    "usecols = [\"youtube_id\", \"start_time\", \"caption\"]\n",
    "file_list = pd.read_csv(os.path.join(dataset_path, \"val.csv\"), index_col=\"youtube_id\", usecols=usecols)\n",
    "\n",
    "d_audios = utils.get_d_audios(file_list, num_to_select=num_samples, dataset_path = os.path.join(dataset_path, dataset_subset), caption_filter=caption_filter, max_f=200)\n",
    "\n",
    "if not os.path.exists(os.path.join(output_path, name_experiment)):\n",
    "    os.mkdir(os.path.join(output_path, name_experiment))\n",
    "    os.mkdir(output_path_generated)\n",
    "    os.mkdir(output_path_noisy)\n",
    "    os.mkdir(output_path_clean)\n",
    "    os.mkdir(output_path_spectrograms)\n",
    "\n",
    "inference_with_mask = utils.get_version(model)\n",
    "\n",
    "for i in range(0, num_samples):\n",
    "  original_audio_paths = [os.path.join(dataset_path, dataset_subset, f\"{d_audios[i]['Index']}_{d_audios[i]['start_time']}.wav\")]\n",
    "  caption = d_audios[i]['caption']\n",
    "  shutil.copyfile(original_audio_paths[0], os.path.join(output_path_clean, f\"{i}.wav\"))\n",
    "  \n",
    "  # Get the embeddings of the original audio\n",
    "  _, original_latents = utils.get_original_latents(original_audio_paths, tango)\n",
    "\n",
    "  # Apply noise to the latents\n",
    "  noisy_latents = utils.apply_noise(original_latents, snr, verbose=False)\n",
    "\n",
    "  # Get the mask (time or mel) to apply to the latents (needed for inpainting)\n",
    "  if sum(time_mask_percentage)==0:\n",
    "      ipainting = True\n",
    "  mask = utils.get_mask(time_mask_percentage, mask_type=mask_type)\n",
    "\n",
    "  # Set the sigma_y value\n",
    "  sigma_y = utils.get_sigma_y(noisy_latents, mask)\n",
    "\n",
    "  with torch.no_grad():\n",
    "      \n",
    "      if model == 'ddnm+':\n",
    "          latents = inference_with_mask(tango.model, [caption], tango.scheduler,\n",
    "                                        num_steps=1000, guidance_scale=3, num_samples_per_prompt=1, disable_progress=False,\n",
    "                                        mask=mask, original= original_latents if mask is None else noisy_latents, \n",
    "                                        sigma_y=sigma_y, travel_length=0,\n",
    "                                        noisy_prompts=noisy_prompts, psnr_prompts=prompt_psnr)\n",
    "      elif model == 'repaint' and mask is not None:\n",
    "          latents = inference_with_mask(tango.model, [caption], tango.scheduler,\n",
    "                                  t_T=1000, guidance_scale=3, num_samples_per_prompt=1, disable_progress=False,\n",
    "                                  mask=mask, original=original_latents)\n",
    "      else:\n",
    "          print(\"Please select a proper method for the desired task\")\n",
    "\n",
    "\n",
    "      output_mels = tango.vae.decode_first_stage(latents)\n",
    "      waves = tango.vae.decode_to_waveform(output_mels)\n",
    "\n",
    "      utils.save_audio(output_path_generated, waves, id_file=f'{i}', descr=audio_description)\n",
    "\n",
    "      original_mels = tango.vae.decode_first_stage(original_latents)\n",
    "      \n",
    "          \n",
    "      noisy_mels = tango.vae.decode_first_stage(noisy_latents)\n",
    "      noisy_waves = tango.vae.decode_to_waveform(noisy_mels)\n",
    "          \n",
    "      if save_noisy_audio:\n",
    "          utils.save_audio(output_path_noisy, noisy_waves, id_file=f'{i}', descr=audio_description)\n",
    "      \n",
    "\n",
    "  mel_mask = utils.get_mask(time_mask_percentage, mask_type='mel')\n",
    "  utils.plot_spectrogram(original_mels, noisy_mels, output_mels, mel_mask=mel_mask, save_fig=True, id_s=i, output_path=output_path_spectrograms)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55b3de60-5992-42a9-8373-d95a15c5b07c",
   "metadata": {},
   "source": [
    "## Evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6cac7b24-9a2e-4dbd-a322-7f0004317003",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in /home/christian/.cache/torch/hub/harritaylor_torchvggish_master\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m             Results \u001b[0m \n",
      "         SNR        SDR      \n",
      "mean   -2.2181   -10.5631  \n",
      "std     0.0895     0.4068    \n",
      "min    -2.3077   -10.9699  \n",
      "max    -2.1286    -10.1563   \n",
      "\n",
      "FAD          4.2125               \n"
     ]
    }
   ],
   "source": [
    "import utils\n",
    "\n",
    "folders = ['generated', 'clean']\n",
    "\n",
    "# specify set of audios to exclude from the metrics computation\n",
    "exclude_audios = []\n",
    "\n",
    "snr_all, sdr_all, fad = utils.compute_metrics(os.path.join(output_path, name_experiment, folders[0]), \n",
    "                                              os.path.join(output_path, name_experiment, folders[1]), exclude=exclude_audios)\n",
    "utils.print_metrics_report(snr_all, sdr_all, fad)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "DiffInp",
   "language": "python",
   "name": "diffinp"
  },
  "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}