add trainer for vae

README.md

@@ -16,21 +16,25 @@ First train your VAE - `VQGanVAE`
 
 ```python
 import torch
-from muse_maskgit_pytorch import VQGanVAE
+from muse_maskgit_pytorch import VQGanVAE, VQGanVAETrainer
 
 vae = VQGanVAE(
     dim = 256,
     vq_codebook_size = 512
-).cuda()
-
-# mock images
+)
 
-images = torch.randn(4, 3, 256, 256).cuda()
+# train on folder of images, as many images as possible
 
-# do this for as many images as possible
+trainer = VQGanVAETrainer(
+    vae = vae,
+    image_size = 128,             # you may want to start with small images, and then curriculum learn to larger ones, but because the vae is all convolution, it should generalize to 512 (as in paper) without training on it
+    folder = '/path/to/images',
+    batch_size = 4,
+    grad_accum_every = 8,
+    num_train_steps = 50000
+).cuda()
 
-loss = vae(images, return_loss = True)
-loss.backward()
+trainer.train()
 ```
 
 Then pass the trained `VQGanVAE` and a `Transformer` to `MaskGit`

muse_maskgit_pytorch/__init__.py

@@ -1,2 +1,4 @@
 from muse_maskgit_pytorch.vqgan_vae import VQGanVAE
 from muse_maskgit_pytorch.muse_maskgit_pytorch import Transformer, MaskGit, Muse
+
+from muse_maskgit_pytorch.trainers import VQGanVAETrainer