add out-painting feature

README.md

@@ -15,8 +15,12 @@ license: gpl-3.0
 
 ---
 
-**UPDATE**: 
+**UPDATES**: 
 
+4/10/2022
+It is now possible to mask parts of the input audio during generation which means you can stitch several samples together (think "out-painting").
+
+27/9/2022
 You can now generate an audio based on a previous one. You can use this to generate variations of the same audio or even to "remix" a track (via a sort of "style transfer"). You can find examples of how to do this in the [`test_model.ipynb`](https://colab.research.google.com/github/teticio/audio-diffusion/blob/master/notebooks/test_model.ipynb) notebook.
 
 ---

audiodiffusion/__init__.py

@@ -9,7 +9,7 @@ from diffusers import DDPMPipeline, DDPMScheduler
 
 from .mel import Mel
 
-VERSION = "1.1.3"
+VERSION = "1.1.4"
 
 
 class AudioDiffusion:
@@ -61,7 +61,9 @@ class AudioDiffusion:
         slice: int = 0,
         start_step: int = 0,
         steps: int = None,
-        generator: torch.Generator = None
+        generator: torch.Generator = None,
+        mask_start_secs: float = 0,
+        mask_end_secs: float = 0
     ) -> Tuple[Image.Image, Tuple[int, np.ndarray]]:
         """Generate random mel spectrogram from audio input and convert to audio.
 
@@ -72,6 +74,8 @@ class AudioDiffusion:
             start_step (int): step to start from
             steps (int): number of de-noising steps to perform (defaults to num_train_timesteps)
             generator (torch.Generator): random number generator or None
+            mask_start_secs (float): number of seconds of audio to mask (not generate) at start
+            mask_end_secs (float): number of seconds of audio to mask (not generate) at end
 
         Returns:
             PIL Image: mel spectrogram
@@ -84,31 +88,51 @@ class AudioDiffusion:
             steps = self.ddpm.scheduler.num_train_timesteps
         scheduler = DDPMScheduler(num_train_timesteps=steps)
         scheduler.set_timesteps(steps)
-        images = torch.randn(
+        mask = None
+        images = noise = torch.randn(
             (1, self.ddpm.unet.in_channels, self.ddpm.unet.sample_size,
              self.ddpm.unet.sample_size),
             generator=generator,
         )
+
         if audio_file is not None or raw_audio is not None:
             self.mel.load_audio(audio_file, raw_audio)
             input_image = self.mel.audio_slice_to_image(slice)
             input_image = np.frombuffer(input_image.tobytes(),
                                         dtype="uint8").reshape(
-                                            (input_image.width,
-                                             input_image.height))
+                                            (input_image.height,
+                                             input_image.width))
             input_image = ((input_image / 255) * 2 - 1)
+
             if start_step > 0:
-                images[0][0] = scheduler.add_noise(
+                images[0, 0] = scheduler.add_noise(
                     torch.tensor(input_image[np.newaxis, np.newaxis, :]),
-                    images, steps - start_step)
+                    noise, steps - start_step)
+
+            mask_start = int(mask_start_secs * self.mel.get_sample_rate() /
+                             self.mel.hop_length)
+            mask_end = int(mask_end_secs * self.mel.get_sample_rate() /
+                           self.mel.hop_length)
+            mask = scheduler.add_noise(
+                torch.tensor(input_image[np.newaxis, np.newaxis, :]), noise,
+                scheduler.timesteps[start_step:])
 
         images = images.to(self.ddpm.device)
-        for t in self.progress_bar(scheduler.timesteps[start_step:]):
+        for step, t in enumerate(
+                self.progress_bar(scheduler.timesteps[start_step:])):
             model_output = self.ddpm.unet(images, t)['sample']
             images = scheduler.step(model_output,
                                     t,
                                     images,
                                     generator=generator)['prev_sample']
+
+            if mask is not None:
+                if mask_start > 0:
+                    images[0, 0, :, :mask_start] = mask[step,
+                                                        0, :, :mask_start]
+                if mask_end > 0:
+                    images[0, 0, :, -mask_end:] = mask[step, 0, :, -mask_end:]
+
         images = (images / 2 + 0.5).clamp(0, 1)
         images = images.cpu().permute(0, 2, 3, 1).numpy()
 

audiodiffusion/mel.py

@@ -37,7 +37,7 @@ class Mel:
         self.slice_size = self.x_res * self.hop_length - 1
         self.fmax = self.sr / 2
         self.top_db = top_db
-        self.y = None
+        self.audio = None
 
     def load_audio(self, audio_file: str = None, raw_audio: np.ndarray = None):
         """Load audio.
@@ -47,11 +47,16 @@ class Mel:
             raw_audio (np.ndarray): audio as numpy array
         """
         if audio_file is not None:
-            self.y, _ = librosa.load(
-                audio_file,
-                mono=True)
+            self.audio, _ = librosa.load(audio_file, mono=True)
         else:
-            self.y = raw_audio
+            self.audio = raw_audio
+
+        # Pad with silence if necessary.
+        if len(self.audio) < self.x_res * self.hop_length:
+            self.audio = np.concatenate([
+                self.audio,
+                np.zeros((self.x_res * self.hop_length - len(self.audio), ))
+            ])
 
     def get_number_of_slices(self) -> int:
         """Get number of slices in audio.
@@ -59,7 +64,7 @@ class Mel:
         Returns:
             int: number of spectograms audio can be sliced into
         """
-        return len(self.y) // self.slice_size
+        return len(self.audio) // self.slice_size
 
     def get_audio_slice(self, slice: int = 0) -> np.ndarray:
         """Get slice of audio.
@@ -70,7 +75,8 @@ class Mel:
         Returns:
             np.ndarray: audio as numpy array
         """
-        return self.y[self.slice_size * slice:self.slice_size * (slice + 1)]
+        return self.audio[self.slice_size * slice:self.slice_size *
+                          (slice + 1)]
 
     def get_sample_rate(self) -> int:
         """Get sample rate:

notebooks/test_model.ipynb

@@ -10,7 +10,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": null,
    "id": "6c7800a6",
    "metadata": {},
    "outputs": [],
@@ -27,7 +27,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": null,
    "id": "b447e2c4",
    "metadata": {},
    "outputs": [],
@@ -39,7 +39,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 41,
+   "execution_count": null,
    "id": "c2fc0e7a",
    "metadata": {},
    "outputs": [],
@@ -63,7 +63,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": null,
    "id": "97f24046",
    "metadata": {},
    "outputs": [],
@@ -79,7 +79,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": null,
    "id": "a3d45c36",
    "metadata": {},
    "outputs": [],
@@ -169,6 +169,39 @@
     "display(Audio(track, rate=sample_rate))"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "11f91ad3",
+   "metadata": {},
+   "source": [
+    "### Generate continuations (\"out-painting\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "756d7af5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "overlap_secs = 2  #@param {type:\"integer\"}\n",
+    "start_step = 0  #@param {type:\"slider\", min:0, max:1000, step:10}\n",
+    "overlap_samples = overlap_secs * sample_rate\n",
+    "track = audio\n",
+    "for variation in range(12):\n",
+    "    image2, (\n",
+    "        sample_rate, audio2\n",
+    "    ) = audio_diffusion.generate_spectrogram_and_audio_from_audio(\n",
+    "        raw_audio=audio[-overlap_samples:],\n",
+    "        start_step=start_step,\n",
+    "        mask_start_secs=overlap_secs)\n",
+    "    display(image2)\n",
+    "    display(Audio(audio2, rate=sample_rate))\n",
+    "    track = np.concatenate([track, audio2[overlap_samples:]])\n",
+    "    audio = audio2\n",
+    "display(Audio(track, rate=sample_rate))"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "b6434d3f",
@@ -182,12 +215,12 @@
    "id": "0da030b2",
    "metadata": {},
    "source": [
-    "Alternatively, you can start from another audio altogether, resulting in a kind of style transfer."
+    "Alternatively, you can start from another audio altogether, resulting in a kind of style transfer. Maintaining the same seed during generation fixes the style, while masking helps stitch consecutive segments together more smoothly."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 50,
+   "execution_count": null,
    "id": "fc620a80",
    "metadata": {},
    "outputs": [],
@@ -207,41 +240,31 @@
    "metadata": {
     "scrolled": true
    },
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6e741e6bd196458fa38f86197bd16378",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/500 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
+   "outputs": [],
    "source": [
-    "start_steps = 500  #@param {type:\"slider\", min:0, max:1000, step:10}\n",
-    "audio_diffusion.mel.load_audio(audio_file)\n",
-    "track = np.array([])\n",
+    "start_step = 500  #@param {type:\"slider\", min:0, max:1000, step:10}\n",
+    "overlap_secs = 1  #@param {type:\"integer\"}\n",
+    "overlap_samples = overlap_secs * sample_rate\n",
+    "mel.load_audio(audio_file)\n",
+    "slice_size = audio_diffusion.mel.x_res * audio_diffusion.mel.hop_length\n",
+    "stride = slice_size - overlap_samples\n",
     "generator = torch.Generator()\n",
     "seed = generator.seed()\n",
-    "for slice in range(audio_diffusion.mel.get_number_of_slices()):\n",
+    "track = np.array([])\n",
+    "for sample in range(len(mel.audio) // stride):\n",
     "    generator.manual_seed(seed)\n",
-    "    audio = audio_diffusion.mel.get_audio_slice(slice)\n",
-    "    _, (\n",
-    "        sample_rate, audio2\n",
-    "    ) = audio_diffusion.generate_spectrogram_and_audio_from_audio(\n",
-    "        audio_file=audio_file,\n",
-    "        slice=slice,\n",
-    "        start_step=start_steps,\n",
-    "        generator=generator)\n",
+    "    audio = mel.audio[sample * stride:sample * stride + slice_size]\n",
+    "    if len(track) > 0:\n",
+    "        audio[:overlap_samples] = audio2[-overlap_samples:]\n",
+    "    _, (sample_rate,\n",
+    "        audio2) = audio_diffusion.generate_spectrogram_and_audio_from_audio(\n",
+    "            raw_audio=audio,\n",
+    "            start_step=start_step,\n",
+    "            generator=generator,\n",
+    "            mask_start_secs=1 if len(track) > 0 else 0)\n",
     "    display(Audio(audio, rate=sample_rate))\n",
     "    display(Audio(audio2, rate=sample_rate))\n",
-    "    track = np.concatenate([track, audio2])"
+    "    track = np.concatenate([track, audio2[overlap_samples:]])"
    ]
   },
   {
@@ -307,7 +330,17 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c59bcc0f",
+   "id": "df112a72",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "len(audio) / mel.hop_length"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ad467206",
    "metadata": {},
    "outputs": [],
    "source": []
@@ -334,7 +367,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.4"
+   "version": "3.10.6"
   },
   "toc": {
    "base_numbering": 1,