initial commit

app.py

@@ -0,0 +1,161 @@
+# Imports
+import gradio as gr
+import matplotlib.pyplot as plt
+import torch
+import torchaudio
+from torch import nn
+import pytorch_lightning as pl
+from ema_pytorch import EMA
+import yaml
+from audio_diffusion_pytorch import DiffusionModel, UNetV0, VDiffusion, VSampler
+
+
+# Load configs
+def load_configs(config_path):
+    with open(config_path, 'r') as file:
+        config = yaml.safe_load(file)
+    pl_configs = config['model']
+    model_configs = config['model']['model']
+    return pl_configs, model_configs
+
+# plot mel spectrogram
+def plot_mel_spectrogram(sample, sr):
+    transform = torchaudio.transforms.MelSpectrogram(
+        sample_rate=sr,
+        n_fft=1024,
+        hop_length=512,
+        n_mels=80,
+        center=True,
+        norm="slaney",
+    )
+
+    spectrogram = transform(torch.mean(sample, dim=0)) # downmix and cal spectrogram
+    spectrogram = torchaudio.functional.amplitude_to_DB(spectrogram, 1.0, 1e-10, 80.0)
+
+    # Plot the Mel spectrogram
+    fig = plt.figure(figsize=(7, 4))
+    plt.imshow(spectrogram, aspect='auto', origin='lower')
+    plt.colorbar(format='%+2.0f dB')
+    plt.xlabel('Frame')
+    plt.ylabel('Mel Bin')
+    plt.title('Mel Spectrogram')
+    plt.tight_layout()
+    
+    return fig
+
+# Define PyTorch Lightning model
+class Model(pl.LightningModule):
+    def __init__(
+        self,
+        lr: float,
+        lr_beta1: float,
+        lr_beta2: float,
+        lr_eps: float,
+        lr_weight_decay: float,
+        ema_beta: float,
+        ema_power: float,
+        model: nn.Module,
+    ):
+        super().__init__()
+        self.lr = lr
+        self.lr_beta1 = lr_beta1
+        self.lr_beta2 = lr_beta2
+        self.lr_eps = lr_eps
+        self.lr_weight_decay = lr_weight_decay
+        self.model = model
+        self.model_ema = EMA(self.model, beta=ema_beta, power=ema_power)
+
+# Instantiate model (must match model that was trained)
+def load_model(model_configs, pl_configs) -> nn.Module:
+    # Diffusion model
+    model = DiffusionModel(
+        net_t=UNetV0, # The model type used for diffusion (U-Net V0 in this case)
+        in_channels=model_configs['in_channels'], # U-Net: number of input/output (audio) channels
+        channels=model_configs['channels'], # U-Net: channels at each layer
+        factors=model_configs['factors'], # U-Net: downsampling and upsampling factors at each layer
+        items=model_configs['items'], # U-Net: number of repeating items at each layer
+        attentions=model_configs['attentions'], # U-Net: attention enabled/disabled at each layer
+        attention_heads=model_configs['attention_heads'], # U-Net: number of attention heads per attention item
+        attention_features=model_configs['attention_features'], # U-Net: number of attention features per attention item
+        diffusion_t=VDiffusion, # The diffusion method used
+        sampler_t=VSampler # The diffusion sampler used
+    )
+
+    # pl model
+    model = Model(
+        lr=pl_configs['lr'],
+        lr_beta1=pl_configs['lr_beta1'],
+        lr_beta2=pl_configs['lr_beta2'],
+        lr_eps=pl_configs['lr_eps'],
+        lr_weight_decay=pl_configs['lr_weight_decay'],
+        ema_beta=pl_configs['ema_beta'],
+        ema_power=pl_configs['ema_power'],
+        model=model
+    )
+
+    return model
+
+# Assign to GPU
+def assign_to_gpu(model):
+    if torch.cuda.is_available():
+        model = model.to('cuda')
+        print(f"Device: {model.device}")
+    return model
+
+# Load model checkpoint
+def load_checkpoint(model, ckpt_path) -> None:
+    checkpoint = torch.load(ckpt_path, map_location='cpu')['state_dict']
+    model.load_state_dict(checkpoint) # should output "<All keys matched successfully>"
+
+
+# Generate Samples
+def generate_samples(model_name, num_samples, num_steps, duration=32768):
+    # load_checkpoint
+    ckpt_path = models[model_name]
+    load_checkpoint(model, ckpt_path)
+
+    with torch.no_grad():
+        all_samples = torch.zeros(2, 0) # initialize all samples
+        for i in range(num_samples):
+            noise = torch.randn((1, 2, int(duration)), device=model.device) # [batch_size, in_channels, length]
+            generated_sample = model.model_ema.ema_model.sample(noise, num_steps=num_steps).squeeze(0).cpu() # Suggested num_steps 10-100
+            
+            # concatenate all samples:
+            all_samples = torch.concat((all_samples, generated_sample), dim=1)
+            
+            torch.cuda.empty_cache()
+    
+    fig = plot_mel_spectrogram(all_samples, sr)
+    plt.title(f"{model_name} Mel Spectrogram")
+
+    return (sr, all_samples.cpu().detach().numpy().T), fig # (sample rate, audio), plot
+
+# load model & configs
+sr = 44100 # sampling rate
+config_path = "saved_models/config.yaml" # config path
+pl_configs, model_configs = load_configs(config_path)
+model = load_model(model_configs, pl_configs)
+model = assign_to_gpu(model)
+
+models = {
+    "Kicks": "saved_models/kicks/kicks_v7.ckpt",
+    "Snares": "saved_models/snares/snares_v0.ckpt",
+    "Hi-hats": "saved_models/hihats/hihats_v2.ckpt",
+    "Percussion": "saved_models/percussion/percussion_v0.ckpt"
+}
+
+demo = gr.Interface(
+    generate_samples,
+    inputs=[
+        gr.Dropdown(choices=list(models.keys()), value=list(models.keys())[0], label="Model"),
+        gr.Slider(1, 25, step=1, label="Number of Samples to Generate", value=1),
+        gr.Slider(1, 100, step=1, label="Number of Diffusion Steps", value=10)
+    ],
+    outputs=[
+        gr.Audio(label="Generated Audio Sample"),
+        gr.Plot(label="Generated Audio Spectrogram")
+    ]
+)
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,34 @@
+torch>=2.0
+torchaudio>=2.0
+pytorch-lightning==1.7.7
+python-dotenv
+hydra-core
+hydra-colorlog
+wandb
+auraloss
+yt-dlp
+datasets
+pyloudnorm
+einops
+omegaconf
+rich
+plotly
+librosa
+transformers
+eng-to-ipa
+ema-pytorch
+py7zr
+notebook
+matplotlib
+ipykernel
+gradio
+
+# k-diffusion
+# v-diffusion-pytorch
+
+audio-diffusion-pytorch==0.1.3
+audio-encoders-pytorch
+audio-data-pytorch
+quantizer-pytorch
+difformer-pytorch
+a-transformers-pytorch

saved_models/config.yaml

@@ -0,0 +1,131 @@
+seed: 12345
+train: true
+ignore_warnings: true
+print_config: false
+work_dir: ${hydra:runtime.cwd}
+logs_dir: ${work_dir}${oc.env:DIR_LOGS}
+data_dir: ${work_dir}${oc.env:DIR_DATA}
+ckpt_dir: ${logs_dir}/ckpts/${now:%Y-%m-%d-%H-%M-%S}
+module: main.module_base
+batch_size: 1
+accumulate_grad_batches: 32
+num_workers: 8
+sampling_rate: 44100
+length: 32768
+channels: 2
+log_every_n_steps: 1000
+model:
+  _target_: ${module}.Model
+  lr: 0.0001
+  lr_beta1: 0.95
+  lr_beta2: 0.999
+  lr_eps: 1.0e-06
+  lr_weight_decay: 0.001
+  ema_beta: 0.995
+  ema_power: 0.7
+  model:
+    _target_: main.DiffusionModel
+    net_t:
+      _target_: ${module}.UNetT
+    in_channels: 2
+    channels:
+    - 32
+    - 32
+    - 64
+    - 64
+    - 128
+    - 128
+    - 256
+    - 256
+    factors:
+    - 1
+    - 2
+    - 2
+    - 2
+    - 2
+    - 2
+    - 2
+    - 2
+    items:
+    - 2
+    - 2
+    - 2
+    - 2
+    - 2
+    - 2
+    - 4
+    - 4
+    attentions:
+    - 0
+    - 0
+    - 0
+    - 0
+    - 0
+    - 1
+    - 1
+    - 1
+    attention_heads: 8
+    attention_features: 64
+datamodule:
+  _target_: main.module_base.Datamodule
+  dataset:
+    _target_: audio_data_pytorch.WAVDataset
+    path: ./data/wav_dataset/kicks
+    recursive: true
+    sample_rate: ${sampling_rate}
+    transforms:
+      _target_: audio_data_pytorch.AllTransform
+      crop_size: ${length}
+      stereo: true
+      source_rate: ${sampling_rate}
+      target_rate: ${sampling_rate}
+      loudness: -20
+  val_split: 0.05
+  batch_size: ${batch_size}
+  num_workers: ${num_workers}
+  pin_memory: true
+callbacks:
+  rich_progress_bar:
+    _target_: pytorch_lightning.callbacks.RichProgressBar
+  model_checkpoint:
+    _target_: pytorch_lightning.callbacks.ModelCheckpoint
+    monitor: valid_loss
+    save_top_k: 1
+    save_last: true
+    mode: min
+    verbose: false
+    dirpath: ${logs_dir}/ckpts/${now:%Y-%m-%d-%H-%M-%S}
+    filename: '{epoch:02d}-{valid_loss:.3f}'
+  model_summary:
+    _target_: pytorch_lightning.callbacks.RichModelSummary
+    max_depth: 2
+  audio_samples_logger:
+    _target_: main.module_base.SampleLogger
+    num_items: 4
+    channels: ${channels}
+    sampling_rate: ${sampling_rate}
+    length: ${length}
+    sampling_steps:
+    - 50
+    use_ema_model: true
+loggers:
+  wandb:
+    _target_: pytorch_lightning.loggers.wandb.WandbLogger
+    project: ${oc.env:WANDB_PROJECT}
+    entity: ${oc.env:WANDB_ENTITY}
+    name: kicks_v7
+    job_type: train
+    group: ''
+    save_dir: ${logs_dir}
+trainer:
+  _target_: pytorch_lightning.Trainer
+  gpus: 1
+  precision: 16
+  accelerator: gpu
+  min_epochs: 0
+  max_epochs: -1
+  enable_model_summary: false
+  log_every_n_steps: 1
+  check_val_every_n_epoch: null
+  val_check_interval: ${log_every_n_steps}
+  accumulate_grad_batches: ${accumulate_grad_batches}

saved_models/hihats/hihats_v2.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fc7245d3d5617bb3a76dcc8534d9cee25030c3986fa80502f19ec3506a68d05c
+size 509086593

saved_models/kicks/kicks_v7.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f3511269e10edc889cfd50393fd5228cdfb069185afc9d92263cef548a18482
+size 509086593

saved_models/percussion/percussion_v0.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f8fe5dc0295738995cb74892a7d70a074abdfd2c7e887951a2bc9814ec9acfaf
+size 509086593

saved_models/snares/snares_v0.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d2f906655666200635267c3a92ff87631f4bb4ef94bf087cfee3e2611da9b30b
+size 509086593