preliminary training code for all attention nets completed, project management

- [x] complete sampling code for both Coarse and Fine Transformers, which will be tricky

- [x] make sure full inference with or without prompting works on the `AudioLM` class

- [x] complete full training code for soundstream, taking care of discriminator training

- [x] add efficient gradient penalty for discriminators for soundstream

- [x] wire up sample hz from sound dataset -> transformers, and have proper resampling within during training - think about whether to allow for dataset to have sound files of varying or enforce same sample hz

- [x] full transformer training code for all three transformers

- [ ] add efficient gradient penalty for discriminators for soundstream

- [ ] figure out how to do the normalization across each dimension mentioned in the paper, but ignore it for v1 of the framework

- [ ] full transformer training code for all three transformers

- [ ] wire up sample hz from sound dataset -> transformers, and have proper resampling within during training - think about whether to allow for dataset to have sound files of varying or enforce same sample hz

- [ ] offer option to weight tie coarse, fine, and semantic embeddings across the 3 hierarchical transformers

- [ ] DRY a little at the end

- [ ] figure out how to suppress logging in fairseq

from audiolm_pytorch.vq_wav2vec import FairseqVQWav2Vec

from audiolm_pytorch.hubert_kmeans import HubertWithKmeans

from audiolm_pytorch.trainer import SoundStreamTrainer, SemanticTransformerTrainer, FineTransformerTrainer

from audiolm_pytorch.trainer import SoundStreamTrainer, SemanticTransformerTrainer, FineTransformerTrainer, CoarseTransformerTrainer

        self.print('training complete')

# semantic transformer trainer

# fine transformer trainer

@typechecked

class CoarseTransformerTrainer(nn.Module):

    def __init__(

        self,

        transformer: CoarseTransformer,

        soundstream: SoundStream,

        wav2vec: Optional[Union[FairseqVQWav2Vec, HubertWithKmeans]],

        *,

        num_train_steps,

        batch_size,

        data_max_length = None,

        folder,

        lr = 3e-4,

        grad_accum_every = 1,

        wd = 0.,

        max_grad_norm = 0.5,

        valid_frac = 0.05,

        random_split_seed = 42,

        save_results_every = 100,

        save_model_every = 1000,

        results_folder = './results',

        accelerate_kwargs: dict = dict()

    ):

        super().__init__()

        self.accelerator = Accelerator(**accelerate_kwargs)

        self.transformer = transformer

        self.soundstream = soundstream

        self.wav2vec = wav2vec

        self.train_wrapper = CoarseTransformerWrapper(

            soundstream = soundstream,

            wav2vec = wav2vec,

            transformer = transformer

        )

        self.register_buffer('steps', torch.Tensor([0]))

        self.num_train_steps = num_train_steps

        self.batch_size = batch_size

        self.grad_accum_every = grad_accum_every

        # optimizers

        self.optim = get_optimizer(transformer.parameters(), lr = lr, wd = wd)

        # max grad norm

        self.max_grad_norm = max_grad_norm

        # create dataset

        self.ds = SoundDataset(

            folder,

            max_length = data_max_length,

            target_sample_hz = (

                wav2vec.target_sample_hz,

                soundstream.target_sample_hz

            ), # need 2 waves resampled differently here

            seq_len_multiple_of = soundstream.seq_len_multiple_of

        )

        # split for validation

        if valid_frac > 0:

            train_size = int((1 - valid_frac) * len(self.ds))

            valid_size = len(self.ds) - train_size

            self.ds, self.valid_ds = random_split(self.ds, [train_size, valid_size], generator = torch.Generator().manual_seed(random_split_seed))

            self.print(f'training with dataset of {len(self.ds)} samples and validating with randomly splitted {len(self.valid_ds)} samples')

        else:

            self.valid_ds = self.ds

            self.print(f'training with shared training and valid dataset of {len(self.ds)} samples')

        # dataloader

        self.dl = get_dataloader(self.ds, batch_size = batch_size, shuffle = True)

        self.valid_dl = get_dataloader(self.valid_ds, batch_size = batch_size, shuffle = True)

        # prepare with accelerator

        (

            self.transformer,

            self.optim,

            self.dl,

            self.valid_dl

        ) = self.accelerator.prepare(

            self.transformer,

            self.optim,

            self.dl,

            self.valid_dl

        )

        # dataloader iterators

        self.dl_iter = cycle(self.dl)

        self.valid_dl_iter = cycle(self.valid_dl)

        self.save_model_every = save_model_every

        self.save_results_every = save_results_every    

        self.results_folder = Path(results_folder)

        if len([*self.results_folder.glob('**/*')]) > 0 and yes_or_no('do you want to clear previous experiment checkpoints and results?'):

            rmtree(str(self.results_folder))

        self.results_folder.mkdir(parents = True, exist_ok = True)

        self.train_wrapper.to(self.device)

    def print(self, msg):

        self.accelerator.print(msg)

    @property

    def device(self):

        return self.accelerator.device

    @property

    def is_distributed(self):

        return not (self.accelerator.distributed_type == DistributedType.NO and self.accelerator.num_processes == 1)

    @property

    def is_main(self):

        return self.accelerator.is_main_process

    @property

    def is_local_main(self):

        return self.accelerator.is_local_main_process

    def train_step(self):

        device = self.device

        steps = int(self.steps.item())

        self.transformer.train()

        # logs

        logs = {}

        # update vae (generator)

        for _ in range(self.grad_accum_every):

            wave_wav2vec, wave_soundstream = tuple(map(lambda t: t.to(device), next(self.dl_iter)))

            loss = self.train_wrapper(

                raw_wave = wave_wav2vec,

                raw_wave_for_soundstream = wave_soundstream,

                return_loss = True

            )

            self.accelerator.backward(loss / self.grad_accum_every)

            accum_log(logs, {'loss': loss.item() / self.grad_accum_every})

        if exists(self.max_grad_norm):

            self.accelerator.clip_grad_norm_(self.transformer.parameters(), self.max_grad_norm)

        self.optim.step()

        self.optim.zero_grad()

        # log

        self.print(f"{steps}: loss: {logs['loss']}")

        # sample results every so often

        if self.is_main and not (steps % self.save_results_every):

            wave_wav2vec, wave_soundstream = tuple(map(lambda t: t.to(device), next(self.valid_dl_iter)))

            with torch.no_grad():

                self.train_wrapper.eval()

                valid_loss = self.train_wrapper(

                    raw_wave = wave_wav2vec,

                    raw_wave_for_soundstream = wave_soundstream,

                    return_loss = True

                )

            self.print(f'{steps}: valid loss {valid_loss}')

        # save model every so often

        if self.is_main and not (steps % self.save_model_every):

            state_dict = self.transformer.state_dict()

            model_path = str(self.results_folder / f'fine.transformer.{steps}.pt')

            torch.save(state_dict, model_path)

            self.print(f'{steps}: saving model to {str(self.results_folder)}')

        self.steps += 1

        return logs

    def train(self, log_fn = noop):

        while self.steps < self.num_train_steps:

            logs = self.train_step()

            log_fn(logs)

        self.print('training complete')

# fine transformer trainer

@typechecked

class FineTransformerTrainer(nn.Module):

        # sample results every so often

        if self.is_main and not (steps % self.save_results_every):

            filename = str(steps)

            self.train_wrapper.eval()

            wave = next(self.valid_dl_iter).to(device)

            with torch.no_grad():

                self.train_wrapper.eval()

                valid_loss = self.train_wrapper(raw_wave = wave, return_loss = True)

            self.print(f'{steps}: valid loss {valid_loss}')

setup(

  name = 'audiolm-pytorch',

  packages = find_packages(exclude=[]),

  version = '0.0.52',

  version = '0.0.53',

  license='MIT',

  description = 'AudioLM - Language Modeling Approach to Audio Generation from Google Research - Pytorch',

  author = 'Phil Wang',