File size: 7,262 Bytes
f8a240c
 
3e8b723
 
 
 
 
 
 
 
 
 
 
 
 
8aa7c27
3e8b723
8aa7c27
3e8b723
 
8aa7c27
3e8b723
 
 
 
13aa297
3e8b723
13aa297
8aa7c27
 
 
 
3e8b723
 
 
 
 
13aa297
 
 
3e8b723
4f552a8
3e8b723
 
 
 
 
 
13aa297
3e8b723
 
13aa297
3e8b723
 
 
 
 
 
 
 
 
 
5bc60f9
 
 
 
 
3e8b723
d76bdef
327bccf
5bc60f9
 
d76bdef
8aa7c27
 
d76bdef
 
 
 
 
327bccf
 
 
 
5bc60f9
 
 
327bccf
3e8b723
d76bdef
 
 
 
3e8b723
d76bdef
3e8b723
 
 
d76bdef
3e8b723
d76bdef
327bccf
 
c5dcd04
3e8b723
 
 
 
327bccf
3e8b723
8aa7c27
 
 
 
 
 
 
 
 
d122744
8aa7c27
 
 
 
 
529c646
 
8aa7c27
 
 
 
3e8b723
 
 
8aa7c27
 
 
 
 
 
 
 
001a426
d122744
 
 
8aa7c27
 
 
 
 
 
 
 
13aa297
5bc60f9
 
529c646
a8d2b62
3e8b723
 
8aa7c27
13aa297
 
 
 
 
 
3e8b723
 
8aa7c27
3e8b723
8aa7c27
 
b3e97c5
8aa7c27
 
529c646
5bc60f9
 
 
8aa7c27
 
 
 
 
 
 
3e8b723
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
# based on https://github.com/CompVis/stable-diffusion/blob/main/main.py

import os
import argparse

import torch
import torchvision
import numpy as np
from PIL import Image
import pytorch_lightning as pl
from omegaconf import OmegaConf
from librosa.util import normalize
from ldm.util import instantiate_from_config
from pytorch_lightning.trainer import Trainer
from torch.utils.data import DataLoader, Dataset
from datasets import load_from_disk, load_dataset
from pytorch_lightning.callbacks import Callback, ModelCheckpoint
from pytorch_lightning.utilities.distributed import rank_zero_only

from audiodiffusion.mel import Mel
from audiodiffusion.utils import convert_ldm_to_hf_vae


class AudioDiffusion(Dataset):

    def __init__(self, model_id, channels=3):
        super().__init__()
        self.channels = channels
        if os.path.exists(model_id):
            self.hf_dataset = load_from_disk(model_id)['train']
        else:
            self.hf_dataset = load_dataset(model_id)['train']

    def __len__(self):
        return len(self.hf_dataset)

    def __getitem__(self, idx):
        image = self.hf_dataset[idx]['image']
        if self.channels == 3:
            image = image.convert('RGB')
        image = np.frombuffer(image.tobytes(), dtype="uint8").reshape(
            (image.height, image.width, self.channels))
        image = ((image / 255) * 2 - 1)
        return {'image': image}


class AudioDiffusionDataModule(pl.LightningDataModule):

    def __init__(self, model_id, batch_size, channels):
        super().__init__()
        self.batch_size = batch_size
        self.dataset = AudioDiffusion(model_id=model_id, channels=channels)
        self.num_workers = 1

    def train_dataloader(self):
        return DataLoader(self.dataset,
                          batch_size=self.batch_size,
                          num_workers=self.num_workers)


class ImageLogger(Callback):

    def __init__(self,
                 every=1000,
                 hop_length=512,
                 sample_rate=22050,
                 n_fft=2048):
        super().__init__()
        self.every = every
        self.hop_length = hop_length
        self.sample_rate = sample_rate
        self.n_fft = n_fft

    @rank_zero_only
    def log_images_and_audios(self, pl_module, batch):
        pl_module.eval()
        with torch.no_grad():
            images = pl_module.log_images(batch, split='train')
        pl_module.train()

        image_shape = next(iter(images.values())).shape
        channels = image_shape[1]
        mel = Mel(x_res=image_shape[2],
                  y_res=image_shape[3],
                  hop_length=self.hop_length,
                  sample_rate=self.sample_rate,
                  n_fft=self.n_fft)

        for k in images:
            images[k] = images[k].detach().cpu()
            images[k] = torch.clamp(images[k], -1., 1.)
            images[k] = (images[k] + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w
            grid = torchvision.utils.make_grid(images[k])

            tag = f"train/{k}"
            pl_module.logger.experiment.add_image(
                tag, grid, global_step=pl_module.global_step)

            images[k] = (images[k].numpy() *
                         255).round().astype("uint8").transpose(0, 2, 3, 1)
            for _, image in enumerate(images[k]):
                audio = mel.image_to_audio(
                    Image.fromarray(image, mode='RGB').convert('L')
                    if channels == 3 else Image.fromarray(image[:, :, 0]))
                pl_module.logger.experiment.add_audio(
                    tag + f"/{_}",
                    normalize(audio),
                    global_step=pl_module.global_step,
                    sample_rate=mel.get_sample_rate())

    def on_train_batch_end(self, trainer, pl_module, outputs, batch,
                           batch_idx):
        if (batch_idx + 1) % self.every != 0:
            return
        self.log_images_and_audios(pl_module, batch)


class HFModelCheckpoint(ModelCheckpoint):

    def __init__(self, ldm_config, hf_checkpoint, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ldm_config = ldm_config
        self.hf_checkpoint = hf_checkpoint

    def on_train_epoch_end(self, trainer, pl_module):
        ldm_checkpoint = self._get_metric_interpolated_filepath_name(
            {'epoch': trainer.current_epoch}, trainer)
        super().on_train_epoch_end(trainer, pl_module)
        convert_ldm_to_hf_vae(ldm_checkpoint, self.ldm_config,
                              self.hf_checkpoint)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train VAE using ldm.")
    parser.add_argument("-d", "--dataset_name", type=str, default=None)
    parser.add_argument("-b", "--batch_size", type=int, default=1)
    parser.add_argument("-c",
                        "--ldm_config_file",
                        type=str,
                        default="config/ldm_autoencoder_kl.yaml")
    parser.add_argument("--ldm_checkpoint_dir",
                        type=str,
                        default="models/ldm-autoencoder-kl")
    parser.add_argument("--hf_checkpoint_dir",
                        type=str,
                        default="models/autoencoder-kl")
    parser.add_argument("-r",
                        "--resume_from_checkpoint",
                        type=str,
                        default=None)
    parser.add_argument("-g",
                        "--gradient_accumulation_steps",
                        type=int,
                        default=1)
    parser.add_argument("--hop_length", type=int, default=512)
    parser.add_argument("--sample_rate", type=int, default=22050)
    parser.add_argument("--n_fft", type=int, default=2048)
    parser.add_argument("--save_images_batches", type=int, default=1000)
    parser.add_argument("--max_epochs", type=int, default=100)
    args = parser.parse_args()

    config = OmegaConf.load(args.ldm_config_file)
    model = instantiate_from_config(config.model)
    model.learning_rate = config.model.base_learning_rate
    data = AudioDiffusionDataModule(
        model_id=args.dataset_name,
        batch_size=args.batch_size,
        channels=config.model.params.ddconfig.in_channels)
    lightning_config = config.pop("lightning", OmegaConf.create())
    trainer_config = lightning_config.get("trainer", OmegaConf.create())
    trainer_config.accumulate_grad_batches = args.gradient_accumulation_steps
    trainer_opt = argparse.Namespace(**trainer_config)
    trainer = Trainer.from_argparse_args(
        trainer_opt,
        max_epochs=args.max_epochs,
        resume_from_checkpoint=args.resume_from_checkpoint,
        callbacks=[
            ImageLogger(every=args.save_images_batches,
                        hop_length=args.hop_length,
                        sample_rate=args.sample_rate,
                        n_fft=args.n_fft),
            HFModelCheckpoint(ldm_config=config,
                              hf_checkpoint=args.hf_checkpoint_dir,
                              dirpath=args.ldm_checkpoint_dir,
                              filename='{epoch:06}',
                              verbose=True,
                              save_last=True)
        ])
    trainer.fit(model, data)