| |
|
|
| 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): |
| super().__init__() |
| self.every = every |
| self.hop_length = hop_length |
|
|
| @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) |
|
|
| 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 |
| 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("--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), |
| 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) |
|
|
