tidy

.gitignore

@@ -1,13 +1,11 @@
 .vscode
 __pycache__
 .ipynb_checkpoints
-data*
-ddpm-ema-audio-*
+data
+models
 flagged
 build
 audiodiffusion.egg-info
 lightning_logs
 taming
 checkpoints
-vae_model
-latent-audio-diffusion-*

README.md

@@ -56,7 +56,7 @@ python scripts/audio_to_images.py \
   --resolution 64 \
   --hop_length 1024 \
   --input_dir path-to-audio-files \
-  --output_dir data-test
+  --output_dir path-to-output-data
 ```
 #### Generate dataset of 256x256 Mel spectrograms and push to hub (you will need to be authenticated with `huggingface-cli login`).
 
@@ -64,7 +64,7 @@ python scripts/audio_to_images.py \
 python scripts/audio_to_images.py \
   --resolution 256 \
   --input_dir path-to-audio-files \
-  --output_dir data-256 \
+  --output_dir data/audio-diffusion-256 \
   --push_to_hub teticio/audio-diffusion-256
 ```
 ## Train model
@@ -72,10 +72,10 @@ python scripts/audio_to_images.py \
 ```bash
 accelerate launch --config_file config/accelerate_local.yaml \
   scripts/train_unconditional.py \
-  --dataset_name data-64 \
+  --dataset_name data/audio-diffusion-64 \
   --resolution 64 \
   --hop_length 1024 \
-  --output_dir ddpm-ema-audio-64 \
+  --output_dir models/ddpm-ema-audio-64 \
   --train_batch_size 16 \
   --num_epochs 100 \
   --gradient_accumulation_steps 1 \
@@ -89,7 +89,7 @@ accelerate launch --config_file config/accelerate_local.yaml \
   scripts/train_unconditional.py \
   --dataset_name teticio/audio-diffusion-256 \
   --resolution 256 \
-  --output_dir audio-diffusion-256 \
+  --output_dir models/audio-diffusion-256 \
   --num_epochs 100 \
   --train_batch_size 2 \
   --eval_batch_size 2 \
@@ -107,7 +107,7 @@ accelerate launch --config_file config/accelerate_sagemaker.yaml \
   scripts/train_unconditional.py \
   --dataset_name teticio/audio-diffusion-256 \
   --resolution 256 \
-  --output_dir ddpm-ema-audio-256 \
+  --output_dir models/ddpm-ema-audio-256 \
   --train_batch_size 16 \
   --num_epochs 100 \
   --gradient_accumulation_steps 1 \

notebooks/test_vae.ipynb

@@ -1,5 +1,17 @@
 {
  "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3c8663ed",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import sys\n",
+    "sys.path.insert(0, os.path.dirname(os.path.abspath(\"\")))"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -12,7 +24,9 @@
     "import numpy as np\n",
     "from PIL import Image\n",
     "from datasets import load_dataset\n",
-    "from diffusers import AutoencoderKL"
+    "from IPython.display import Audio\n",
+    "from diffusers import AutoencoderKL\n",
+    "from audiodiffusion.mel import Mel"
    ]
   },
   {
@@ -22,7 +36,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "vae = AutoencoderKL.from_pretrained('../vae_model')"
+    "mel = Mel()\n",
+    "vae = AutoencoderKL.from_pretrained('../models/autoencoder-kl')"
    ]
   },
   {
@@ -42,7 +57,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ds = load_dataset('teticio/audio-diffusion-256')"
+    "ds = load_dataset('teticio/audio-diffusion-breaks-256')"
    ]
   },
   {
@@ -53,7 +68,8 @@
    "outputs": [],
    "source": [
     "image = random.choice(ds['train'])['image']\n",
-    "image"
+    "display(image)\n",
+    "Audio(data=mel.image_to_audio(image), rate=mel.get_sample_rate())"
    ]
   },
   {
@@ -84,7 +100,9 @@
     "output_image = (output_image + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w\n",
     "output_image = (output_image.detach().cpu().numpy() *\n",
     "                255).round().astype(\"uint8\").transpose(0, 2, 3, 1)[0]\n",
-    "Image.fromarray(output_image)"
+    "output_image = Image.fromarray(output_image).convert('L')\n",
+    "display(output_image)\n",
+    "Audio(data=mel.image_to_audio(output_image), rate=mel.get_sample_rate())"
    ]
   },
   {
@@ -100,7 +118,9 @@
     "output_image = (output_image + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w\n",
     "output_image = (output_image.detach().cpu().numpy() *\n",
     "                255).round().astype(\"uint8\").transpose(0, 2, 3, 1)[0]\n",
-    "Image.fromarray(output_image)"
+    "output_image = Image.fromarray(output_image).convert('L')\n",
+    "display(output_image)\n",
+    "Audio(data=mel.image_to_audio(output_image), rate=mel.get_sample_rate())"
    ]
   },
   {

scripts/train_unconditional.py

@@ -73,11 +73,12 @@ def main(args):
         )
 
     if args.scheduler == "ddpm":
-        noise_scheduler = DDPMScheduler(num_train_timesteps=1000,
-                                        tensor_format="pt")
+        noise_scheduler = DDPMScheduler(
+            num_train_timesteps=args.num_train_steps, tensor_format="pt")
     else:
-        noise_scheduler = DDIMScheduler(num_train_timesteps=1000,
-                                        tensor_format="pt")
+        noise_scheduler = DDIMScheduler(
+            num_train_timesteps=args.num_train_steps, tensor_format="pt")
+
     optimizer = torch.optim.AdamW(
         model.parameters(),
         lr=args.learning_rate,
@@ -305,7 +306,6 @@ if __name__ == "__main__":
     parser.add_argument("--overwrite_output_dir", type=bool, default=False)
     parser.add_argument("--cache_dir", type=str, default=None)
     parser.add_argument("--resolution", type=int, default=256)
-    parser.add_argument("--latent_resolution", type=int, default=64)
     parser.add_argument("--train_batch_size", type=int, default=16)
     parser.add_argument("--eval_batch_size", type=int, default=16)
     parser.add_argument("--num_epochs", type=int, default=100)
@@ -342,6 +342,8 @@ if __name__ == "__main__":
     parser.add_argument("--hop_length", type=int, default=512)
     parser.add_argument("--from_pretrained", type=str, default=None)
     parser.add_argument("--start_epoch", type=int, default=0)
+    parser.add_argument("--num_train_steps", type=int, default=1000)
+    parser.add_argument("--latent_resolution", type=int, default=64)
     parser.add_argument("--scheduler",
                         type=str,
                         default="ddpm",

scripts/train_vae.py

@@ -107,7 +107,7 @@ class ImageLogger(Callback):
 
 class HFModelCheckpoint(ModelCheckpoint):
 
-    def __init__(self, ldm_config, hf_checkpoint='vae_model', *args, **kwargs):
+    def __init__(self, ldm_config, hf_checkpoint='models/autoencoder-kl', *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.ldm_config = ldm_config
         self.hf_checkpoint = hf_checkpoint
@@ -130,7 +130,7 @@ if __name__ == "__main__":
                         default="config/ldm_autoencoder_kl.yaml")
     parser.add_argument("--ldm_checkpoint_dir",
                         type=str,
-                        default="checkpoints")
+                        default="models/ldm-autoencoder-kl")
     parser.add_argument("--hf_checkpoint_dir", type=str, default="vae_model")
     parser.add_argument("-r",
                         "--resume_from_checkpoint",