Upload folder using huggingface_hub

.gitattributes

@@ -124,3 +124,7 @@ f32c4_down/checkpoint/checkpoint.tmp filter=lfs diff=lfs merge=lfs -text
 f32c4_down/checkpoint.tmp filter=lfs diff=lfs merge=lfs -text
 f32c4_down/checkpointbest.tmp/checkpointbest.tmp.tmp filter=lfs diff=lfs merge=lfs -text
 f32c4_down/checkpointbest.tmp.tmp filter=lfs diff=lfs merge=lfs -text
+f32c32_full/checkpoint/checkpoint.tmp filter=lfs diff=lfs merge=lfs -text
+f32c32_full/checkpoint.tmp filter=lfs diff=lfs merge=lfs -text
+f32c32_full/checkpointbest.tmp/checkpointbest.tmp.tmp filter=lfs diff=lfs merge=lfs -text
+f32c32_full/checkpointbest.tmp.tmp filter=lfs diff=lfs merge=lfs -text

f32c32_full/checkpoint.tmp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:640d4065303f113997eb89cface839503937a7827652069bd0c82c2a101ed6a9
+size 1676176195

f32c32_full/checkpoint/checkpoint.tmp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:29a85dabf5004a4405402384a3b1a0c0ee81894d022ba01cc726ed8f4254a9e0
+size 1676176195

f32c32_full/checkpointbest.tmp.tmp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:79318fec959fe175e33826cab87069568360411d0b915f84f3fc94d8f4b6db75
+size 1676176195

f32c32_full/checkpointbest.tmp/checkpointbest.tmp.tmp

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd482a1de53f1ebc027b706f0a1bf3ca1a8ffbaa26f59d138161987dcaa49174
+size 1676176195

f32c32_full/train.py

@@ -0,0 +1,433 @@
+try: # For debugging
+    from localutils.debugger import enable_debug
+    enable_debug()
+except ImportError:
+    pass
+
+import flax.linen as nn
+import jax.numpy as jnp
+from absl import app, flags
+from functools import partial
+import numpy as np
+import tqdm
+import jax
+import jax.numpy as jnp
+import flax
+import optax
+import wandb
+from ml_collections import config_flags
+import ml_collections
+import tensorflow_datasets as tfds
+import tensorflow as tf
+tf.config.set_visible_devices([], "GPU")
+tf.config.set_visible_devices([], "TPU")
+import matplotlib.pyplot as plt
+from typing import Any
+import os
+
+from utils.wandb import setup_wandb, default_wandb_config
+from utils.train_state import TrainState, target_update
+from utils.checkpoint import Checkpoint
+from utils.pretrained_resnet import get_pretrained_embs, get_pretrained_model
+from utils.fid import get_fid_network, fid_from_stats
+from models.vqvae import VQVAE
+from models.discriminator import Discriminator
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string('dataset_name', 'imagenet256', 'Environment name.')
+flags.DEFINE_string('save_dir', "/home/lambda/jax-vqvae-vqgan/chkpts/checkpoint", 'Save dir (if not None, save params).')
+flags.DEFINE_string('load_dir', "/home/lambda/jax-vqvae-vqgan/chkpts/checkpoint.tmp" , 'Load dir (if not None, load params from here).')
+flags.DEFINE_integer('seed', 0, 'Random seed.')
+flags.DEFINE_integer('log_interval', 1000, 'Logging interval.')
+flags.DEFINE_integer('eval_interval', 1000, 'Eval interval.')
+flags.DEFINE_integer('save_interval', 1000, 'Save interval.')
+flags.DEFINE_integer('batch_size', 64, 'Total Batch size.')
+flags.DEFINE_integer('max_steps', int(1_000_000), 'Number of training steps.')
+
+model_config = ml_collections.ConfigDict({
+    # VQVAE
+    'lr': 0.0001,
+    'beta1': 0.0,#.5
+    'beta2': 0.99,#.9
+    'lr_warmup_steps': 2000,
+    'lr_decay_steps': 500_000,#They use 'lambdalr'
+    'filters': 128,
+    'num_res_blocks': 2,
+    'channel_multipliers': (1, 1, 2, 2, 4, 4),#We want 5 blocks for downsample 4x
+    'embedding_dim': 32, # For FSQ, a good default is 4.
+    'norm_type': 'GN',
+    'weight_decay': 0.05,#None maybe?
+    'clip_gradient': 1.0,
+    'l2_loss_weight': 1.0,#They use L1 actually
+    'eps_update_rate': 0.9999,
+    # Quantizer
+    'quantizer_type': 'vq', # or 'fsq', 'kl'
+    # Quantizer (VQ)
+    'quantizer_loss_ratio': 1,
+    'codebook_size': 1024,
+    'entropy_loss_ratio': 0.1,
+    'entropy_loss_type': 'softmax',
+    'entropy_temperature': 0.01,
+    'commitment_cost': 0.25,
+    # Quantizer (FSQ)
+    'fsq_levels': 5, # Bins per dimension.
+    # Quantizer (KL)
+    'kl_weight': 0.000001,#They use 1e-6 on their stuff LUL. .001 is the default
+    # GAN
+    'g_adversarial_loss_weight': 0.5,
+    'g_grad_penalty_cost': 10,
+    'perceptual_loss_weight': 0.5,
+    'gan_warmup_steps': 25000,
+})
+
+wandb_config = default_wandb_config()
+wandb_config.update({
+    'project': 'vqvae',
+    'name': 'vqvae_{dataset_name}',
+})
+
+config_flags.DEFINE_config_dict('wandb', wandb_config, lock_config=False)
+config_flags.DEFINE_config_dict('model', model_config, lock_config=False)
+
+##############################################
+## Model Definitions.
+##############################################
+
+@jax.vmap
+def sigmoid_cross_entropy_with_logits(*, labels: jnp.ndarray, logits: jnp.ndarray) -> jnp.ndarray:
+    """https://github.com/google-research/maskgit/blob/main/maskgit/libml/losses.py
+    """
+    zeros = jnp.zeros_like(logits, dtype=logits.dtype)
+    condition = (logits >= zeros)
+    relu_logits = jnp.where(condition, logits, zeros)
+    neg_abs_logits = jnp.where(condition, -logits, logits)
+    return relu_logits - logits * labels + jnp.log1p(jnp.exp(neg_abs_logits))
+
+class VQGANModel(flax.struct.PyTreeNode):
+    rng: Any
+    config: dict = flax.struct.field(pytree_node=False)
+    vqvae: TrainState
+    vqvae_eps: TrainState
+    discriminator: TrainState
+
+    # Train G and D.
+    @partial(jax.pmap, axis_name='data', in_axes=(0, 0))
+    def update(self, images, pmap_axis='data'):
+        new_rng, curr_key = jax.random.split(self.rng, 2)
+
+        resnet, resnet_params = get_pretrained_model('resnet50', 'data/resnet_pretrained.npy')
+
+        is_gan_training = 1.0 - (self.vqvae.step < self.config['gan_warmup_steps']).astype(jnp.float32)
+
+        def loss_fn(params_vqvae, params_disc):
+            # Reconstruct image
+            reconstructed_images, result_dict = self.vqvae(images, params=params_vqvae, rngs={'noise': curr_key})
+            print("Reconstructed images shape", reconstructed_images.shape)
+            print("Input images shape", images.shape)
+            assert reconstructed_images.shape == images.shape
+
+            # GAN loss on VQVAE output.
+            discriminator_fn = lambda x: self.discriminator(x, params=params_disc)
+            real_logit, vjp_fn = jax.vjp(discriminator_fn, images, has_aux=False)
+            gradient = vjp_fn(jnp.ones_like(real_logit))[0] # Gradient of discriminator output wrt. real images.
+            gradient = gradient.reshape((images.shape[0], -1))
+            gradient = jnp.asarray(gradient, jnp.float32)
+            penalty = jnp.sum(jnp.square(gradient), axis=-1)
+            penalty = jnp.mean(penalty) # Gradient penalty for training D.
+            fake_logit = discriminator_fn(reconstructed_images)
+            d_loss_real = sigmoid_cross_entropy_with_logits(labels=jnp.ones_like(real_logit), logits=real_logit).mean()
+            d_loss_fake = sigmoid_cross_entropy_with_logits(labels=jnp.zeros_like(fake_logit), logits=fake_logit).mean()
+            loss_d = d_loss_real + d_loss_fake + (penalty * self.config['g_grad_penalty_cost'])
+
+            d_loss_for_vae = sigmoid_cross_entropy_with_logits(labels=jnp.ones_like(fake_logit), logits=fake_logit).mean()
+            d_loss_for_vae = d_loss_for_vae * is_gan_training
+
+            real_pools, _ = get_pretrained_embs(resnet_params, resnet, images=images)
+            fake_pools, _ = get_pretrained_embs(resnet_params, resnet, images=reconstructed_images)
+            perceptual_loss = jnp.mean((real_pools - fake_pools)**2)
+
+            l2_loss = jnp.mean((reconstructed_images - images) ** 2)
+            quantizer_loss = result_dict['quantizer_loss'] if 'quantizer_loss' in result_dict else 0.0
+            if self.config['quantizer_type'] == 'kl' or self.config["quantizer_type"] == "kl_two":
+                quantizer_loss = quantizer_loss * self.config['kl_weight']
+            loss_vae = (l2_loss * FLAGS.model['l2_loss_weight']) \
+                + (quantizer_loss * FLAGS.model['quantizer_loss_ratio']) \
+                + (d_loss_for_vae * FLAGS.model['g_adversarial_loss_weight']) \
+                + (perceptual_loss * FLAGS.model['perceptual_loss_weight'])
+            codebook_usage = result_dict['usage'] if 'usage' in result_dict else 0.0
+            return (loss_vae, loss_d), {
+                'loss_vae': loss_vae,
+                'loss_d': loss_d,
+                'l2_loss': l2_loss,
+                'd_loss_for_vae': d_loss_for_vae,
+                'perceptual_loss': perceptual_loss,
+                'quantizer_loss': quantizer_loss,
+                'codebook_usage': codebook_usage,
+            }
+        
+        # This is a fancy way to do 'jax.grad' so (loss_vae, params_vqvae) and (loss_d, params_disc) are differentiated.
+        _, grad_fn, info = jax.vjp(loss_fn, self.vqvae.params, self.discriminator.params, has_aux=True)
+        vae_grads, _ = grad_fn((1., 0.))
+        _, d_grads = grad_fn((0., 1.))
+
+        vae_grads = jax.lax.pmean(vae_grads, axis_name=pmap_axis)
+        d_grads = jax.lax.pmean(d_grads, axis_name=pmap_axis)
+        d_grads = jax.tree_map(lambda x: x * is_gan_training, d_grads)
+
+        info = jax.lax.pmean(info, axis_name=pmap_axis)
+        if self.config['quantizer_type'] == 'fsq':
+            info['codebook_usage'] = jnp.sum(info['codebook_usage'] > 0) / info['codebook_usage'].shape[-1]
+
+        updates, new_opt_state = self.vqvae.tx.update(vae_grads, self.vqvae.opt_state, self.vqvae.params)
+        new_params = optax.apply_updates(self.vqvae.params, updates)
+        new_vqvae = self.vqvae.replace(step=self.vqvae.step + 1, params=new_params, opt_state=new_opt_state)
+
+        updates, new_opt_state = self.discriminator.tx.update(d_grads, self.discriminator.opt_state, self.discriminator.params)
+        new_params = optax.apply_updates(self.discriminator.params, updates)
+        new_discriminator = self.discriminator.replace(step=self.discriminator.step + 1, params=new_params, opt_state=new_opt_state)
+
+        info['grad_norm_vae'] = optax.global_norm(vae_grads)
+        info['grad_norm_d'] = optax.global_norm(d_grads)
+        info['update_norm'] = optax.global_norm(updates)
+        info['param_norm'] = optax.global_norm(new_params)
+        info['is_gan_training'] = is_gan_training
+
+        new_vqvae_eps = target_update(new_vqvae, self.vqvae_eps, 1-self.config['eps_update_rate'])
+
+        new_model = self.replace(rng=new_rng, vqvae=new_vqvae, vqvae_eps=new_vqvae_eps, discriminator=new_discriminator)
+        return new_model, info
+    
+    @partial(jax.pmap, axis_name='data', in_axes=(0, 0))
+    def reconstruction(self, images, pmap_axis='data', sampling = False):
+        if not sampling:
+            reconstructed_images, _ = self.vqvae_eps(images)
+        else:
+            new_rng, curr_key = jax.random.split(self.rng, 2)
+            reconstructed_images, _ = self.vqvae(images, rngs={'noise': curr_key})
+
+        reconstructed_images = jnp.clip(reconstructed_images, 0, 1)
+        return reconstructed_images
+
+##############################################
+## Training Code.
+##############################################
+def main(_):
+    np.random.seed(FLAGS.seed)
+    print("Using devices", jax.local_devices())
+    device_count = len(jax.local_devices())
+    global_device_count = jax.device_count()
+    local_batch_size = FLAGS.batch_size // (global_device_count // device_count)
+    print("Device count", device_count)
+    print("Global device count", global_device_count)
+    print("Global Batch: ", FLAGS.batch_size)
+    print("Node Batch: ", local_batch_size)
+    print("Device Batch:", local_batch_size // device_count)
+
+    # Create wandb logger
+    if jax.process_index() == 0:
+        setup_wandb(FLAGS.model.to_dict(), **FLAGS.wandb)
+
+    def get_dataset(is_train):
+        if 'imagenet' in FLAGS.dataset_name:
+            def deserialization_fn(data):
+                image = data['image']
+                min_side = tf.minimum(tf.shape(image)[0], tf.shape(image)[1])
+                image = tf.image.resize_with_crop_or_pad(image, min_side, min_side)
+                if 'imagenet256' in FLAGS.dataset_name:
+                    image = tf.image.resize(image, (256, 256))
+                elif 'imagenet128' in FLAGS.dataset_name:
+                    image = tf.image.resize(image, (128, 128))
+                else:
+                    raise ValueError(f"Unknown dataset {FLAGS.dataset_name}")
+                if is_train:
+                    image = tf.image.random_flip_left_right(image)
+                image = tf.cast(image, tf.float32) / 255.0
+                return image
+
+            
+            split = tfds.split_for_jax_process('train' if is_train else 'validation', drop_remainder=True)
+            print(split)
+            dataset = tfds.load('imagenet2012', split=split, data_dir = "/dev/shm")
+            dataset = dataset.map(deserialization_fn, num_parallel_calls=tf.data.AUTOTUNE)
+            dataset = dataset.shuffle(10000, seed=42, reshuffle_each_iteration=True)
+            dataset = dataset.repeat()
+            dataset = dataset.batch(local_batch_size)
+            dataset = dataset.prefetch(tf.data.AUTOTUNE)
+            dataset = tfds.as_numpy(dataset)
+            dataset = iter(dataset)
+            return dataset
+        else:
+            raise ValueError(f"Unknown dataset {FLAGS.dataset_name}")
+        
+    dataset = get_dataset(is_train=True)
+    dataset_valid = get_dataset(is_train=False)
+    example_obs = next(dataset)[:1]
+
+    get_fid_activations = get_fid_network()
+    if not os.path.exists('./data/imagenet256_fidstats_openai.npz'):
+        raise ValueError("Please download the FID stats file! See the README.")
+#    truth_fid_stats = np.load('data/imagenet256_fidstats_openai.npz')
+    truth_fid_stats = np.load("./base_stats.npz")
+
+    rng = jax.random.PRNGKey(FLAGS.seed)
+    rng, param_key = jax.random.split(rng)
+    print("Total Memory on device:", float(jax.local_devices()[0].memory_stats()['bytes_limit']) / 1024**3, "GB")
+
+    ###################################
+    # Creating Model and put on devices.
+    ###################################
+    FLAGS.model.image_channels = example_obs.shape[-1]
+    FLAGS.model.image_size = example_obs.shape[1]
+    vqvae_def = VQVAE(FLAGS.model, train=True)
+    vqvae_params = vqvae_def.init({'params': param_key, 'noise': param_key}, example_obs)['params']
+    tx = optax.adam(learning_rate=FLAGS.model['lr'], b1=FLAGS.model['beta1'], b2=FLAGS.model['beta2'])
+    vqvae_ts = TrainState.create(vqvae_def, vqvae_params, tx=tx)
+    vqvae_def_eps = VQVAE(FLAGS.model, train=False)
+    vqvae_eps_ts = TrainState.create(vqvae_def_eps, vqvae_params)
+    print("Total num of VQVAE parameters:", sum(x.size for x in jax.tree_util.tree_leaves(vqvae_params)))
+
+    discriminator_def = Discriminator(FLAGS.model)
+    discriminator_params = discriminator_def.init(param_key, example_obs)['params']
+    tx = optax.adam(learning_rate=FLAGS.model['lr'], b1=FLAGS.model['beta1'], b2=FLAGS.model['beta2'])
+    discriminator_ts = TrainState.create(discriminator_def, discriminator_params, tx=tx)
+    print("Total num of Discriminator parameters:", sum(x.size for x in jax.tree_util.tree_leaves(discriminator_params)))
+
+    model = VQGANModel(rng=rng, vqvae=vqvae_ts, vqvae_eps=vqvae_eps_ts, discriminator=discriminator_ts, config=FLAGS.model)
+
+    if FLAGS.load_dir is not None:
+        try:
+            cp = Checkpoint(FLAGS.load_dir)
+            model = cp.load_model(model)
+            print("Loaded model with step", model.vqvae.step)
+        except:
+            print("Random init")
+    else:
+        print("Random init")
+
+    model = flax.jax_utils.replicate(model, devices=jax.local_devices())
+    jax.debug.visualize_array_sharding(model.vqvae.params['decoder']['Conv_0']['bias'])
+
+    ###################################
+    # Train Loop
+    ###################################
+    
+    best_fid = 100000
+
+    for i in tqdm.tqdm(range(1, FLAGS.max_steps + 1),
+                       smoothing=0.1,
+                       dynamic_ncols=True):
+
+        batch_images = next(dataset)
+        batch_images = batch_images.reshape((len(jax.local_devices()), -1, *batch_images.shape[1:])) # [devices, batch//devices, etc..]
+
+        model, update_info = model.update(batch_images)
+
+        if i % FLAGS.log_interval == 0:
+            update_info = jax.tree_map(lambda x: x.mean(), update_info)
+            train_metrics = {f'training/{k}': v for k, v in update_info.items()}
+            if jax.process_index() == 0:
+                wandb.log(train_metrics, step=i)
+
+        if i % FLAGS.eval_interval == 0:
+            # Print some images
+            reconstructed_images = model.reconstruction(batch_images) # [devices, 8, 256, 256, 3]
+            valid_images = next(dataset_valid)
+            valid_images = valid_images.reshape((len(jax.local_devices()), -1, *valid_images.shape[1:])) # [devices, batch//devices, etc..]
+            valid_reconstructed_images = model.reconstruction(valid_images) # [devices, 8, 256, 256, 3]
+
+            if jax.process_index() == 0:
+                wandb.log({'batch_image_mean': batch_images.mean()}, step=i)
+                wandb.log({'reconstructed_images_mean': reconstructed_images.mean()}, step=i)
+                wandb.log({'batch_image_std': batch_images.std()}, step=i)
+                wandb.log({'reconstructed_images_std': reconstructed_images.std()}, step=i)
+
+                # plot comparison witah matplotlib. put each reconstruction side by side.
+                fig, axs = plt.subplots(2, 8, figsize=(30, 15))
+                #print("batch shape", batch_images.shape)#batch shape (4, 32, 256, 256, 3) #THE FIRST SHAPE IS DEVICES
+                #print("recon shape", reconstructed_images.shape)#it's all the same lol
+                #print("valid shape", valid_images.shape)
+                #it seems to be made for 8 device, aka tpuv3 instead
+                for j in range(4):#fuck it
+                    axs[0, j].imshow(batch_images[j, 0], vmin=0, vmax=1)
+                    axs[1, j].imshow(reconstructed_images[j, 0], vmin=0, vmax=1)
+                wandb.log({'reconstruction': wandb.Image(fig)}, step=i)
+                plt.close(fig)
+                fig, axs = plt.subplots(2, 8, figsize=(30, 15))
+                for j in range(4):
+                    axs[0, j].imshow(valid_images[j, 0], vmin=0, vmax=1)
+                    axs[1, j].imshow(valid_reconstructed_images[j, 0], vmin=0, vmax=1)
+                wandb.log({'reconstruction_valid': wandb.Image(fig)}, step=i)
+                plt.close(fig)
+
+            # Validation Losses
+            _, valid_update_info = model.update(valid_images)
+            valid_update_info = jax.tree_map(lambda x: x.mean(), valid_update_info)
+            valid_metrics = {f'validation/{k}': v for k, v in valid_update_info.items()}
+            if jax.process_index() == 0:
+                wandb.log(valid_metrics, step=i)
+
+            # FID measurement.
+            activations = []
+            activations2 = []
+            for _ in range(780):#This is apprximately 40k
+                valid_images = next(dataset_valid)
+                valid_images = valid_images.reshape((len(jax.local_devices()), -1, *valid_images.shape[1:])) # [devices, batch//devices, etc..]
+                valid_reconstructed_images = model.reconstruction(valid_images) # [devices, 8, 256, 256, 3]
+
+                valid_reconstructed_images = jax.image.resize(valid_reconstructed_images, (valid_images.shape[0], valid_images.shape[1], 299, 299, 3),
+                                                               method='bilinear', antialias=False)
+                valid_reconstructed_images = 2 * valid_reconstructed_images - 1
+                activations += [np.array(get_fid_activations(valid_reconstructed_images))[..., 0, 0, :]]
+
+                
+                #Only needed when we save
+                #valid_reconstructed_images = jax.image.resize(valid_images, (valid_images.shape[0], valid_images.shape[1], 299, 299, 3),
+                                                               #method='bilinear', antialias=False)
+                #valid_reconstructed_images = 2 * valid_reconstructed_images - 1
+                #activations2 += [np.array(get_fid_activations(valid_reconstructed_images))[..., 0, 0, :]]
+
+
+                # TODO: use all_gather to get activations from all devices.
+            #This seems to be FID with only 64 images?
+            activations = np.concatenate(activations, axis=0)
+            activations = activations.reshape((-1, activations.shape[-1]))
+
+ #           activations2 = np.concatenate(activations2, axis = 0)
+ #           activations2 = activations2.reshape((-1, activations2.shape[-1]))
+
+            print("doing this much FID", activations.shape)#8192, 2048 should be 2048 items then I guess
+            mu1 = np.mean(activations, axis=0)
+            sigma1 = np.cov(activations, rowvar=False)
+            fid = fid_from_stats(mu1, sigma1, truth_fid_stats['mu'], truth_fid_stats['sigma'])
+            
+#            mu2 = np.mean(activations2, axis = 0)
+#            sigma2 = np.cov(activations2, rowvar = False)
+
+            #save mu2 and sigma2
+            #And then exit for now
+#            np.savez("base.npz", mu = mu2, sigma = sigma2)
+#            exit()
+
+            #Used with loading base
+            #fid = fid_from_stats(mu1, sigma1, mu2, sigma2)
+
+            if jax.process_index() == 0:
+                wandb.log({'validation/fid': fid}, step=i)
+                print("validation FID at step", i, fid)
+                #Then if fid is smaller than previous best FID, save new FID
+                if fid < best_fid:
+                    model_single = flax.jax_utils.unreplicate(model)
+                    cp = Checkpoint(FLAGS.save_dir + "best.tmp")
+                    cp.set_model(model_single)
+                    cp.save()
+                    best_fid = fid
+
+        if (i % FLAGS.save_interval == 0) and (FLAGS.save_dir is not None):
+            if jax.process_index() == 0:
+                model_single = flax.jax_utils.unreplicate(model)
+                cp = Checkpoint(FLAGS.save_dir)
+                cp.set_model(model_single)
+                cp.save()
+
+if __name__ == '__main__':
+    app.run(main)

f32c32_full/vqvae.py

@@ -0,0 +1,453 @@
+from typing import Any
+import flax.linen as nn
+import jax.numpy as jnp
+import functools
+import ml_collections
+import jax
+
+###########################
+### Helper Modules
+### https://github.com/google-research/maskgit/blob/main/maskgit/nets/layers.py
+###########################
+
+def get_norm_layer(norm_type):
+    """Normalization layer."""
+    if norm_type == 'BN':
+        raise NotImplementedError
+    elif norm_type == 'LN':
+        norm_fn = functools.partial(nn.LayerNorm)
+    elif norm_type == 'GN':
+        norm_fn = functools.partial(nn.GroupNorm)
+    else:
+        raise NotImplementedError
+    return norm_fn
+
+
+def tensorflow_style_avg_pooling(x, window_shape, strides, padding: str):
+    pool_sum = jax.lax.reduce_window(x, 0.0, jax.lax.add,
+                                   (1,) + window_shape + (1,),
+                                   (1,) + strides + (1,), padding)
+    pool_denom = jax.lax.reduce_window(
+        jnp.ones_like(x), 0.0, jax.lax.add, (1,) + window_shape + (1,),
+        (1,) + strides + (1,), padding)
+    return pool_sum / pool_denom
+
+def upsample(x, factor=2):
+    n, h, w, c = x.shape
+    x = jax.image.resize(x, (n, h * factor, w * factor, c), method='nearest')
+    return x
+
+def dsample(x):
+    return tensorflow_style_avg_pooling(x, (2, 2), strides=(2, 2), padding='same')
+
+def squared_euclidean_distance(a: jnp.ndarray,
+                               b: jnp.ndarray,
+                               b2: jnp.ndarray = None) -> jnp.ndarray:
+    """Computes the pairwise squared Euclidean distance.
+
+    Args:
+        a: float32: (n, d): An array of points.
+        b: float32: (m, d): An array of points.
+        b2: float32: (d, m): b square transpose.
+
+    Returns:
+        d: float32: (n, m): Where d[i, j] is the squared Euclidean distance between
+        a[i] and b[j].
+    """
+    if b2 is None:
+        b2 = jnp.sum(b.T**2, axis=0, keepdims=True)
+    a2 = jnp.sum(a**2, axis=1, keepdims=True)
+    ab = jnp.matmul(a, b.T)
+    d = a2 - 2 * ab + b2
+    return d
+
+def entropy_loss_fn(affinity, loss_type="softmax", temperature=1.0):
+    """Calculates the entropy loss. Affinity is the similarity/distance matrix."""
+    flat_affinity = affinity.reshape(-1, affinity.shape[-1])
+    flat_affinity /= temperature
+    probs = jax.nn.softmax(flat_affinity, axis=-1)
+    log_probs = jax.nn.log_softmax(flat_affinity + 1e-5, axis=-1)
+    if loss_type == "softmax":
+        target_probs = probs
+    elif loss_type == "argmax":
+        codes = jnp.argmax(flat_affinity, axis=-1)
+        onehots = jax.nn.one_hot(
+            codes, flat_affinity.shape[-1], dtype=flat_affinity.dtype)
+        onehots = probs - jax.lax.stop_gradient(probs - onehots)
+        target_probs = onehots
+    else:
+        raise ValueError("Entropy loss {} not supported".format(loss_type))
+    avg_probs = jnp.mean(target_probs, axis=0)
+    avg_entropy = -jnp.sum(avg_probs * jnp.log(avg_probs + 1e-5))
+    sample_entropy = -jnp.mean(jnp.sum(target_probs * log_probs, axis=-1))
+    loss = sample_entropy - avg_entropy
+    return loss
+
+def sg(x):
+    return jax.lax.stop_gradient(x)
+
+
+
+
+###########################
+### Modules
+###########################
+
+class ResBlock(nn.Module):
+    """Basic Residual Block."""
+    filters: int
+    norm_fn: Any
+    activation_fn: Any
+
+    @nn.compact
+    def __call__(self, x):
+        input_dim = x.shape[-1]
+        residual = x
+        x = self.norm_fn()(x)
+        x = self.activation_fn(x)
+        x = nn.Conv(self.filters, kernel_size=(3, 3), use_bias=False)(x)
+        x = self.norm_fn()(x)
+        x = self.activation_fn(x)
+        x = nn.Conv(self.filters, kernel_size=(3, 3), use_bias=False)(x)
+
+        if input_dim != self.filters:#Basically if input doesn't match output, use a skip
+            residual = nn.Conv(self.filters, kernel_size=(1, 1), use_bias=False)(x)
+        return x + residual
+    
+class Encoder(nn.Module):
+    """From [H,W,D] image to [H',W',D'] embedding. Using Conv layers."""
+    config: ml_collections.ConfigDict
+
+    def setup(self):
+        self.filters = self.config.filters#filters is the original setup
+        self.num_res_blocks = self.config.num_res_blocks
+        self.channel_multipliers = self.config.channel_multipliers
+        self.embedding_dim = self.config.embedding_dim
+        self.norm_type = self.config.norm_type
+        self.activation_fn = nn.swish
+
+    def pixels(self, x):
+        #print("pixel shuffle x shape", x.shape)
+        x = pixel_unshuffle(x, 2)
+        #print(x.shape)
+        B, H, W, C = x.shape
+        x = jnp.reshape(x, (B, H, W, int(C/4), 4))
+        #print(x.shape)
+        x = jnp.mean(x, axis = -1)
+        #print(x.shape)
+        #exit()
+        return x
+
+
+    @nn.compact
+    def __call__(self, x):
+        print("Initializing encoder.")
+        norm_fn = get_norm_layer(norm_type=self.norm_type)
+        block_args = dict(norm_fn=norm_fn, activation_fn=self.activation_fn)
+        print("Incoming encoder shape", x.shape)
+        x = nn.Conv(self.filters, kernel_size=(3, 3), use_bias=False)(x)
+        print('Encoder layer', x.shape)
+        num_blocks = len(self.channel_multipliers)
+
+        #The way SD works, is it does 2x resnet, not changing anything, then downsample
+        #It does this 3 times, leading to 8x downsample
+        #Then it has an extra resnet block, and THEN from 512 to 8 / 4
+        
+        #So the DCAE architecture is like 4x resnet, down
+        #And then efficient vit down
+        for i in range(num_blocks):
+            filters = self.filters * self.channel_multipliers[i]
+            for _ in range(self.num_res_blocks):
+                x = ResBlock(filters, **block_args)(x)
+
+            if i < num_blocks - 1:#For each block *except end* do downsample
+                print("doing downsample")
+                #If we want to do it DCAE style, they do channel averaging between before downsample and after
+                if self.channel_multipliers[i] != -1:
+                    print("pre pixels", x.shape)
+                    pixel_x = self.pixels(x)
+                    print("pixel_x", pixel_x.shape)
+                    x = dsample(x) + pixel_x
+                    print("post", x.shape)
+                else:
+                    x = dsample(x)
+                    print("other post", x.shape)
+                
+            print('Encoder layer', x.shape)
+        
+        #After we are done downsampling, we do the 2 resnet, and down below here, we have the 2 midblock?
+
+        for _ in range(self.num_res_blocks):
+            x = ResBlock(filters, **block_args)(x)
+            print('Encoder layer final', x.shape)
+
+        x = norm_fn()(x)
+        x = self.activation_fn(x)
+        last_dim = self.embedding_dim*2 if self.config['quantizer_type'] == 'kl' else self.embedding_dim
+        x = nn.Conv(last_dim, kernel_size=(1, 1))(x)
+        print("Final embeddings are size", x.shape)
+        return x
+    
+class Decoder(nn.Module):
+    """From [H',W',D'] embedding to [H,W,D] embedding. Using Conv layers."""
+
+    config: ml_collections.ConfigDict
+
+    def setup(self):
+        self.filters = self.config.filters
+        self.num_res_blocks = self.config.num_res_blocks
+        self.channel_multipliers = self.config.channel_multipliers
+        self.norm_type = self.config.norm_type
+        self.image_channels = self.config.image_channels
+        self.activation_fn = nn.swish
+
+    def pixels(self, x):
+        print("pixels shape", x.shape)
+        x = jnp.repeat(x, 4, axis = -1)
+        print(x.shape)
+        x = pixel_shuffle(x, 2)
+        print(x.shape)
+        print("done duplicating")
+        return x
+
+    @nn.compact
+    def __call__(self, x):
+        norm_fn = get_norm_layer(norm_type=self.norm_type)
+        block_args = dict(norm_fn=norm_fn, activation_fn=self.activation_fn,)
+        num_blocks = len(self.channel_multipliers)
+        filters = self.filters * self.channel_multipliers[-1]
+        print("Decoder incoming shape", x.shape)
+
+        #We don't need to do anything here because it'll put it back to 512
+
+        x = nn.Conv(filters, kernel_size=(3, 3), use_bias=True)(x)
+        print("Decoder input", x.shape)
+        
+
+        #This is the mid block
+        for _ in range(self.num_res_blocks):
+            x = ResBlock(filters, **block_args)(x)
+            print('Mid Block Decoder layer', x.shape)
+
+        #First two SET of blocks is just 3 resnet, no channel changes, we are already at 4x = 512
+        
+        for i in reversed(range(num_blocks)):
+            filters = self.filters * self.channel_multipliers[i]
+            for _ in range(self.num_res_blocks + 1):
+                x = ResBlock(filters, **block_args)(x)
+            if i > 0:
+                #We do pixel channel downsampling every time we downsample spatially.
+                pixel = self.pixels(x)
+                print("pre up", x.shape)
+                x = upsample(x, 2)
+                print("post up", x.shape)
+                x = x + pixel
+                x = nn.Conv(filters, kernel_size=(3, 3))(x)
+            print('Decoder layer', x.shape)
+        x = norm_fn()(x)
+        x = self.activation_fn(x)
+        x = nn.Conv(self.image_channels, kernel_size=(3, 3))(x)
+        return x
+    
+class VectorQuantizer(nn.Module):
+    """Basic vector quantizer."""
+    config: ml_collections.ConfigDict
+    train: bool
+
+    @nn.compact
+    def __call__(self, x):
+        codebook_size = self.config.codebook_size
+        emb_dim = x.shape[-1]
+        codebook = self.param(
+            "codebook",
+            jax.nn.initializers.variance_scaling(scale=1.0, mode="fan_in", distribution="uniform"),
+            (codebook_size, emb_dim))
+        codebook = jnp.asarray(codebook) # (codebook_size, emb_dim)
+        distances = jnp.reshape(
+            squared_euclidean_distance(jnp.reshape(x, (-1, emb_dim)), codebook),
+            x.shape[:-1] + (codebook_size,)) # [x, codebook_size] similarity matrix.
+        encoding_indices = jnp.argmin(distances, axis=-1)
+        encoding_onehot = jax.nn.one_hot(encoding_indices, codebook_size)
+        quantized = self.quantize(encoding_onehot)
+        result_dict = dict()
+        if self.train:
+            e_latent_loss = jnp.mean((sg(quantized) - x)**2) * self.config.commitment_cost
+            q_latent_loss = jnp.mean((quantized - sg(x))**2)
+            entropy_loss = 0.0
+            if self.config.entropy_loss_ratio != 0:
+                entropy_loss = entropy_loss_fn(
+                    -distances,
+                    loss_type=self.config.entropy_loss_type,
+                    temperature=self.config.entropy_temperature
+                ) * self.config.entropy_loss_ratio
+            e_latent_loss = jnp.asarray(e_latent_loss, jnp.float32)
+            q_latent_loss = jnp.asarray(q_latent_loss, jnp.float32)
+            entropy_loss = jnp.asarray(entropy_loss, jnp.float32)
+            loss = e_latent_loss + q_latent_loss + entropy_loss
+            result_dict = dict(
+                quantizer_loss=loss,
+                e_latent_loss=e_latent_loss,
+                q_latent_loss=q_latent_loss,
+                entropy_loss=entropy_loss)
+            quantized = x + jax.lax.stop_gradient(quantized - x)
+
+        result_dict.update({
+            "z_ids": encoding_indices,
+        })
+        return quantized, result_dict
+
+    def quantize(self, encoding_onehot: jnp.ndarray) -> jnp.ndarray:
+        codebook = jnp.asarray(self.variables["params"]["codebook"])
+        return jnp.dot(encoding_onehot, codebook)
+
+    def decode_ids(self, ids: jnp.ndarray) -> jnp.ndarray:
+        codebook = self.variables["params"]["codebook"]
+        return jnp.take(codebook, ids, axis=0)
+
+class KLQuantizer(nn.Module):
+    config: ml_collections.ConfigDict
+    train: bool
+
+    @nn.compact
+    def __call__(self, x):
+        emb_dim = x.shape[-1] // 2 # Use half as means, half as logvars.
+        means = x[..., :emb_dim]
+        logvars = x[..., emb_dim:]
+        if not self.train:
+            result_dict = dict()
+            return means, result_dict
+        else:
+            noise = jax.random.normal(self.make_rng("noise"), means.shape)
+            stds = jnp.exp(0.5 * logvars)
+            z = means + stds * noise
+            kl_loss = -0.5 * jnp.mean(1 + logvars - means**2 - jnp.exp(logvars))
+            result_dict = dict(quantizer_loss=kl_loss)
+            return z, result_dict
+        
+class AEQuantizer(nn.Module): #cooking
+    config: ml_collections.ConfigDict
+    train: bool
+
+    @nn.compact
+    def __call__(self, x):
+        result_dict = dict()
+        return x, result_dict
+
+from einops import rearrange
+
+def pixel_unshuffle(x, factor):
+
+    x = rearrange(x, '... (h b1) (w b2) c -> ... h w (c b1 b2)', b1=factor, b2=factor)
+    return x
+def pixel_shuffle(x, factor):
+    x = rearrange(x, '... h w (c b1 b2) -> ... (h b1) (w b2) c', b1=factor, b2=factor)
+    return x
+
+class KLQuantizerTwo(nn.Module):
+    config: ml_collections.ConfigDict
+    train: bool
+
+    @nn.compact
+    def __call__(self, x):
+        #emb_dim = x.shape[-1] // 2 # Use half as means, half as logvars.
+        #means = x[..., :emb_dim]
+        #logvars = x[..., emb_dim:]
+
+        #Wwe actually wanna do mean and STD on the batch axis?
+
+
+        #we start as b hw 8, go to b hw 4, with mean and std over those.
+
+        if not self.train:
+            result_dict = dict()
+            return x, result_dict
+        else:
+            #Previous run is mean over axis 0..
+            means = jnp.mean(x, axis = [1,2,3])
+            stds = jnp.std(x, axis = [1,2,3])
+
+            noise = jax.random.normal(self.make_rng("noise"), means.shape)
+
+            logvars = .5 * jnp.log(stds)
+
+            z = means + stds * noise
+            #We just... don't need to return Z for this, but instead we return X
+            #This is the denoising version
+            kl_loss = -0.5 * jnp.mean(1 + logvars - means**2 - jnp.exp(logvars))
+            result_dict = dict(quantizer_loss=kl_loss)
+            return x, result_dict
+
+        
+class FSQuantizer(nn.Module):
+    config: ml_collections.ConfigDict
+    train: bool
+
+    @nn.compact
+    def __call__(self, x):
+        assert self.config['fsq_levels'] % 2 == 1, "FSQ levels must be odd."
+        z = jnp.tanh(x) # [-1, 1]
+        z = z * (self.config['fsq_levels']-1) / 2 # [-fsq_levels/2, fsq_levels/2]
+        zhat = jnp.round(z) # e.g. [-2, -1, 0, 1, 2]
+        quantized = z + jax.lax.stop_gradient(zhat - z)
+        quantized = quantized / (self.config['fsq_levels'] // 2) # [-1, 1], but quantized.
+        result_dict = dict()
+
+        # Diagnostics for codebook usage.
+        zhat_scaled = zhat + self.config['fsq_levels'] // 2
+        basis = jnp.concatenate((jnp.array([1]), jnp.cumprod(jnp.array([self.config['fsq_levels']] * (x.shape[-1]-1))))).astype(jnp.uint32)
+        idx = (zhat_scaled * basis).sum(axis=-1).astype(jnp.uint32)
+        idx_flat = idx.reshape(-1)
+        usage = jnp.bincount(idx_flat, length=self.config['fsq_levels']**x.shape[-1])
+
+        result_dict.update({
+            "z_ids": zhat,
+            'usage': usage
+        })
+        return quantized, result_dict
+
+class VQVAE(nn.Module):
+    """VQVAE model."""
+    config: ml_collections.ConfigDict
+    train: bool
+
+    def setup(self):
+        """VQVAE setup."""
+        if self.config['quantizer_type'] == 'vq':
+            self.quantizer = VectorQuantizer(config=self.config, train=self.train)
+        elif self.config['quantizer_type'] == 'kl':
+            self.quantizer = KLQuantizer(config=self.config, train=self.train)
+        elif self.config['quantizer_type'] == 'fsq':
+            self.quantizer = FSQuantizer(config=self.config, train=self.train)
+        elif self.config['quantizer_type'] == 'ae':
+            self.quantizer = AEQuantizer(config=self.config, train=self.train)
+        elif self.config["quantizer_type"] == "kl_two":
+            self.quantizer = KLQuantizerTwo(config=self.config, train=self.train)
+        self.encoder = Encoder(config=self.config)
+        self.decoder = Decoder(config=self.config)
+
+    def encode(self, image):
+        encoded_feature = self.encoder(image)
+        quantized, result_dict = self.quantizer(encoded_feature)
+        print("After quant", quantized.shape)
+        return quantized, result_dict
+
+    def decode(self, z_vectors):
+        print("z_vectors shape", z_vectors.shape)
+        reconstructed = self.decoder(z_vectors)
+        return reconstructed
+
+    def decode_from_indices(self, z_ids):
+        z_vectors = self.quantizer.decode_ids(z_ids)
+        reconstructed_image = self.decode(z_vectors)
+        return reconstructed_image
+
+    def encode_to_indices(self, image):
+        encoded_feature = self.encoder(image)
+        _, result_dict = self.quantizer(encoded_feature)
+        ids = result_dict["z_ids"]
+        return ids
+
+    def __call__(self, input_dict):
+        quantized, result_dict = self.encode(input_dict)
+        outputs = self.decoder(quantized)
+        return outputs, result_dict