File size: 7,054 Bytes
f8a240c
 
3e8b723
f29faf1
3e8b723
 
 
f29faf1
 
 
 
3e8b723
f29faf1
 
 
 
3e8b723
f29faf1
3e8b723
f29faf1
8a3fb2e
3e8b723
 
 
13aa297
3e8b723
13aa297
8aa7c27
f29faf1
8aa7c27
f29faf1
3e8b723
 
 
 
 
f29faf1
13aa297
f29faf1
 
 
 
3e8b723
 
 
13aa297
3e8b723
 
13aa297
3e8b723
 
 
f29faf1
3e8b723
 
 
f29faf1
3e8b723
d76bdef
327bccf
5bc60f9
 
d76bdef
8aa7c27
 
d76bdef
 
f29faf1
d76bdef
 
327bccf
 
f29faf1
 
 
 
 
 
 
327bccf
3e8b723
d76bdef
f29faf1
d76bdef
 
3e8b723
d76bdef
f29faf1
3e8b723
f29faf1
d76bdef
327bccf
f29faf1
 
 
 
3e8b723
 
 
 
f29faf1
 
3e8b723
f29faf1
8aa7c27
 
 
 
 
 
d122744
8aa7c27
 
 
f29faf1
 
 
 
 
8aa7c27
 
f29faf1
8aa7c27
f29faf1
 
8aa7c27
3e8b723
 
 
8aa7c27
 
f29faf1
 
 
 
 
13aa297
5bc60f9
 
529c646
a8d2b62
3e8b723
 
8aa7c27
13aa297
 
 
 
 
f29faf1
 
3e8b723
 
8aa7c27
3e8b723
8aa7c27
 
b3e97c5
8aa7c27
 
f29faf1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
# based on https://github.com/CompVis/stable-diffusion/blob/main/main.py

import argparse
import os

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

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.0, 1.0)
            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
        self.sample_size = None

    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
        if self.sample_size is None:
            self.sample_size = list(batch["image"].shape[1:3])

    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)
        self.ldm_config.model.params.ddconfig.resolution = self.sample_size
        convert_ldm_to_hf_vae(ldm_checkpoint, self.ldm_config, self.hf_checkpoint, self.sample_size)


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)