Upload folder using huggingface_hub

whiten/calc_means.py

@@ -0,0 +1,438 @@
+from typing import Any
+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
+
+from utils.wandb import setup_wandb, default_wandb_config
+from utils.train_state import TrainStateEma
+from utils.checkpoint import Checkpoint
+from utils.stable_vae import StableVAE
+from utils.sharding import create_sharding, all_gather
+from utils.datasets import get_dataset
+from model import DiT
+from helper_eval import eval_model
+from helper_inference import do_inference
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string('dataset_name', 'imagenet256', 'Environment name.')
+flags.DEFINE_string('load_dir', None, 'Logging dir (if not None, save params).')
+flags.DEFINE_string('save_dir', './checkpoints/', 'Logging dir (if not None, save params).')
+flags.DEFINE_string('fid_stats', None, 'FID stats file.')
+flags.DEFINE_integer('seed', 10, 'Random seed.') # Must be the same across all processes.
+flags.DEFINE_integer('log_interval', 1000, 'Logging interval.')
+flags.DEFINE_integer('eval_interval', 1000000, 'Eval interval.')
+flags.DEFINE_integer('save_interval', 10000, 'Save interval.')
+flags.DEFINE_integer('batch_size', 512, 'Mini batch size.')
+flags.DEFINE_integer('max_steps', int(810_000), 'Number of training steps.')
+flags.DEFINE_integer('debug_overfit', 0, 'Debug overfitting.')
+flags.DEFINE_string('mode', 'train', 'train or inference.')
+
+model_config = ml_collections.ConfigDict({
+    'lr': 0.0001,
+    'beta1': 0.9,
+    'beta2': 0.999,
+    'weight_decay': 0.1,
+    'use_cosine': 0,
+    'warmup': 0,
+    'dropout': 0.0,
+    'hidden_size': 768, # change this!
+    'patch_size': 2, # change this!
+    'depth': 12, # change this!
+    'num_heads': 12, # change this!
+    'mlp_ratio': 4, # change this!
+    'class_dropout_prob': 0.1,
+    'num_classes': 1000,
+    'denoise_timesteps': 128,
+    'cfg_scale': 4.0,
+    'target_update_rate': 0.999,
+    'use_ema': 0,
+    'use_stable_vae': 1,
+    'sharding': 'dp', # dp or fsdp.
+    't_sampling': 'discrete-dt',
+    'dt_sampling': 'uniform',
+    'bootstrap_cfg': 1,
+    'bootstrap_every': 8, # Make sure its a divisor of batch size.
+    'bootstrap_ema': 1,
+    'bootstrap_dt_bias': 0,
+    'train_type': 'shortcut' # or naive.
+})
+
+
+#wandb_config = default_wandb_config()
+#wandb_config.update({
+#    'project': 'shortcut',
+#    'name': 'shortcut_{dataset_name}',
+#})
+
+#config_flags.DEFINE_config_dict('wandb', wandb_config, lock_config=False)
+config_flags.DEFINE_config_dict('model', model_config, lock_config=False)
+    
+##############################################
+## 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()
+    print("Device count", device_count)
+    print("Global device count", global_device_count)
+    local_batch_size = FLAGS.batch_size // (global_device_count // device_count)
+    print("Global Batch: ", FLAGS.batch_size)
+    print("Node Batch: ", local_batch_size)
+    print("Device Batch:", local_batch_size // device_count)
+
+        
+    dataset = get_dataset(FLAGS.dataset_name, local_batch_size, True, FLAGS.debug_overfit)
+    dataset_valid = get_dataset(FLAGS.dataset_name, local_batch_size, False, FLAGS.debug_overfit)
+    example_obs, example_labels = next(dataset)
+    example_obs = example_obs[:1]
+    example_obs_shape = example_obs.shape
+
+    if FLAGS.model.use_stable_vae:
+        vae = StableVAE.create()
+        if 'latent' in FLAGS.dataset_name:
+            example_obs = example_obs[:, :, :, example_obs.shape[-1] // 2:]
+            example_obs_shape = example_obs.shape
+        else:
+            example_obs = vae.encode(jax.random.PRNGKey(0), example_obs)
+        example_obs_shape = example_obs.shape
+        vae_rng = jax.random.PRNGKey(42)
+        vae_encode = jax.jit(vae.encode)
+        vae_decode = jax.jit(vae.decode)
+
+    if FLAGS.fid_stats is not None:
+        from utils.fid import get_fid_network, fid_from_stats
+        get_fid_activations = get_fid_network() 
+        truth_fid_stats = np.load(FLAGS.fid_stats)
+    else:
+        get_fid_activations = None
+        truth_fid_stats = None
+
+    ###################################
+    # Creating Model and put on devices.
+    ###################################
+    FLAGS.model.image_channels = example_obs_shape[-1]
+    FLAGS.model.image_size = example_obs_shape[1]
+    dit_args = {
+        'patch_size': FLAGS.model['patch_size'],
+        'hidden_size': FLAGS.model['hidden_size'],
+        'depth': FLAGS.model['depth'],
+        'num_heads': FLAGS.model['num_heads'],
+        'mlp_ratio': FLAGS.model['mlp_ratio'],
+        'out_channels': example_obs_shape[-1],
+        'class_dropout_prob': FLAGS.model['class_dropout_prob'],
+        'num_classes': FLAGS.model['num_classes'],
+        'dropout': FLAGS.model['dropout'],
+        'ignore_dt': False if (FLAGS.model['train_type'] in ('shortcut', 'livereflow')) else True,
+    }
+    model_def = DiT(**dit_args)
+    tabulate_fn = flax.linen.tabulate(model_def, jax.random.PRNGKey(0))
+    print(tabulate_fn(example_obs, jnp.zeros((1,)), jnp.zeros((1,)), jnp.zeros((1,), dtype=jnp.int32)))
+
+    if FLAGS.model.use_cosine:
+        lr_schedule = optax.warmup_cosine_decay_schedule(0.0, FLAGS.model['lr'], FLAGS.model['warmup'], FLAGS.max_steps)
+    elif FLAGS.model.warmup > 0:
+        lr_schedule = optax.linear_schedule(0.0, FLAGS.model['lr'], FLAGS.model['warmup'])
+    else:
+        lr_schedule = lambda x: FLAGS.model['lr']
+    adam = optax.adamw(learning_rate=lr_schedule, b1=FLAGS.model['beta1'], b2=FLAGS.model['beta2'], weight_decay=FLAGS.model['weight_decay'])
+    tx = optax.chain(adam)
+    
+    start_step = 1
+    
+    def log_param_shapes(params, label=""):
+        flat = flax.traverse_util.flatten_dict(params)
+
+        squeezed_flat = {k: jnp.squeeze(v, axis = 0) for k, v in flat.items() if v.shape[0] == 1}
+        print(f"\n{label} parameter shapes:")
+        for k, v in flat.items():
+            print(f"{k}: {v.shape}")
+        return flax.traverse_util.unflatten_dict(squeezed_flat)
+
+
+    def init(rng):
+        param_key, dropout_key, dropout2_key = jax.random.split(rng, 3)
+        example_t = jnp.zeros((1,))
+        example_dt = jnp.zeros((1,))
+        example_label = jnp.zeros((1,), dtype=jnp.int32)
+        example_obs = jnp.zeros(example_obs_shape)
+        model_rngs = {'params': param_key, 'label_dropout': dropout_key, 'dropout': dropout2_key}
+        params = model_def.init(model_rngs, example_obs, example_t, example_dt, example_label)['params']
+        opt_state = tx.init(params)
+            
+        ts = TrainStateEma.create(model_def, params, rng=rng, tx=tx, opt_state=opt_state)
+
+    
+        if FLAGS.load_dir is not None:
+
+            cp = Checkpoint(FLAGS.load_dir)
+            train_state_load = cp.load_as_dict()["train_state"]
+            
+            log_param_shapes(ts.params)
+            flat = log_param_shapes(train_state_load["params"])
+            flat_ema = log_param_shapes(train_state_load["params_ema"])
+            flat_mu = log_param_shapes(train_state_load["opt_state"][0][0].mu)
+            flat_nu = log_param_shapes(train_state_load["opt_state"][0][0].nu)
+
+            from optax import ScaleByAdamState
+            opt_state = train_state_load["opt_state"]
+            new_state = ScaleByAdamState(
+                opt_state[0][0].count,
+                mu=flat_mu,
+                nu=flat_nu
+            )
+            opt_state = list(opt_state)
+            opt_state[0] = list(opt_state[0])
+            opt_state[0][0] = new_state
+
+            opt_state[0] = tuple(opt_state[0])
+            opt_state = tuple(opt_state)
+
+            train_state_load = TrainStateEma.create(model_def, params = flat, rng = rng, tx = tx, opt_state=opt_state)
+            
+
+            #Need to replace EMA because we have a separate ema
+            log_param_shapes(train_state_load.params)
+
+            train_state_load = train_state_load.replace(params_ema = flat_ema)
+            start_step = train_state_load.step
+
+            ts = train_state_load
+
+        return ts
+    
+    rng = jax.random.PRNGKey(FLAGS.seed)
+    train_state_shape = jax.eval_shape(init, rng)
+
+    data_sharding, train_state_sharding, no_shard, shard_data, global_to_local = create_sharding(FLAGS.model.sharding, train_state_shape)
+    train_state = jax.jit(init, out_shardings=train_state_sharding)(rng)
+    jax.debug.visualize_array_sharding(train_state.params['FinalLayer_0']['Dense_0']['kernel'])
+    jax.debug.visualize_array_sharding(train_state.params['TimestepEmbedder_1']['Dense_0']['kernel'])
+    jax.experimental.multihost_utils.assert_equal(train_state.params['TimestepEmbedder_1']['Dense_0']['kernel'])
+
+
+    if FLAGS.model.train_type == 'progressive' or FLAGS.model.train_type == 'consistency-distillation':
+        train_state_teacher = jax.jit(lambda x : x, out_shardings=train_state_sharding)(train_state)
+    else:
+        train_state_teacher = None
+
+    visualize_labels = example_labels
+    visualize_labels = shard_data(visualize_labels)
+    visualize_labels = jax.experimental.multihost_utils.process_allgather(visualize_labels)
+    imagenet_labels = open('data/imagenet_labels.txt').read().splitlines()
+
+    ###################################
+    # Update Function
+    ###################################
+
+    @partial(jax.jit, out_shardings=(train_state_sharding, no_shard))
+    def update(train_state, train_state_teacher, images, labels, force_t=-1, force_dt=-1):
+        new_rng, targets_key, dropout_key, perm_key = jax.random.split(train_state.rng, 4)
+        info = {}
+
+        id_perm = jax.random.permutation(perm_key, images.shape[0])
+        images = images[id_perm]
+        labels = labels[id_perm]
+        images = jax.lax.with_sharding_constraint(images, data_sharding)
+        labels = jax.lax.with_sharding_constraint(labels, data_sharding)
+
+        if FLAGS.model['cfg_scale'] == 0: # For unconditional generation.
+            labels = jnp.ones(labels.shape[0], dtype=jnp.int32) * FLAGS.model['num_classes']
+
+        if FLAGS.model['train_type'] == 'naive':
+            from baselines.targets_naive import get_targets
+            x_t, v_t, t, dt_base, labels, info = get_targets(FLAGS, targets_key, train_state, images, labels, force_t, force_dt)
+        elif FLAGS.model['train_type'] == 'shortcut':
+            from targets_shortcut import get_targets
+            x_t, v_t, t, dt_base, labels, info = get_targets(FLAGS, targets_key, train_state, images, labels, force_t, force_dt)
+        elif FLAGS.model['train_type'] == 'progressive':
+            from baselines.targets_progressive import get_targets
+            x_t, v_t, t, dt_base, labels, info = get_targets(FLAGS, targets_key, train_state, train_state_teacher, images, labels, force_t, force_dt)
+        elif FLAGS.model['train_type'] == 'consistency-distillation':
+            from baselines.targets_consistency_distillation import get_targets
+            x_t, v_t, t, dt_base, labels, info = get_targets(FLAGS, targets_key, train_state, train_state_teacher, images, labels, force_t, force_dt)
+        elif FLAGS.model['train_type'] == 'consistency':
+            from baselines.targets_consistency_training import get_targets
+            x_t, v_t, t, dt_base, labels, info = get_targets(FLAGS, targets_key, train_state, images, labels, force_t, force_dt)
+        elif FLAGS.model['train_type'] == 'livereflow':
+            from baselines.targets_livereflow import get_targets
+            x_t, v_t, t, dt_base, labels, info = get_targets(FLAGS, targets_key, train_state, images, labels, force_t, force_dt)
+
+        def loss_fn(grad_params):
+            v_prime, logvars, activations = train_state.call_model(x_t, t, dt_base, labels, train=True, rngs={'dropout': dropout_key}, params=grad_params, return_activations=True)
+            mse_v = jnp.mean((v_prime - v_t) ** 2, axis=(1, 2, 3))
+            loss = jnp.mean(mse_v)
+
+            if True:#cosine direction velocity
+                cos_loss = 1-optax.cosine_distance(v_prime, v_t, axis = 3, epsilon = 1e-5)
+                cos_v = jnp.mean(cos_loss, axis = [1,2])
+                cos_loss = cos_v.mean()
+
+            info = {
+                'loss': loss,
+                'v_magnitude_prime': jnp.sqrt(jnp.mean(jnp.square(v_prime))),
+                **{'activations/' + k : jnp.sqrt(jnp.mean(jnp.square(v))) for k, v in activations.items()},
+                'cosine_loss': cos_loss,
+            }
+
+            if FLAGS.model['train_type'] == 'shortcut' or FLAGS.model['train_type'] == 'livereflow':
+                bootstrap_size = FLAGS.batch_size // FLAGS.model['bootstrap_every']
+                info['loss_flow'] = jnp.mean(mse_v[bootstrap_size:])
+                info['loss_bootstrap'] = jnp.mean(mse_v[:bootstrap_size])
+                info['cosine_loss_flow'] = jnp.mean(cos_v[bootstrap_size:])
+                info['cosine_loss_boostrap'] = jnp.mean(cos_v[:bootstrap_size])
+            if True:
+                loss = loss + cos_loss
+
+            return loss, info
+        
+        grads, new_info = jax.grad(loss_fn, has_aux=True)(train_state.params)
+        info = {**info, **new_info}
+        updates, new_opt_state = train_state.tx.update(grads, train_state.opt_state, train_state.params)
+        new_params = optax.apply_updates(train_state.params, updates)
+
+        info['grad_norm'] = optax.global_norm(grads)
+        info['update_norm'] = optax.global_norm(updates)
+        info['param_norm'] = optax.global_norm(new_params)
+        info['lr'] = lr_schedule(train_state.step)
+
+        train_state = train_state.replace(rng=new_rng, step=train_state.step + 1, params=new_params, opt_state=new_opt_state)
+        train_state = train_state.update_ema(FLAGS.model['target_update_rate'])
+        return train_state, info
+    
+
+    ###################################
+    # Train Loop
+   ###################################
+    global_mean = None    
+    class_means = {}
+    total = 1281167
+    #Do we need to do global means more often?
+    print("starting this shit")
+    i = 0
+    cpus = jax.devices("cpu")
+    images = []
+    for i in range(0, int(total/512)):
+        print(i)
+        i += 1
+
+        batch_images, batch_labels = shard_data(*next(dataset))
+        vae_rng, vae_key = jax.random.split(vae_rng)
+        batch_images = vae_encode(vae_key, batch_images)
+        #print(batch_images.shape)#512x32x32x4
+        if global_mean == None:
+            global_mean = batch_images.mean(axis = 0)/total
+        else:
+            global_mean += batch_images.mean(axis = 0)/total
+
+        for key, bimage in zip(batch_labels, batch_images):            
+            key = str(int(key))
+            if key in class_means.keys():
+                class_means[key] = class_means[key] + bimage/total
+            else:
+                class_means[key] = np.asarray(bimage/total)
+        
+#        z = jax.device_put(batch_images, cpus[0])
+        images.append(batch_images)
+
+    images = jnp.asarray(images)
+    #maybe just save images and exit?
+    np.savez("images.npz", images)
+    exit()
+
+    """
+    #Get per channel stats.
+    batch_shape = images.shape[0] * images.shape[1]
+    H, W = images.shape[2], images.shape[3]
+    images_white = jnp.zeros(images.shape)
+    stats = []
+    
+    for c in range(images.shape[-1]):
+        
+        x = images[:,:,:,:,c].reshape(batch_shape, -1)#Get h*w by batch
+        mean = x.mean(axis = 0, keepdims = True)
+        x_centered = x - mean
+        cov = x_centered.T @ x_centered / (batch_shape - 1)  # shape: (H*W, H*W)
+        U, S, _ = jnp.linalg.svd(cov, full_matrices=False)
+        S_inv_root = jnp.diag(1.0 / jnp.sqrt(S + 1e-5))
+        zca = U @ S_inv_root @ U.T
+
+        x_whitened = (zca @ x_centered.T).T  # shape: (B, H*W)
+        images_whitened[:, :, :, :, c] = x_whitened.view(B, H, W)
+        
+        stats.append((mean, zca))  # Save stats for unwhitening
+    
+    #Now we need to save stats?
+    np.savez("stats.npz", stats)
+    """
+
+#    jnp.save("global_mean", global_mean)
+#    np.savez("classes.npz", **class_means)
+    exit()
+    for i in tqdm.tqdm(range(1 + start_step, FLAGS.max_steps + 1 + start_step),
+                       smoothing=0.1,
+                       dynamic_ncols=True):
+        
+        # Sample data.
+        if not FLAGS.debug_overfit or i == 1:
+            batch_images, batch_labels = shard_data(*next(dataset))
+            if FLAGS.model.use_stable_vae and 'latent' not in FLAGS.dataset_name:
+                vae_rng, vae_key = jax.random.split(vae_rng)
+                batch_images = vae_encode(vae_key, batch_images)
+
+        
+
+        # Train update.
+        train_state, update_info = update(train_state, train_state_teacher, batch_images, batch_labels)
+
+        if i % FLAGS.log_interval == 0 or i == 1:
+            update_info = jax.device_get(update_info)
+            update_info = jax.tree_map(lambda x: np.array(x), update_info)
+            update_info = jax.tree_map(lambda x: x.mean(), update_info)
+            train_metrics = {f'training/{k}': v for k, v in update_info.items()}
+
+            valid_images, valid_labels = shard_data(*next(dataset_valid))
+            if FLAGS.model.use_stable_vae and 'latent' not in FLAGS.dataset_name:
+                valid_images = vae_encode(vae_rng, valid_images)
+            _, valid_update_info = update(train_state, train_state_teacher, valid_images, valid_labels)
+            valid_update_info = jax.device_get(valid_update_info)
+            valid_update_info = jax.tree_map(lambda x: x.mean(), valid_update_info)
+            train_metrics['training/loss_valid'] = valid_update_info['loss']
+            train_metrics['training/loss_cosine'] = valid_update_info['cosine_loss']
+
+            if jax.process_index() == 0:
+                wandb.log(train_metrics, step=i)
+
+        if FLAGS.model['train_type'] == 'progressive':
+            num_sections = np.log2(FLAGS.model['denoise_timesteps']).astype(jnp.int32)
+            if i % (FLAGS.max_steps // num_sections) == 0:
+                train_state_teacher = jax.jit(lambda x : x, out_shardings=train_state_sharding)(train_state)
+
+        if i % FLAGS.eval_interval == 0:
+            eval_model(FLAGS, train_state, train_state_teacher, i, dataset, dataset_valid, shard_data, vae_encode, vae_decode, update,
+                       get_fid_activations, imagenet_labels, visualize_labels, 
+                       fid_from_stats, truth_fid_stats)
+
+        if i % FLAGS.save_interval == 0 and FLAGS.save_dir is not None:
+            train_state_gather = jax.experimental.multihost_utils.process_allgather(train_state)
+            if jax.process_index() == 0:
+                cp = Checkpoint(FLAGS.save_dir+str(train_state_gather.step+1), parallel=False)
+                cp.train_state = train_state_gather
+                cp.save()
+                del cp
+            del train_state_gather
+
+if __name__ == '__main__':
+    app.run(main)
+

whiten/stable_vae.py

@@ -0,0 +1,101 @@
+
+from functools import partial, cached_property
+
+import jax
+from diffusers import FlaxAutoencoderKL
+from einops import rearrange
+from flax import struct
+
+from jaxtyping import Array, PyTree, Key, Float, Shaped, Int, UInt8, jaxtyped
+from typeguard import typechecked
+from functools import partial
+typecheck = partial(jaxtyped, typechecker=typechecked)
+
+import jax.numpy as jnp
+
+import pickle
+def load_stats(path="stats.pkl"):
+    with open(path, "rb") as f:
+        return pickle.load(f)
+
+try:
+    stats = load_stats()#mean, zca
+except:
+    pass
+@struct.dataclass
+class StableVAE:
+    params: PyTree[Float[Array, "..."]]
+    module: FlaxAutoencoderKL = struct.field(pytree_node=False)
+
+    @classmethod
+    def create(cls) -> "VAE":
+        # module, params = FlaxAutoencoderKL.from_pretrained(
+        #     "stabilityai/stable-diffusion-xl-base-1.0", subfolder="vae"
+        # )
+        module, params = FlaxAutoencoderKL.from_pretrained(
+            "pcuenq/sd-vae-ft-mse-flax"
+        )
+        params = jax.device_get(params)
+        return cls(
+            params=params,
+            module=module,
+        )
+
+    @partial(jax.jit, static_argnames="scale")
+    def encode(
+        self, key: Key[Array, ""], images: Float[Array, "b h w 3"], scale: bool = True
+    ) -> Float[Array, "b lh lw 4"]:
+        images = rearrange(images, "b h w c -> b c h w")
+        latents = self.module.apply(
+            {"params": self.params}, images, method=self.module.encode
+        ).latent_dist.sample(key)
+        
+#        return latents
+        B, H, W, C = latents.shape
+        latents_whitened = jnp.zeros(latents.shape)
+        for c in range(C):
+            x = latents[:, :, :, c].reshape(B, -1)#We are channels last probably
+            mean, zca = stats[c]
+
+            x_centered = x - mean
+            x_whitened = (zca @ x_centered.T).T
+            latents_whitened = latents_whitened.at[:, :, :, c].set(x_whitened.reshape(B, H, W))
+
+#        if scale:
+#            latents *= self.module.config.scaling_factor
+        return latents_whitened
+
+    @partial(jax.jit, static_argnames="scale")
+    def decode(
+        self, latents: Float[Array, "b lh lw 4"], scale: bool = True
+    ) -> Float[Array, "b h w 3"]:
+        #if scale:
+        #    latents /= self.module.config.scaling_factor
+        
+#        latents = latents.reshape(1)#256x32x32x4
+        #Not sure these latents are correct shape, but whatever
+        B, H, W, C = latents.shape
+        latents_unwhitened = jnp.zeros(latents.shape)
+        
+        for c in range(C):
+            x = latents[:, :, :, c].reshape(B, -1)
+            mean, zca = stats[c]
+            zca_inv = jnp.linalg.inv(zca)
+
+            x_unwhitened = (zca_inv @ x.T).T + mean
+            latents_unwhitened = latents_unwhitened.at[:, : ,: ,c].set(x_unwhitened.reshape(B,H,W))
+            
+        latents = latents_unwhitened
+        #I don't think you need to sample to encode and sample to decode.
+        images = self.module.apply(
+            {"params": self.params}, latents, method=self.module.decode
+        ).sample
+
+        # convert to channels-last
+        #This actually just converts to channels FIRST, which is needed to convert to image
+        images = rearrange(images, "b c h w -> b h w c")
+        return images
+
+    @cached_property
+    def downscale_factor(self) -> int:
+        return 2 ** (len(self.module.block_out_channels) - 1)

whiten/whiten.py

@@ -0,0 +1,41 @@
+import numpy as np
+import jax
+import jax.numpy as jnp
+import gc
+images = np.load("images.npz")["arr_0"]
+print(images.shape)
+if True:
+    batch_shape = images.shape[0] * images.shape[1]
+    H, W = images.shape[2], images.shape[3]
+    images_white = jnp.zeros(images.shape)
+    stats = []
+    
+    for c in range(images.shape[-1]):
+        print(c)
+        x = images[:,:,:,:,c].reshape(batch_shape, -1)#Get h*w by batch
+        print(x.shape)
+        mean = x.mean(axis = 0, keepdims = True)
+        print(mean.shape)#It's like 1024, because it's reshaped.
+        x = x - mean
+        cov = x.T @ x / (batch_shape - 1)  # shape: (H*W, H*W)
+        U, S, _ = jnp.linalg.svd(cov, full_matrices=False)
+        S_inv_root = jnp.diag(1.0 / jnp.sqrt(S + 1e-5))
+        zca = U @ S_inv_root @ U.T
+        del cov
+        del U
+        del S
+        del _
+        del S_inv_root
+        x = (zca @ x.T).T  # shape: (B, H*W)
+        gc.collect()
+        #images_whitened[:, :, :, :, c] = x.reshape(images.shape[0], images.shape[1],images.shape[2], images.shape[3])
+        
+        #only need mean and zca..
+        stats.append((mean, zca))  # Save stats for unwhitening
+
+    #Now we need to save stats?
+#    np.savez("stats.npz", stats)
+    import pickle
+    with open("stats.pkl","wb") as f:
+        pickle.dump(stats, f)
+