process_audio.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchaudio\n",
    "import librosa\n",
    "import soundfile\n",
    "from interactive_spectrogram_inpainting.vqvae.vqvae import VQVAE\n",
    "from interactive_spectrogram_inpainting.utils.misc import get_spectrograms_helper\n",
    "from IPython.display import display, Audio\n",
    "\n",
    "\n",
    "def extend_duration_like(this: torch.Tensor, reference: torch.Tensor, dim: int = -1) -> torch.Tensor:\n",
    "    padding_duration = reference.shape[dim] - this.shape[dim]\n",
    "    if padding_duration < 0:\n",
    "        return this.narrow(start=0, length=reference.shape[dim], dim=dim)\n",
    "    padding_shape = list(this.shape)\n",
    "    padding_shape[dim] = padding_duration\n",
    "    return torch.cat([this, torch.zeros_like(this.flatten()[0]).expand(padding_shape)], dim=dim)\n",
    "\n",
    "device = 'cuda'\n",
    "\n",
    "# downsampling factor 8 then 4\n",
    "model_large_model_parameters_json_path = '/home/theis/bachbach/code/vq-vae-2-pytorch/checkpoints/20200310-105659-1b55c7/model_parameters.json'\n",
    "model_large_model_weights_path = '/home/theis/bachbach/code/vq-vae-2-pytorch/checkpoints/20200310-105659-1b55c7/vqvae_nsynth_223.pt'\n",
    "\n",
    "model_large = VQVAE.from_parameters_and_weights(\n",
    "    model_large_model_parameters_json_path,\n",
    "    model_large_model_weights_path\n",
    ")\n",
    "\n",
    "# downsampling factor 16 then 2\n",
    "model_small_model_parameters_json_path = '/home/theis/code/vq-vae-2-pytorch/data/checkpoints/20200117-205357-5b1d22/model_parameters.json'\n",
    "model_small_model_weights_path = '/home/theis/code/vq-vae-2-pytorch/data/checkpoints/20200117-205357-5b1d22/vqvae_nsynth_471.pt'\n",
    "\n",
    "model_small = VQVAE.from_parameters_and_weights(\n",
    "    model_small_model_parameters_json_path,\n",
    "    model_small_model_weights_path\n",
    ")\n",
    "models = {\n",
    "    'small': {\n",
    "            'model': model_small,\n",
    "            'model_parameters_json_path': model_small_model_parameters_json_path,\n",
    "            'model_weights_path': model_small_model_weights_path\n",
    "        },\n",
    "    'large': {\n",
    "        'model': model_large,\n",
    "        'model_parameters_json_path': model_large_model_parameters_json_path,\n",
    "        'model_weights_path': model_large_model_weights_path\n",
    "        } \n",
    "}\n",
    "\n",
    "for model_name, model_dict in models.items():\n",
    "    model = models[model_name]['model']\n",
    "    models[model_name]['model'].adapt_quantized_durations = True\n",
    "    models[model_name]['model'] = model.to(device)\n",
    "    print(model_name + ' model uses GANSynth normalization:',\n",
    "          model.use_gansynth_normalization)\n",
    "\n",
    "assert models['small']['model'].adapt_quantized_durations\n",
    "\n",
    "VQVAE_TRAINING_PARAMETERS_PATH = '/home/theis/code/vq-vae-2-pytorch/data/checkpoints/20200117-205357-5b1d22/command_line_parameters.json'\n",
    "with open(VQVAE_TRAINING_PARAMETERS_PATH, 'r') as f:\n",
    "        vqvae_training_parameters = json.load(f)\n",
    "spectrograms_helper = get_spectrograms_helper(\n",
    "        device=device, **vqvae_training_parameters)\n",
    "spectrograms_helper.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def process(audio: torch.Tensor, sr: int, model: VQVAE,\n",
    "            device: str = 'cuda') -> torch.Tensor:\n",
    "    if audio.ndim == 3:\n",
    "        audio = audio.squeeze(0)\n",
    "    if sr != 16000:\n",
    "        resampler_to_model = torchaudio.transforms.Resample(\n",
    "            orig_freq=sr, new_freq=16000)\n",
    "        resampler_from_model = torchaudio.transforms.Resample(\n",
    "            orig_freq=16000, new_freq=sr)\n",
    "    else:\n",
    "        resampler_to_model = resampler_from_model = lambda x: x\n",
    "\n",
    "    audio_original_sr = audio\n",
    "    audio_original_sr = audio_original_sr.to(device)\n",
    "    audio_resampled = resampler_to_model(audio_original_sr)\n",
    "    spec_resampled = spectrograms_helper.to_spectrogram(audio_resampled)\n",
    "    spec_reconstruction_resampled = model(spec_resampled)[0]\n",
    "    audio_reconstruction_resampled = spectrograms_helper.to_audio(\n",
    "        spec_reconstruction_resampled)\n",
    "    audio_reconstruction_original_sr = resampler_from_model(\n",
    "        audio_reconstruction_resampled)\n",
    "    audio_reconstruction_original_sr_same_duration = extend_duration_like(\n",
    "        audio_reconstruction_original_sr, audio_original_sr)\n",
    "    return audio_reconstruction_original_sr_same_duration, sr\n",
    "\n",
    "def process_all(audio: torch.Tensor, sr: int) -> torch.Tensor:\n",
    "    return {model_name: process(audio, sr, model_dict['model'])\n",
    "            for model_name, model_dict in models.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_player(audio: torch.Tensor, rate: int = 16000):\n",
    "    try:\n",
    "        audio = audio.cpu().numpy()\n",
    "    except AttributeError:\n",
    "        pass\n",
    "    display(Audio(data=audio[0], rate=rate,\n",
    "                  normalize=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch import nn\n",
    "from glob import glob\n",
    "import os\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm\n",
    "import shutil\n",
    "from typing import Optional\n",
    "\n",
    "class AudioSet(Dataset):\n",
    "    def __init__(self, directory: Path, extension: str = 'wav',\n",
    "                 fs_hz: int = 16000, min_duration_s: Optional[float] = None):\n",
    "        self.directory = directory\n",
    "        self.extension = extension\n",
    "        self.filenames = glob(str(self.directory / f'**/*.{self.extension}'), recursive=False)\n",
    "        self.filenames = [filename for filename in self.filenames\n",
    "                         if os.path.isfile(filename)]\n",
    "        self.fs_hz = fs_hz\n",
    "        self.min_duration_s = min_duration_s\n",
    "        self.min_duration_n = None\n",
    "        if self.min_duration_s is not None:\n",
    "            self.min_duration_n = self.min_duration_s * self.fs_hz\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        filename = self.filenames[index]\n",
    "        \n",
    "        audio_original_sr, original_sr = librosa.load(filename, sr=None, mono=False)\n",
    "        if audio_original_sr.ndim == 1:\n",
    "            audio_original_sr = audio_original_sr[None, :]\n",
    "\n",
    "        audio_stereo = librosa.resample(audio_original_sr, original_sr, self.fs_hz)\n",
    "        audio_tensor = torch.as_tensor(audio_stereo)\n",
    "        \n",
    "        duration_n = audio_tensor.shape[-1] \n",
    "        if (self.min_duration_n is not None and duration_n  < self.min_duration_n):\n",
    "            padding_duration = self.min_duration_n - duration_n\n",
    "            padding_shape = list(audio_tensor.shape)\n",
    "            padding_shape[-1] = padding_duration\n",
    "            padding = torch.zeros(padding_shape, dtype=audio_tensor.dtype,\n",
    "                                  layout=audio_tensor.layout,\n",
    "                                  device=audio_tensor.device,)\n",
    "            audio_tensor = torch.cat([audio_tensor,\n",
    "                                      padding],\n",
    "                                     dim=-1)\n",
    "        return audio_tensor, filename\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.filenames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root_dir = Path('/home/theis/sounds_for_processing')\n",
    "sample_pack_names = [\n",
    "    'Gentleman'\n",
    "]\n",
    "output_dir_name = 'VQ-VAE-PROCESSED-stereo/'\n",
    "\n",
    "for model_size, model_dict in models.items():\n",
    "    parallel_model = nn.DataParallel(model_dict['model'])\n",
    "\n",
    "    root_output_dir = root_dir / output_dir_name / (model_size + '_model')\n",
    "    root_output_dir.mkdir(exist_ok=True, parents=True)\n",
    "\n",
    "    shutil.copy2(model_dict['model_parameters_json_path'], root_output_dir)\n",
    "    shutil.copy2(model_dict['model_weights_path'], root_output_dir)\n",
    "\n",
    "    for sample_pack_name in sample_pack_names:\n",
    "        samples_dir = root_dir / sample_pack_name\n",
    "        print(samples_dir)\n",
    "        output_dir = root_output_dir / (sample_pack_name + '/')\n",
    "        output_dir.mkdir(parents=True, exist_ok=True)\n",
    "        print(output_dir)\n",
    "        audio_set = AudioSet(samples_dir, min_duration_s=1)\n",
    "        audio_loader = DataLoader(audio_set, num_workers=8, batch_size=1, shuffle=False)\n",
    "\n",
    "        DEVICE = 'cuda'\n",
    "        parallel_model.to(DEVICE)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for i, (samples, filenames) in enumerate(tqdm(audio_loader)):\n",
    "                sample = samples[0]\n",
    "                filename = filenames[0]\n",
    "                output_file_path = output_dir / Path(filename).relative_to(samples_dir)\n",
    "                output_file_path.parent.mkdir(exist_ok=True, parents=True)\n",
    "                processed_channels, processed_sr = process(sample, 16000, parallel_model, DEVICE)\n",
    "\n",
    "                with soundfile.SoundFile(filename) as sf:\n",
    "                    soundfile.write(\n",
    "                        file=str(output_file_path), data=processed_channels.cpu().numpy().T,\n",
    "                        samplerate=processed_sr, format=sf.format, subtype=sf.subtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with SoundFile(sample_path) as sf:\n",
    "    print(sf)\n",
    "    \n",
    "    # get the audio in channels-first format\n",
    "    audio_stereo = sf.read(dtype='float32').T\n",
    "    # resample to the model's sample-rate\n",
    "    audio_stereo_model_sr = librosa.resample(audio_stereo, sf.samplerate, 16000)\n",
    "    \n",
    "    audio_tensor_model_sr = torch.as_tensor(audio_stereo_model_sr).sum(0, keepdim=True)\n",
    "    make_player(audio_model_sr, rate=16000)\n",
    "    \n",
    "    for model_name in models:\n",
    "        print(model_name)\n",
    "        model = models[model_name]['model']\n",
    "        processed_tensor, processed_sr = process(audio_tensor_model_sr, 16000,\n",
    "                                                 model=model)\n",
    "        make_player(processed_tensor, rate=processed_sr)\n",
    "        \n",
    "        # resample to the original sample-rate \n",
    "        processed_original_sr = librosa.resample(processed_tensor.cpu().numpy(),\n",
    "                                                 processed_sr, sf.samplerate)\n",
    "        \n",
    "        processed_soundfile = processed_original_sr.T\n",
    "        write(file=f'./soundfile_output-{model_name}.wav', data=processed_soundfile,\n",
    "              samplerate=sf.samplerate, format=sf.format, subtype=sf.subtype)\n",
    "        \n",
    "        processed_soundfile_normalized = processed_soundfile / processed_soundfile.max()\n",
    "        write(file=f'./soundfile_output_normalized-{model_name}.wav',\n",
    "              data=processed_soundfile_normalized,\n",
    "              samplerate=sf.samplerate, format=sf.format, subtype=sf.subtype)"
   ]
  }
 ],
 "metadata": {
  "file_extension": ".py",
  "kernelspec": {
   "display_name": "Python 3.7.6 64-bit ('vqvae-synth': conda)",
   "language": "python",
   "name": "python37664bitvqvaesynthcondaf340fd92c0314539ac2eeac103b85658"
  },
  "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.7.9"
  },
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython3",
  "version": 3
 },
 "nbformat": 4,
 "nbformat_minor": 4
}