Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,164 @@
+wandb/
+train_logs/
+slurms/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 MedARC
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,12 @@
+# MindEyeV2
+
+In-progress -- this repo is under active development in the MedARC discord server (feel free to join us and help develop MindEyeV2!)
+
+1. Download all of https://huggingface.co/datasets/pscotti/mindeyev2 and place them in a folder. You will need to specify the path to this folder as "data_path" variable.
+
+2. Run setup.sh to install a new "fmri" conda environment.
+
+3. Activate the conda environment with "conda activate fmri"
+
+4. Run Train.ipynb or Train.py (they are the same code)
+

src/Train-with-memory-cat.py

@@ -0,0 +1,1056 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[3]:
+
+
+# # Code to convert this notebook to .py if you want to run it via command line or with Slurm
+# from subprocess import call
+# command = "jupyter nbconvert Train-with-memory-cat.ipynb --to python"
+# call(command,shell=True)
+
+
+# # Import packages & functions
+
+# In[4]:
+
+
+import os
+import sys
+import json
+import argparse
+import numpy as np
+import math
+from einops import rearrange
+import time
+import random
+import h5py
+from tqdm import tqdm
+
+import webdataset as wds
+import gc
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+
+from accelerate import Accelerator, DeepSpeedPlugin
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+
+global_batch_size = 128 #128
+
+
+# In[5]:
+
+
+### Multi-GPU config ###
+local_rank = os.getenv('RANK')
+if local_rank is None: 
+    local_rank = 0
+else:
+    local_rank = int(local_rank)
+print("LOCAL RANK ", local_rank)  
+
+num_devices = torch.cuda.device_count()
+if num_devices==0: num_devices = 1
+
+accelerator = Accelerator(split_batches=False)
+
+### UNCOMMENT BELOW STUFF TO USE DEEPSPEED (also comment out the immediately above "accelerator = " line) ###
+
+# if num_devices <= 1 and utils.is_interactive():
+#     # can emulate a distributed environment for deepspeed to work in jupyter notebook
+#     os.environ["MASTER_ADDR"] = "localhost"
+#     os.environ["MASTER_PORT"] = str(np.random.randint(10000)+9000)
+#     os.environ["RANK"] = "0"
+#     os.environ["LOCAL_RANK"] = "0"
+#     os.environ["WORLD_SIZE"] = "1"
+#     os.environ["GLOBAL_BATCH_SIZE"] = str(global_batch_size) # set this to your batch size!
+#     global_batch_size = os.environ["GLOBAL_BATCH_SIZE"]
+
+# # alter the deepspeed config according to your global and local batch size
+# if local_rank == 0:
+#     with open('deepspeed_config_stage2.json', 'r') as file:
+#         config = json.load(file)
+#     config['train_batch_size'] = int(os.environ["GLOBAL_BATCH_SIZE"])
+#     config['train_micro_batch_size_per_gpu'] = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices
+#     with open('deepspeed_config_stage2.json', 'w') as file:
+#         json.dump(config, file)
+# else:
+#     # give some time for the local_rank=0 gpu to prep new deepspeed config file
+#     time.sleep(10)
+# deepspeed_plugin = DeepSpeedPlugin("deepspeed_config_stage2.json")
+# accelerator = Accelerator(split_batches=False, deepspeed_plugin=deepspeed_plugin)
+
+
+# In[6]:
+
+
+print("PID of this process =",os.getpid())
+device = accelerator.device
+print("device:",device)
+num_workers = num_devices
+print(accelerator.state)
+world_size = accelerator.state.num_processes
+distributed = not accelerator.state.distributed_type == 'NO'
+print("distributed =",distributed, "num_devices =", num_devices, "local rank =", local_rank, "world size =", world_size)
+print = accelerator.print # only print if local_rank=0
+
+
+# # Configurations
+
+# In[7]:
+
+
+# if running this interactively, can specify jupyter_args here for argparser to use
+if utils.is_interactive():
+    # Example use
+    jupyter_args = f"--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset \
+                    --model_name=test \
+                    --subj=1 --batch_size={global_batch_size} --n_samples_save=0 \
+                    --max_lr=3e-5 --mixup_pct=.66 --num_epochs=12 --ckpt_interval=999 --no-use_image_aug"
+
+    jupyter_args = jupyter_args.split()
+    print(jupyter_args)
+    
+    from IPython.display import clear_output # function to clear print outputs in cell
+    get_ipython().run_line_magic('load_ext', 'autoreload')
+    # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions
+    get_ipython().run_line_magic('autoreload', '2')
+
+
+# In[8]:
+
+
+parser = argparse.ArgumentParser(description="Model Training Configuration")
+parser.add_argument(
+    "--model_name", type=str, default="testing",
+    help="name of model, used for ckpt saving and wandb logging (if enabled)",
+)
+parser.add_argument(
+    "--data_path", type=str, default="/fsx/proj-fmri/shared/natural-scenes-dataset",
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,5,7],
+)
+parser.add_argument(
+    "--batch_size", type=int, default=32,
+    help="Batch size can be increased by 10x if only training v2c and not diffusion diffuser",
+)
+parser.add_argument(
+    "--wandb_log",action=argparse.BooleanOptionalAction,default=False,
+    help="whether to log to wandb",
+)
+parser.add_argument(
+    "--resume_from_ckpt",action=argparse.BooleanOptionalAction,default=False,
+    help="if not using wandb and want to resume from a ckpt",
+)
+parser.add_argument(
+    "--wandb_project",type=str,default="stability",
+    help="wandb project name",
+)
+parser.add_argument(
+    "--mixup_pct",type=float,default=.33,
+    help="proportion of way through training when to switch from BiMixCo to SoftCLIP",
+)
+parser.add_argument(
+    "--use_image_aug",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to use image augmentation",
+)
+parser.add_argument(
+    "--num_epochs",type=int,default=240,
+    help="number of epochs of training",
+)
+parser.add_argument(
+    "--lr_scheduler_type",type=str,default='cycle',choices=['cycle','linear'],
+)
+parser.add_argument(
+    "--ckpt_saving",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--ckpt_interval",type=int,default=5,
+    help="save backup ckpt and reconstruct every x epochs",
+)
+parser.add_argument(
+    "--seed",type=int,default=42,
+)
+parser.add_argument(
+    "--max_lr",type=float,default=3e-4,
+)
+parser.add_argument(
+    "--n_samples_save",type=int,default=0,choices=[0,1],
+    help="Number of reconstructions for monitoring progress, 0 will speed up training",
+)
+
+if utils.is_interactive():
+    args = parser.parse_args(jupyter_args)
+else:
+    args = parser.parse_args()
+
+# create global variables without the args prefix
+for attribute_name in vars(args).keys():
+    globals()[attribute_name] = getattr(args, attribute_name)
+
+print("global batch_size", batch_size)
+batch_size = int(batch_size / num_devices)
+print("batch_size", batch_size)
+
+
+# In[9]:
+
+
+outdir = os.path.abspath(f'../train_mem_logs/{model_name}')
+if not os.path.exists(outdir):
+    os.makedirs(outdir,exist_ok=True)
+if use_image_aug:
+    import kornia
+    from kornia.augmentation.container import AugmentationSequential
+    img_augment = AugmentationSequential(
+        kornia.augmentation.RandomResizedCrop((224,224), (0.6,1), p=0.3),
+        kornia.augmentation.Resize((224, 224)),
+        kornia.augmentation.RandomHorizontalFlip(p=0.3),
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.3),
+        kornia.augmentation.RandomGrayscale(p=0.3),
+        same_on_batch=False,
+        data_keys=["input"],
+    )
+
+
+# # Prep data, models, and dataloaders
+
+# ## Dataloader
+
+# In[10]:
+
+
+if subj==1:
+    num_train = 24958
+    num_test = 2770
+test_batch_size = num_test
+
+def my_split_by_node(urls): return urls
+    
+train_url = f"{data_path}/wds/subj0{subj}/train/" + "{0..36}.tar"
+print(train_url)
+
+train_data = wds.WebDataset(train_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+print(test_url)
+
+test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+test_dl = torch.utils.data.DataLoader(test_data, batch_size=test_batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+
+# ### check dataloaders are working
+
+# In[9]:
+
+
+# test_indices = []
+# test_images = []
+# for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+#     test_indices = np.append(test_indices, behav[:,0,5].numpy())
+#     test_images = np.append(test_images, behav[:,0,0].numpy())
+# test_indices = test_indices.astype(np.int16)
+# print(test_i, (test_i+1) * test_batch_size, len(test_indices))
+# print("---\n")
+
+# train_indices = []
+# train_images = []
+# for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+#     train_indices = np.append(train_indices, behav[:,0,5].long().numpy())
+#     train_images = np.append(train_images, behav[:,0,0].numpy())
+# train_indices = train_indices.astype(np.int16)
+# print(train_i, (train_i+1) * batch_size, len(train_indices))
+
+# # train_images = np.hstack((train_images, test_images))
+# # print("WARNING: ADDED TEST IMAGES TO TRAIN IMAGES")
+
+
+# ## Load data and images
+
+# In[ ]:
+
+
+# load betas
+f = h5py.File(f'{data_path}/betas_all_subj0{subj}.hdf5', 'r')
+voxels = f['betas'][:]
+print(f"subj0{subj} betas loaded into memory")
+voxels = torch.Tensor(voxels).to("cpu").half()
+if subj==1:
+    voxels = torch.hstack((voxels, torch.zeros((len(voxels), 5))))
+print("voxels", voxels.shape)
+num_voxels = voxels.shape[-1]
+
+# load orig images
+f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+images = f['images'][:]
+images = torch.Tensor(images).to("cpu").half()
+print("images", images.shape)
+
+
+# ## Load models
+
+# ### CLIP image embeddings  model
+
+# In[ ]:
+
+
+from models import Clipper
+clip_model = Clipper("ViT-L/14", device=torch.device(f"cuda:{local_rank}"), hidden_state=True, norm_embs=True)
+
+clip_seq_dim = 257
+clip_emb_dim = 768
+hidden_dim = 4096
+
+
+# ### SD VAE (blurry images)
+
+# In[ ]:
+
+
+from diffusers import AutoencoderKL
+autoenc = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, cache_dir="/fsx/proj-fmri/shared/cache")
+# autoenc.load_state_dict(torch.load('../train_logs/sdxl_vae_normed/best.pth')["model_state_dict"])
+autoenc.eval()
+autoenc.requires_grad_(False)
+autoenc.to(device)
+utils.count_params(autoenc)
+
+
+# ### MindEye modules
+
+# In[13]:
+
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+model
+
+
+# In[14]:
+
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer
+    def __init__(self, input_size, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linear = torch.nn.Linear(input_size, out_features)
+    def forward(self, x):
+        return self.linear(x)
+        
+model.ridge = RidgeRegression(voxels.shape[1], out_features=hidden_dim)
+utils.count_params(model.ridge)
+utils.count_params(model)
+
+b = torch.randn((2,1,voxels.shape[1]))
+print(b.shape, model.ridge(b).shape)
+
+
+# In[22]:
+
+
+from functools import partial
+from diffusers.models.vae import Decoder
+class BrainNetwork(nn.Module):
+    def __init__(self, out_dim=768, in_dim=15724, clip_size=768, h=4096, n_blocks=4, norm_type='ln', act_first=False, drop=.15, blurry_dim=16):
+        super().__init__()
+        self.blurry_dim = blurry_dim
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        self.lin0 = nn.Linear(in_dim, h)
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, out_dim, bias=True)
+        self.blin1 = nn.Linear(out_dim, blurry_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.clip_size = clip_size
+        self.clip_proj = nn.Sequential(
+            nn.LayerNorm(clip_size),
+            nn.GELU(),
+            nn.Linear(clip_size, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, clip_size)
+        )
+        self.upsampler = Decoder(
+                in_channels=64,
+                out_channels=4,
+                up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D"],
+                block_out_channels=[64, 128, 256],
+                layers_per_block=1,
+            )
+        
+    def forward(self, x):
+        x = self.lin0(x)
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        b = self.blin1(x)
+        b = self.upsampler(b.reshape(len(b), -1, 7, 7))
+        c = self.clip_proj(x.reshape(len(x), -1, self.clip_size))
+        return c, b
+
+model.backbone = BrainNetwork(h=2048, in_dim=hidden_dim*2, clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, blurry_dim=64*7*7) 
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+b = torch.randn((2,8192))
+print(b.shape)
+clip_, blur_ = model.backbone(b)
+print(clip_.shape, blur_.shape)
+
+
+# In[23]:
+
+
+# memory model
+
+from timm.layers.mlp import Mlp
+
+class MemoryEncoder(nn.Module):
+    def __init__(self, in_dim=15279, out_dim=768, h=4096, num_past_voxels=15, embedding_time_dim = 512, n_blocks=4, norm_type='ln', act_first=False, drop=.15):
+        super().__init__()
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        self.out_dim = out_dim
+        self.embedding_time = nn.Embedding(num_past_voxels, embedding_time_dim)
+        self.final_input_dim = in_dim + embedding_time_dim
+        self.lin0 = nn.Linear(self.final_input_dim, h)
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, out_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.num_past_voxels = num_past_voxels
+        self.embedding_time_dim = embedding_time_dim
+        self.memory = nn.Parameter(torch.randn((self.num_past_voxels, self.embedding_time_dim)))
+
+
+    def forward(self, x, time):
+        time = time.long()
+        time = self.embedding_time(time)
+        x = torch.cat((x, time), dim=-1)
+        x = self.lin0(x)
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        return x
+    
+# # test the memory encoder
+# memory_encoder = MemoryEncoder(in_dim=voxels.shape[1], out_dim=hidden_dim, num_past_voxels=15, embedding_time_dim=512)
+
+# device = torch.device("cpu")
+# memory_encoder.to(device)
+
+# # count params
+# total_parameters = 0
+# for parameter in memory_encoder.parameters():
+#     total_parameters += parameter.numel()
+
+# rand_input = torch.randn((2, 15279)).to(device)
+# rand_time = torch.randint(0, 15, (2,)).to(device)
+# print(rand_input.shape, rand_time.shape)
+# memory_encoder(rand_input, rand_time).shape
+
+class MemoryCompressor(nn.Module):
+    def __init__(self, in_dim=768, num_past = 15, output_dim=768, h=4096, n_blocks=4, norm_type='ln', act_first=False, drop=.15):
+        super().__init__()
+        self.num_past = num_past
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        self.final_input_dim = in_dim * num_past
+        self.lin0 = nn.Linear(self.final_input_dim, h)
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, output_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.num_past = num_past
+        self.output_dim = output_dim
+
+    def forward(self, x):
+        # x is (batch_size, num_past, in_dim)
+        x = x.reshape(len(x), -1)
+        x = self.lin0(x)
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        return x
+    
+# # test the memory compressor
+# memory_compressor = MemoryCompressor(in_dim=768, num_past=15, output_dim=768)
+
+# device = torch.device("cpu")
+# memory_compressor.to(device)
+
+# # count params
+# total_parameters = 0
+# for parameter in memory_compressor.parameters():
+#     total_parameters += parameter.numel()
+
+# rand_input = torch.randn((2, 15, 768)).to(device)
+# print(rand_input.shape)
+# memory_compressor(rand_input).shape
+
+model.memory_encoder = MemoryEncoder(in_dim=voxels.shape[1], out_dim=4096, num_past_voxels=15, embedding_time_dim=512)
+model.memory_compressor = MemoryCompressor(in_dim=model.memory_encoder.out_dim, num_past=15, output_dim=4096)
+
+utils.count_params(model.memory_encoder)
+utils.count_params(model.memory_compressor)
+utils.count_params(model)
+
+
+
+# In[24]:
+
+
+no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+opt_grouped_parameters = [
+    {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
+    {'params': [p for n, p in model.memory_encoder.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.memory_compressor.named_parameters()], 'weight_decay': 1e-2},
+]
+
+optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr, betas=(0.9, 0.95))
+
+if lr_scheduler_type == 'linear':
+    lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+        optimizer,
+        total_iters=int(num_epochs*(num_train*num_devices//batch_size)),
+        last_epoch=-1
+    )
+elif lr_scheduler_type == 'cycle':
+    total_steps=int(num_epochs*(num_train*num_devices//batch_size))
+    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer, 
+        max_lr=max_lr,
+        total_steps=total_steps,
+        final_div_factor=1000,
+        last_epoch=-1, pct_start=2/num_epochs
+    )
+    
+def save_ckpt(tag):    
+    ckpt_path = outdir+f'/{tag}.pth'
+    print(f'saving {ckpt_path}',flush=True)
+    unwrapped_model = accelerator.unwrap_model(model)
+    try:
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': unwrapped_model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'lr_scheduler': lr_scheduler.state_dict(),
+            'train_losses': losses,
+            'test_losses': test_losses,
+            'lrs': lrs,
+            }, ckpt_path)
+    except:
+        print("Couldn't save... moving on to prevent crashing.")
+    del unwrapped_model
+        
+print("\nDone with model preparations!")
+utils.count_params(model)
+
+
+# In[18]:
+
+
+
+
+
+# # Weights and Biases
+
+# In[25]:
+
+
+# params for wandb
+wandb_log = True
+if local_rank==0 and wandb_log: # only use main process for wandb logging
+    import wandb
+    
+    wandb_project = 'stability'
+    wandb_run = model_name
+    wandb_notes = ''
+    
+    print(f"wandb {wandb_project} run {wandb_run}")
+    wandb.login(host='https://stability.wandb.io')#, relogin=True)
+    wandb_config = {
+      "model_name": model_name,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "use_image_aug": use_image_aug,
+      "max_lr": max_lr,
+      "lr_scheduler_type": lr_scheduler_type,
+      "mixup_pct": mixup_pct,
+      "num_train": num_train,
+      "num_test": num_test,
+      "seed": seed,
+      "distributed": distributed,
+      "num_devices": num_devices,
+      "world_size": world_size,
+    }
+    print("wandb_config:\n",wandb_config)
+    if False: # wandb_auto_resume
+        print("wandb_id:",model_name)
+        wandb.init(
+            id = model_name,
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+            resume="allow",
+        )
+    else:
+        wandb.init(
+            project=wandb_project,
+            name=model_name,
+            config=wandb_config,
+            notes=wandb_notes,
+        )
+else:
+    wandb_log = False
+
+
+# # More custom functions
+
+# In[26]:
+
+
+# using the same preprocessing as was used in MindEye + BrainDiffuser
+pixcorr_preprocess = transforms.Compose([
+    transforms.Resize(425, interpolation=transforms.InterpolationMode.BILINEAR),
+])
+def pixcorr(images,brains):
+    # Flatten images while keeping the batch dimension
+    all_images_flattened = pixcorr_preprocess(images).reshape(len(images), -1)
+    all_brain_recons_flattened = pixcorr_preprocess(brains).view(len(brains), -1)
+    corrmean = torch.diag(utils.batchwise_pearson_correlation(all_images_flattened, all_brain_recons_flattened)).mean()
+    return corrmean
+
+
+# # Main
+
+# In[27]:
+
+
+epoch = 0
+losses, test_losses, lrs = [], [], []
+best_test_loss = 1e9
+soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))
+
+# Optionally resume from checkpoint #
+if resume_from_ckpt:
+    print("\n---resuming from last.pth ckpt---\n")
+    try:
+        checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+    except:
+        print('last.pth failed... trying last_backup.pth')
+        checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+    epoch = checkpoint['epoch']
+    print("Epoch",epoch)
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])
+    del checkpoint
+elif wandb_log:
+    if wandb.run.resumed:
+        print("\n---resuming from last.pth ckpt---\n")
+        try:
+            checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+        except:
+            print('last.pth failed... trying last_backup.pth')
+            checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+        epoch = checkpoint['epoch']
+        print("Epoch",epoch)
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+        diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])
+        del checkpoint
+torch.cuda.empty_cache()
+
+
+# In[28]:
+
+
+model, optimizer, train_dl, test_dl, lr_scheduler = accelerator.prepare(
+model, optimizer, train_dl, test_dl, lr_scheduler
+)
+
+
+# In[29]:
+no_more = False
+
+
+print(f"{model_name} starting with epoch {epoch} / {num_epochs}")
+progress_bar = tqdm(range(0,num_epochs), ncols=1200, disable=(local_rank!=0))
+test_image, test_voxel = None, None
+mse = nn.MSELoss()
+for epoch in progress_bar:
+    model.train()
+    
+    fwd_percent_correct = 0.
+    bwd_percent_correct = 0.
+    test_fwd_percent_correct = 0.
+    test_bwd_percent_correct = 0.
+
+    loss_clip_total = 0.
+    loss_blurry_total = 0.
+    test_loss_clip_total = 0.
+    test_loss_blurry_total = 0.
+
+    blurry_pixcorr = 0.
+    test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1
+    
+    for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+        #if epoch == 0 or epoch == 1:
+        #    break
+        with torch.cuda.amp.autocast():
+            optimizer.zero_grad()
+
+            voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+            
+            image = images[behav[:,0,0].cpu().long()].to(device).float()
+
+            past_15_voxels = voxels[past_behav[:,:,5].cpu().long()].to(device) # batch_size, 15, 15279
+            past_15_times = torch.Tensor([i for i in range(15)]).to(device) # 15
+
+            blurry_image_enc = autoenc.encode(image).latent_dist.mode()
+            
+            if use_image_aug: image = img_augment(image)
+            
+            clip_target = clip_model.embed_image(image)
+            assert not torch.any(torch.isnan(clip_target))
+  
+            if epoch < int(mixup_pct * num_epochs):
+                voxel, perm, betas, select = utils.mixco(voxel)
+
+            # reshape past voxels to be (batch_size * 15, 15279)
+            past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+            past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+            past_15_times = past_15_times.reshape(-1)
+            
+            #print(past_15_voxels.shape, past_15_times.shape)
+
+            embeds_past_voxels = model.memory_encoder(past_15_voxels, past_15_times)
+            #print(embeds_past_voxels.shape)
+            embeds_past_voxels = embeds_past_voxels.reshape(voxel.shape[0], 15, -1)
+            #print(embeds_past_voxels.shape)
+            information_past_voxels = model.memory_compressor(embeds_past_voxels)
+
+
+            voxel_ridge = torch.cat([model.ridge(voxel), information_past_voxels], dim=-1)
+            
+            clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+            
+            clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+            clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+            if epoch < int(mixup_pct * num_epochs):                
+                loss_clip = utils.mixco_nce(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006, 
+                    perm=perm, betas=betas, select=select)
+            else:
+                epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=epoch_temp)
+
+            loss_blurry = mse(blurry_image_enc_, blurry_image_enc) 
+
+            loss_clip_total += loss_clip.item()
+            loss_blurry_total += loss_blurry.item()
+
+            loss = loss_blurry + loss_clip
+
+            
+            print('voxel', voxel.shape, voxel)
+            print('image', image.shape, image)
+            print('past_15_voxels', past_15_voxels.shape)
+            print('past_15_times', past_15_times.shape)
+            print('embeds_past_voxels', embeds_past_voxels.shape)
+            print('information_past_voxels', information_past_voxels.shape)
+            print('voxel_ridge', voxel_ridge.shape)
+            print('clip_target', clip_target.shape)
+            print('clip_voxels', clip_voxels.shape)
+            print('clip_voxels_norm', clip_voxels_norm.shape)
+            print('clip_target_norm', clip_target_norm.shape)
+            print('loss_clip_total', loss_clip_total)
+            print('loss_blurry_total', loss_blurry_total)
+            print('loss_clip', loss_clip)
+            print('loss_blurry', loss_blurry)
+            print('train_i', train_i)
+            print('epoch', epoch)
+            
+            e = utils.check_loss(loss)
+            if e == 'NaN loss':
+                no_more = True
+                print("saving ckpt")
+                save_ckpt(f'error')
+                # save all the tensors
+                torch.save({
+                    'voxel': voxel,
+                    'image': image,
+                    'past_15_voxels': past_15_voxels,
+                    'past_15_times': past_15_times,
+                    'embeds_past_voxels': embeds_past_voxels,
+                    'information_past_voxels': information_past_voxels,
+                    'voxel_ridge': voxel_ridge,
+                    'blurry_image_enc': blurry_image_enc,
+                    'clip_target': clip_target,
+                    'clip_voxels': clip_voxels,
+                    'blurry_image_enc_': blurry_image_enc_,
+                    'clip_voxels_norm': clip_voxels_norm,
+                    'clip_target_norm': clip_target_norm,
+                    'loss': loss,
+                    'loss_clip': loss_clip,
+                    'loss_blurry': loss_blurry,
+                    'loss_clip_total': loss_clip_total,
+                    'loss_blurry_total': loss_blurry_total,
+                    'train_i': train_i,
+                    'epoch': epoch,
+                    'model_state_dict': model.state_dict(),
+                }, outdir+f'/error_tensors.pth')
+
+                print("Error with loss here")
+                print('voxel', voxel.shape, voxel)
+                print('image', image.shape, image)
+                print('past_15_voxels', past_15_voxels.shape, past_15_voxels)
+                print('past_15_times', past_15_times.shape, past_15_times)
+                print('embeds_past_voxels', embeds_past_voxels.shape, embeds_past_voxels)
+                print('information_past_voxels', information_past_voxels.shape, information_past_voxels)
+                print('voxel_ridge', voxel_ridge.shape, voxel_ridge)
+                print('clip_target', clip_target.shape, clip_target)
+                print('clip_voxels', clip_voxels.shape, clip_voxels)
+                print('clip_voxels_norm', clip_voxels_norm.shape, clip_voxels_norm)
+                print('clip_target_norm', clip_target_norm.shape, clip_target_norm)
+                print('loss_clip_total', loss_clip_total)
+                print('loss_blurry_total', loss_blurry_total)
+                print('loss_clip', loss_clip)
+                print('loss_blurry', loss_blurry)
+                print('train_i', train_i)
+                print('epoch', epoch)
+
+
+
+                
+
+            accelerator.backward(loss)
+            optimizer.step()
+    
+            losses.append(loss.item())
+            lrs.append(optimizer.param_groups[0]['lr'])
+    
+            # forward and backward top 1 accuracy        
+            labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+            fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)
+            bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)
+
+            with torch.no_grad():
+                # only doing pixcorr eval on a subset (8) of the samples per batch because its costly & slow to compute autoenc.decode()
+                random_samps = np.random.choice(np.arange(len(voxel)), size=2, replace=False)
+                blurry_recon_images = autoenc.decode(blurry_image_enc_[random_samps]).sample.clamp(0,1)
+                blurry_pixcorr += pixcorr(image[random_samps], blurry_recon_images)
+
+            if lr_scheduler_type is not None:
+                lr_scheduler.step()
+    
+    model.eval()
+    for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+        print('test')
+        with torch.cuda.amp.autocast():
+            with torch.no_grad():   
+                # all test samples should be loaded per batch such that test_i should never exceed 0
+                if len(behav) != num_test: print("!",len(behav),num_test)
+
+                
+                ## Average same-image repeats ##
+                if test_image is None:
+                    voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+                    
+                    image = behav[:,0,0].cpu().long()
+                    
+                    unique_image, sort_indices = torch.unique(image, return_inverse=True)
+                    for im in unique_image:
+                        locs = torch.where(im == image)[0]
+                        if test_image is None:
+                            test_image = images[im][None]
+                            test_voxel = torch.mean(voxel[locs],axis=0)[None]
+                        else:
+                            test_image = torch.vstack((test_image, images[im][None]))
+                            test_voxel = torch.vstack((test_voxel, torch.mean(voxel[locs],axis=0)[None]))
+    
+                # sample of batch_size
+                random_indices = torch.arange(len(test_voxel))[:batch_size] #torch.randperm(len(test_voxel))[:300]
+                voxel = test_voxel[random_indices].to(device)
+                image = test_image[random_indices].to(device)
+
+                current_past_behav = past_behav[random_indices]
+
+                past_15_voxels = voxels[current_past_behav[:,:,5].cpu().long()].to(device) # batch_size, 15, 15279
+                past_15_times = torch.Tensor([i for i in range(15)]).to(device) # 15
+
+                assert len(image) == batch_size
+    
+                blurry_image_enc = autoenc.encode(image).latent_dist.mode()
+        
+                clip_target = clip_model.embed_image(image.float())
+
+                past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+                past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+                past_15_times = past_15_times.reshape(-1)
+                
+                print(past_15_voxels.shape, past_15_times.shape)
+
+                embeds_past_voxels = model.memory_encoder(past_15_voxels, past_15_times)
+                embeds_past_voxels = embeds_past_voxels.reshape(batch_size, 15, -1)
+                information_past_voxels = model.memory_compressor(embeds_past_voxels)
+
+    
+                voxel_ridge = torch.cat([model.ridge(voxel), information_past_voxels], dim=-1)
+                
+                clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+                
+                clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+                clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+        
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006)
+
+                loss_blurry = mse(blurry_image_enc_, blurry_image_enc)
+                
+                loss = loss_blurry + loss_clip
+                
+                utils.check_loss(loss)
+        
+                test_losses.append(loss.item())
+        
+                # forward and backward top 1 accuracy        
+                labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+                test_fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)
+                test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)
+
+                # halving the batch size because the decoder is computationally heavy
+                blurry_recon_images = autoenc.decode(blurry_image_enc_[:len(voxel)//2]).sample.clamp(0,1)
+                blurry_recon_images = torch.vstack((blurry_recon_images, autoenc.decode(blurry_image_enc_[len(voxel)//2:]).sample.clamp(0,1)))
+                test_blurry_pixcorr += pixcorr(image, blurry_recon_images)
+
+                # transform blurry recon latents to images and plot it
+                fig, axes = plt.subplots(1, 4, figsize=(8, 4))
+                axes[0].imshow(utils.torch_to_Image(image[[0]]))
+                axes[1].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[0]]).sample.clamp(0,1)))
+                axes[2].imshow(utils.torch_to_Image(image[[1]]))
+                axes[3].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[1]]).sample.clamp(0,1)))
+                axes[0].axis('off'); axes[1].axis('off'); axes[2].axis('off'); axes[3].axis('off')
+                plt.show()
+
+    if local_rank==0:      
+        # if utils.is_interactive(): clear_output(wait=True)
+        assert (test_i+1) == 1
+        logs = {"train/loss": np.mean(losses[-(train_i+1):]),
+            "test/loss": np.mean(test_losses[-(test_i+1):]),
+            "train/lr": lrs[-1],
+            "train/num_steps": len(losses),
+            "test/num_steps": len(test_losses),
+            "train/fwd_pct_correct": fwd_percent_correct / (train_i + 1),
+            "train/bwd_pct_correct": bwd_percent_correct / (train_i + 1),
+            "test/test_fwd_pct_correct": test_fwd_percent_correct / (test_i + 1),
+            "test/test_bwd_pct_correct": test_bwd_percent_correct / (test_i + 1),
+            "train/loss_clip_total": loss_clip_total / (train_i + 1),
+            "train/loss_blurry_total": loss_blurry_total / (train_i + 1),
+            "test/loss_clip_total": test_loss_clip_total / (test_i + 1),
+            "test/loss_blurry_total": test_loss_blurry_total / (test_i + 1),
+            "train/blurry_pixcorr": blurry_pixcorr / (train_i + 1),
+            "test/blurry_pixcorr": test_blurry_pixcorr / (test_i + 1),
+            }
+        progress_bar.set_postfix(**logs)
+
+        # Save model checkpoint and reconstruct
+        if no_more:
+            print("not writing more")
+        if ((epoch % ckpt_interval == 0) and (not no_more)):
+            if not utils.is_interactive():
+                save_ckpt(f'last')
+                
+        if wandb_log: wandb.log(logs)
+
+    # wait for other GPUs to catch up if needed
+    accelerator.wait_for_everyone()
+    torch.cuda.empty_cache()
+    gc.collect()
+
+print("\n===Finished!===\n")
+if ckpt_saving:
+    save_ckpt(f'last')
+if not utils.is_interactive():
+    sys.exit(0)
+
+
+# In[ ]:
+
+
+plt.plot(losses)
+plt.show()
+plt.plot(test_losses)
+plt.show()
+

src/Train-with-memory-rr-dropout.py

@@ -0,0 +1,1040 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[ ]:
+
+
+# # Code to convert this notebook to .py if you want to run it via command line or with Slurm
+#from subprocess import call
+#command = "jupyter nbconvert Train-with-memory-rr.ipynb --to python"
+#call(command,shell=True)
+
+
+# In[24]:
+
+
+#get_ipython().system('nvidia-smi')
+
+
+# # Import packages & functions
+
+# In[3]:
+
+
+import os
+import sys
+import json
+import argparse
+import numpy as np
+import math
+from einops import rearrange
+import time
+import random
+import h5py
+from tqdm import tqdm
+
+import webdataset as wds
+import gc
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+
+from accelerate import Accelerator, DeepSpeedPlugin
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+
+global_batch_size = 512 #128
+
+
+# In[4]:
+
+
+### Multi-GPU config ###
+local_rank = os.getenv('RANK')
+if local_rank is None: 
+    local_rank = 0
+else:
+    local_rank = int(local_rank)
+print("LOCAL RANK ", local_rank)  
+
+num_devices = torch.cuda.device_count()
+if num_devices==0: num_devices = 1
+
+accelerator = Accelerator(split_batches=False)
+
+### UNCOMMENT BELOW STUFF TO USE DEEPSPEED (also comment out the immediately above "accelerator = " line) ###
+
+# if num_devices <= 1 and utils.is_interactive():
+#     # can emulate a distributed environment for deepspeed to work in jupyter notebook
+#     os.environ["MASTER_ADDR"] = "localhost"
+#     os.environ["MASTER_PORT"] = str(np.random.randint(10000)+9000)
+#     os.environ["RANK"] = "0"
+#     os.environ["LOCAL_RANK"] = "0"
+#     os.environ["WORLD_SIZE"] = "1"
+#     os.environ["GLOBAL_BATCH_SIZE"] = str(global_batch_size) # set this to your batch size!
+#     global_batch_size = os.environ["GLOBAL_BATCH_SIZE"]
+
+# # alter the deepspeed config according to your global and local batch size
+# if local_rank == 0:
+#     with open('deepspeed_config_stage2.json', 'r') as file:
+#         config = json.load(file)
+#     config['train_batch_size'] = int(os.environ["GLOBAL_BATCH_SIZE"])
+#     config['train_micro_batch_size_per_gpu'] = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices
+#     with open('deepspeed_config_stage2.json', 'w') as file:
+#         json.dump(config, file)
+# else:
+#     # give some time for the local_rank=0 gpu to prep new deepspeed config file
+#     time.sleep(10)
+# deepspeed_plugin = DeepSpeedPlugin("deepspeed_config_stage2.json")
+# accelerator = Accelerator(split_batches=False, deepspeed_plugin=deepspeed_plugin)
+
+
+# In[5]:
+
+
+print("PID of this process =",os.getpid())
+device = accelerator.device
+print("device:",device)
+num_workers = num_devices
+print(accelerator.state)
+world_size = accelerator.state.num_processes
+distributed = not accelerator.state.distributed_type == 'NO'
+print("distributed =",distributed, "num_devices =", num_devices, "local rank =", local_rank, "world size =", world_size)
+print = accelerator.print # only print if local_rank=0
+
+
+# # Configurations
+
+# In[6]:
+
+
+# if running this interactively, can specify jupyter_args here for argparser to use
+if utils.is_interactive():
+    # Example use
+    jupyter_args = f"--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset \
+                    --model_name=test \
+                    --subj=1 --batch_size={global_batch_size} --n_samples_save=0 \
+                    --max_lr=3e-5 --mixup_pct=.66 --num_epochs=12 --ckpt_interval=999 --no-use_image_aug"
+
+    jupyter_args = jupyter_args.split()
+    print(jupyter_args)
+    
+    from IPython.display import clear_output # function to clear print outputs in cell
+    get_ipython().run_line_magic('load_ext', 'autoreload')
+    # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions
+    get_ipython().run_line_magic('autoreload', '2')
+
+
+# In[7]:
+
+
+parser = argparse.ArgumentParser(description="Model Training Configuration")
+parser.add_argument(
+    "--model_name", type=str, default="memory_cat_rr",
+    help="name of model, used for ckpt saving and wandb logging (if enabled)",
+)
+parser.add_argument(
+    "--data_path", type=str, default="/fsx/proj-fmri/shared/natural-scenes-dataset",
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,5,7],
+)
+parser.add_argument(
+    "--batch_size", type=int, default=32,
+    help="Batch size can be increased by 10x if only training v2c and not diffusion diffuser",
+)
+parser.add_argument(
+    "--wandb_log",action=argparse.BooleanOptionalAction,default=False,
+    help="whether to log to wandb",
+)
+parser.add_argument(
+    "--resume_from_ckpt",action=argparse.BooleanOptionalAction,default=False,
+    help="if not using wandb and want to resume from a ckpt",
+)
+parser.add_argument(
+    "--wandb_project",type=str,default="stability",
+    help="wandb project name",
+)
+parser.add_argument(
+    "--mixup_pct",type=float,default=.33,
+    help="proportion of way through training when to switch from BiMixCo to SoftCLIP",
+)
+parser.add_argument(
+    "--use_image_aug",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to use image augmentation",
+)
+parser.add_argument(
+    "--num_epochs",type=int,default=240,
+    help="number of epochs of training",
+)
+parser.add_argument(
+    "--lr_scheduler_type",type=str,default='cycle',choices=['cycle','linear'],
+)
+parser.add_argument(
+    "--ckpt_saving",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--ckpt_interval",type=int,default=5,
+    help="save backup ckpt and reconstruct every x epochs",
+)
+parser.add_argument(
+    "--seed",type=int,default=42,
+)
+parser.add_argument(
+    "--max_lr",type=float,default=3e-4,
+)
+parser.add_argument(
+    "--n_samples_save",type=int,default=0,choices=[0,1],
+    help="Number of reconstructions for monitoring progress, 0 will speed up training",
+)
+
+if utils.is_interactive():
+    args = parser.parse_args(jupyter_args)
+else:
+    args = parser.parse_args()
+
+# create global variables without the args prefix
+for attribute_name in vars(args).keys():
+    globals()[attribute_name] = getattr(args, attribute_name)
+
+print("global batch_size", batch_size)
+batch_size = int(batch_size / num_devices)
+print("batch_size", batch_size)
+
+
+# In[8]:
+
+
+outdir = os.path.abspath(f'../train_mem_logs/{model_name}')
+if not os.path.exists(outdir):
+    os.makedirs(outdir,exist_ok=True)
+if use_image_aug:
+    import kornia
+    from kornia.augmentation.container import AugmentationSequential
+    img_augment = AugmentationSequential(
+        kornia.augmentation.RandomResizedCrop((224,224), (0.6,1), p=0.3),
+        kornia.augmentation.Resize((224, 224)),
+        kornia.augmentation.RandomHorizontalFlip(p=0.3),
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.3),
+        kornia.augmentation.RandomGrayscale(p=0.3),
+        same_on_batch=False,
+        data_keys=["input"],
+    )
+
+
+# # Prep data, models, and dataloaders
+
+# ## Dataloader
+
+# In[9]:
+
+
+if subj==1:
+    num_train = 24958
+    num_test = 2770
+test_batch_size = num_test
+
+def my_split_by_node(urls): return urls
+    
+train_url = f"{data_path}/wds/subj0{subj}/train/" + "{0..36}.tar"
+print(train_url)
+
+train_data = wds.WebDataset(train_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+print(test_url)
+
+test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+test_dl = torch.utils.data.DataLoader(test_data, batch_size=test_batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+
+# ### check dataloaders are working
+
+# In[10]:
+
+
+# test_indices = []
+# test_images = []
+# for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+#     test_indices = np.append(test_indices, behav[:,0,5].numpy())
+#     test_images = np.append(test_images, behav[:,0,0].numpy())
+# test_indices = test_indices.astype(np.int16)
+# print(test_i, (test_i+1) * test_batch_size, len(test_indices))
+# print("---\n")
+
+# train_indices = []
+# train_images = []
+# for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+#     train_indices = np.append(train_indices, behav[:,0,5].long().numpy())
+#     train_images = np.append(train_images, behav[:,0,0].numpy())
+# train_indices = train_indices.astype(np.int16)
+# print(train_i, (train_i+1) * batch_size, len(train_indices))
+
+# # train_images = np.hstack((train_images, test_images))
+# # print("WARNING: ADDED TEST IMAGES TO TRAIN IMAGES")
+
+
+# ## Load data and images
+
+# In[11]:
+
+
+# load betas
+f = h5py.File(f'{data_path}/betas_all_subj0{subj}.hdf5', 'r')
+voxels = f['betas'][:]
+print(f"subj0{subj} betas loaded into memory")
+voxels = torch.Tensor(voxels).to("cpu").half()
+if subj==1:
+    voxels = torch.hstack((voxels, torch.zeros((len(voxels), 5))))
+print("voxels", voxels.shape)
+num_voxels = voxels.shape[-1]
+
+# load orig images
+f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+images = f['images'][:]
+images = torch.Tensor(images).to("cpu").half()
+print("images", images.shape)
+
+
+# ## Load models
+
+# ### CLIP image embeddings  model
+
+# In[12]:
+
+
+from models import Clipper
+clip_model = Clipper("ViT-L/14", device=torch.device(f"cuda:{local_rank}"), hidden_state=True, norm_embs=True)
+
+clip_seq_dim = 257
+clip_emb_dim = 768
+hidden_dim = 4096
+
+
+# ### SD VAE (blurry images)
+
+# In[13]:
+
+
+from diffusers import AutoencoderKL
+autoenc = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, cache_dir="/fsx/proj-fmri/shared/cache")
+# autoenc.load_state_dict(torch.load('../train_logs/sdxl_vae_normed/best.pth')["model_state_dict"])
+autoenc.eval()
+autoenc.requires_grad_(False)
+autoenc.to(device)
+utils.count_params(autoenc)
+
+
+# ### MindEye modules
+
+# In[14]:
+
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+model
+
+
+# In[15]:
+
+
+time_embedding_dim = 512
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer
+    def __init__(self, input_size, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linear = torch.nn.Linear(input_size, out_features)
+    def forward(self, x):
+        return self.linear(x)
+        
+model.ridge = RidgeRegression(voxels.shape[1] + time_embedding_dim, out_features=hidden_dim)
+utils.count_params(model.ridge)
+utils.count_params(model)
+
+b = torch.randn((2,1,voxels.shape[1]))
+time_emb_test = torch.randn((2,1,time_embedding_dim))
+print(b.shape, model.ridge(torch.cat((b,time_emb_test),dim=-1)).shape)
+
+
+# In[16]:
+
+
+from functools import partial
+from diffusers.models.vae import Decoder
+class BrainNetwork(nn.Module):
+    def __init__(self, out_dim=768, in_dim=15724, clip_size=768, h=4096, n_blocks=4, norm_type='ln', act_first=False, drop=.75, blurry_dim=16):
+        super().__init__()
+        self.blurry_dim = blurry_dim
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        self.lin0 = nn.Linear(in_dim, h)
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, out_dim, bias=True)
+        self.blin1 = nn.Linear(out_dim, blurry_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.clip_size = clip_size
+        self.clip_proj = nn.Sequential(
+            nn.LayerNorm(clip_size),
+            nn.GELU(),
+            nn.Linear(clip_size, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, clip_size)
+        )
+        self.upsampler = Decoder(
+                in_channels=64,
+                out_channels=4,
+                up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D"],
+                block_out_channels=[64, 128, 256],
+                layers_per_block=1,
+            )
+        
+    def forward(self, x):
+        x = self.lin0(x)
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        b = self.blin1(x)
+        b = self.upsampler(b.reshape(len(b), -1, 7, 7))
+        c = self.clip_proj(x.reshape(len(x), -1, self.clip_size))
+        return c, b
+
+model.backbone = BrainNetwork(h=2048, in_dim=hidden_dim*2, clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, blurry_dim=64*7*7) 
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+b = torch.randn((2,8192))
+print(b.shape)
+clip_, blur_ = model.backbone(b)
+print(clip_.shape, blur_.shape)
+
+
+# In[17]:
+
+
+# memory model
+
+from timm.layers.mlp import Mlp
+
+
+
+class MemoryDropout(nn.Module):
+    def __init__(self, p):
+        super(MemoryDropout, self).__init__()
+        self.p = p
+
+    def forward(self, x):
+        if self.training:
+            mask = torch.zeros(x.size(0), x.size(1)).bernoulli_(1 - self.p).unsqueeze(2).expand_as(x)
+            mask = mask.to(x.device)
+            x = x * mask / (1 - self.p)
+        return x
+
+memory_dropout_percentage = 0.75
+memory_dropout = MemoryDropout(memory_dropout_percentage)
+
+class MemoryEncoder(nn.Module):
+    def __init__(self, in_dim=15279, out_dim=768, h=4096, num_past_voxels=15, embedding_time_dim = 512, n_blocks=4, norm_type='ln', act_first=False, drop=.85):
+        super().__init__()
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        self.out_dim = out_dim
+        self.embedding_time = nn.Embedding(num_past_voxels, embedding_time_dim)
+        self.final_input_dim = in_dim + embedding_time_dim
+        self.lin0 = nn.Linear(self.final_input_dim, h)
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, out_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.num_past_voxels = num_past_voxels
+        self.embedding_time_dim = embedding_time_dim
+        self.memory = nn.Parameter(torch.randn((self.num_past_voxels, self.embedding_time_dim)))
+
+
+    def forward(self, x, time):
+        time = time.long()
+        time = self.embedding_time(time)
+        x = torch.cat((x, time), dim=-1)
+        x = self.lin0(x)
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        return x
+    
+# # test the memory encoder
+# memory_encoder = MemoryEncoder(in_dim=voxels.shape[1], out_dim=hidden_dim, num_past_voxels=15, embedding_time_dim=512)
+
+# device = torch.device("cpu")
+# memory_encoder.to(device)
+
+# # count params
+# total_parameters = 0
+# for parameter in memory_encoder.parameters():
+#     total_parameters += parameter.numel()
+
+# rand_input = torch.randn((2, 15279)).to(device)
+# rand_time = torch.randint(0, 15, (2,)).to(device)
+# print(rand_input.shape, rand_time.shape)
+# memory_encoder(rand_input, rand_time).shape
+
+class MemoryCompressor(nn.Module):
+    def __init__(self, in_dim=768, num_past = 15, output_dim=768, h=4096, n_blocks=4, norm_type='ln', act_first=False, drop=.75):
+        super().__init__()
+        self.num_past = num_past
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        self.final_input_dim = in_dim * num_past
+        self.lin0 = nn.Linear(self.final_input_dim, h)
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, output_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.num_past = num_past
+        self.output_dim = output_dim
+
+    def forward(self, x):
+        # x is (batch_size, num_past, in_dim)
+        x = x.reshape(len(x), -1)
+        x = self.lin0(x)
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        return x
+    
+# # test the memory compressor
+# memory_compressor = MemoryCompressor(in_dim=768, num_past=15, output_dim=768)
+
+# device = torch.device("cpu")
+# memory_compressor.to(device)
+
+# # count params
+# total_parameters = 0
+# for parameter in memory_compressor.parameters():
+#     total_parameters += parameter.numel()
+
+# rand_input = torch.randn((2, 15, 768)).to(device)
+# print(rand_input.shape)
+# memory_compressor(rand_input).shape
+
+class TimeEmbedding(nn.Module):
+    def __init__(self, embedding_time_dim=512, num_past_voxels=15):
+        super().__init__()
+        self.embedding_time = nn.Embedding(num_past_voxels, embedding_time_dim)
+        self.num_past_voxels = num_past_voxels
+        self.embedding_time_dim = embedding_time_dim
+
+    def forward(self, time):
+        # time is (batch_size,)
+        time = time.long()
+        time = self.embedding_time(time)
+        return time # (batch_size, embedding_time_dim)
+    
+
+#model.memory_encoder = MemoryEncoder(in_dim=voxels.shape[1], out_dim=4096, num_past_voxels=15, embedding_time_dim=512)
+model.time_embedding = TimeEmbedding(embedding_time_dim=512, num_past_voxels=15)
+model.memory_compressor = MemoryCompressor(in_dim=model.ridge.out_features, num_past=15, output_dim=4096)
+
+#utils.count_params(model.memory_encoder)
+utils.count_params(model.memory_compressor)
+utils.count_params(model)
+
+
+
+# In[18]:
+
+
+no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+opt_grouped_parameters = [
+    {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
+    #{'params': [p for n, p in model.memory_encoder.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.memory_compressor.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.time_embedding.named_parameters()], 'weight_decay': 0.0},
+]
+
+optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr, betas=(0.9, 0.95))
+
+if lr_scheduler_type == 'linear':
+    lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+        optimizer,
+        total_iters=int(num_epochs*(num_train*num_devices//batch_size)),
+        last_epoch=-1
+    )
+elif lr_scheduler_type == 'cycle':
+    total_steps=int(num_epochs*(num_train*num_devices//batch_size))
+    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer, 
+        max_lr=max_lr,
+        total_steps=total_steps,
+        final_div_factor=1000,
+        last_epoch=-1, pct_start=2/num_epochs
+    )
+    
+def save_ckpt(tag):    
+    ckpt_path = outdir+f'/{tag}.pth'
+    print(f'saving {ckpt_path}',flush=True)
+    unwrapped_model = accelerator.unwrap_model(model)
+    try:
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': unwrapped_model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'lr_scheduler': lr_scheduler.state_dict(),
+            'train_losses': losses,
+            'test_losses': test_losses,
+            'lrs': lrs,
+            }, ckpt_path)
+    except:
+        print("Couldn't save... moving on to prevent crashing.")
+    del unwrapped_model
+        
+print("\nDone with model preparations!")
+utils.count_params(model)
+
+
+# In[ ]:
+
+
+
+
+
+# # Weights and Biases
+
+# In[19]:
+
+
+# params for wandb
+wandb_log = True
+if local_rank==0 and wandb_log: # only use main process for wandb logging
+    import wandb
+    
+    wandb_project = 'stability'
+    wandb_run = model_name
+    wandb_notes = ''
+    
+    print(f"wandb {wandb_project} run {wandb_run}")
+    wandb.login(host='https://stability.wandb.io')#, relogin=True)
+    wandb_config = {
+      "model_name": model_name,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "use_image_aug": use_image_aug,
+      "max_lr": max_lr,
+      "lr_scheduler_type": lr_scheduler_type,
+      "mixup_pct": mixup_pct,
+      "num_train": num_train,
+      "num_test": num_test,
+      "seed": seed,
+      "distributed": distributed,
+      "num_devices": num_devices,
+      "world_size": world_size,
+    }
+    print("wandb_config:\n",wandb_config)
+    if False: # wandb_auto_resume
+        print("wandb_id:",model_name)
+        wandb.init(
+            id = model_name,
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+            resume="allow",
+        )
+    else:
+        wandb.init(
+            project=wandb_project,
+            name=model_name,
+            config=wandb_config,
+            notes=wandb_notes,
+        )
+else:
+    wandb_log = False
+
+
+# # More custom functions
+
+# In[20]:
+
+
+# using the same preprocessing as was used in MindEye + BrainDiffuser
+pixcorr_preprocess = transforms.Compose([
+    transforms.Resize(425, interpolation=transforms.InterpolationMode.BILINEAR),
+])
+def pixcorr(images,brains):
+    # Flatten images while keeping the batch dimension
+    all_images_flattened = pixcorr_preprocess(images).reshape(len(images), -1)
+    all_brain_recons_flattened = pixcorr_preprocess(brains).view(len(brains), -1)
+    corrmean = torch.diag(utils.batchwise_pearson_correlation(all_images_flattened, all_brain_recons_flattened)).mean()
+    return corrmean
+
+
+# # Main
+
+# In[21]:
+
+
+epoch = 0
+losses, test_losses, lrs = [], [], []
+best_test_loss = 1e9
+soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))
+
+# Optionally resume from checkpoint #
+if resume_from_ckpt:
+    print("\n---resuming from last.pth ckpt---\n")
+    try:
+        checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+    except:
+        print('last.pth failed... trying last_backup.pth')
+        checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+    epoch = checkpoint['epoch']
+    print("Epoch",epoch)
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])
+    del checkpoint
+elif wandb_log:
+    if wandb.run.resumed:
+        print("\n---resuming from last.pth ckpt---\n")
+        try:
+            checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+        except:
+            print('last.pth failed... trying last_backup.pth')
+            checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+        epoch = checkpoint['epoch']
+        print("Epoch",epoch)
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+        diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])
+        del checkpoint
+torch.cuda.empty_cache()
+
+
+# In[22]:
+
+
+model, optimizer, train_dl, test_dl, lr_scheduler = accelerator.prepare(
+model, optimizer, train_dl, test_dl, lr_scheduler
+)
+
+
+# In[23]:
+
+
+print(f"{model_name} starting with epoch {epoch} / {num_epochs}")
+progress_bar = tqdm(range(0,num_epochs), ncols=1200, disable=(local_rank!=0))
+test_image, test_voxel = None, None
+mse = nn.MSELoss()
+for epoch in progress_bar:
+    model.train()
+    
+    fwd_percent_correct = 0.
+    bwd_percent_correct = 0.
+    test_fwd_percent_correct = 0.
+    test_bwd_percent_correct = 0.
+
+    loss_clip_total = 0.
+    loss_blurry_total = 0.
+    test_loss_clip_total = 0.
+    test_loss_blurry_total = 0.
+
+    blurry_pixcorr = 0.
+    test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1
+    
+    for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+        #if epoch == 0 or epoch == 1:
+        #    break
+        with torch.cuda.amp.autocast():
+            optimizer.zero_grad()
+
+            voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+            
+            image = images[behav[:,0,0].cpu().long()].to(device).float()
+
+            past_15_voxels = voxels[past_behav[:,:,5].cpu().long()].to(device) # batch_size, 15, 15279
+            past_15_times = torch.Tensor([i for i in range(15)]).to(device) # 15
+
+            blurry_image_enc = autoenc.encode(image).latent_dist.mode()
+            
+            if use_image_aug: image = img_augment(image)
+            
+            clip_target = clip_model.embed_image(image)
+            assert not torch.any(torch.isnan(clip_target))
+  
+            if epoch < int(mixup_pct * num_epochs):
+                voxel, perm, betas, select = utils.mixco(voxel)
+
+            # reshape past voxels to be (batch_size * 15, 15279)
+            past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+            past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+            past_15_times = past_15_times.reshape(-1)
+            
+            #print(past_15_voxels.shape, past_15_times.shape)
+            time_embeddings = model.time_embedding(past_15_times)
+            past_info_full = torch.cat((past_15_voxels, time_embeddings), dim=-1)
+            embeds_past_voxels = model.ridge(past_info_full)
+            #print(embeds_past_voxels.shape)
+            embeds_past_voxels = embeds_past_voxels.reshape(voxel.shape[0], 15, -1)
+            embeds_past_voxels = memory_dropout(embeds_past_voxels)
+            #print(embeds_past_voxels.shape)
+            information_past_voxels = model.memory_compressor(embeds_past_voxels)
+
+            positional_current_voxel = torch.zeros((voxel.shape[0], time_embeddings.shape[-1])).to(voxel.device)
+            #print(torch.cat((voxel, positional_current_voxel), dim=-1).shape, positional_current_voxel.shape, voxel.shape)
+            voxel_ridge = torch.cat([model.ridge(torch.cat((voxel, positional_current_voxel), dim=-1)), information_past_voxels], dim=-1)
+            
+            clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+            
+            clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+            clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+            if epoch < int(mixup_pct * num_epochs):                
+                loss_clip = utils.mixco_nce(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006, 
+                    perm=perm, betas=betas, select=select)
+            else:
+                epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=epoch_temp)
+
+            loss_blurry = mse(blurry_image_enc_, blurry_image_enc) 
+
+            loss_clip_total += loss_clip.item()
+            loss_blurry_total += loss_blurry.item()
+
+            loss = loss_blurry + loss_clip
+            
+            utils.check_loss(loss)
+
+            accelerator.backward(loss)
+            optimizer.step()
+    
+            losses.append(loss.item())
+            lrs.append(optimizer.param_groups[0]['lr'])
+    
+            # forward and backward top 1 accuracy        
+            labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+            fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)
+            bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)
+
+            with torch.no_grad():
+                # only doing pixcorr eval on a subset (8) of the samples per batch because its costly & slow to compute autoenc.decode()
+                random_samps = np.random.choice(np.arange(len(voxel)), size=2, replace=False)
+                blurry_recon_images = autoenc.decode(blurry_image_enc_[random_samps]).sample.clamp(0,1)
+                blurry_pixcorr += pixcorr(image[random_samps], blurry_recon_images)
+
+            if lr_scheduler_type is not None:
+                lr_scheduler.step()
+    
+    model.eval()
+    for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+        print('test')
+        with torch.cuda.amp.autocast():
+            with torch.no_grad():   
+                # all test samples should be loaded per batch such that test_i should never exceed 0
+                if len(behav) != num_test: print("!",len(behav),num_test)
+
+                
+                ## Average same-image repeats ##
+                if test_image is None:
+                    voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+                    
+                    image = behav[:,0,0].cpu().long()
+                    
+                    unique_image, sort_indices = torch.unique(image, return_inverse=True)
+                    for im in unique_image:
+                        locs = torch.where(im == image)[0]
+                        if test_image is None:
+                            test_image = images[im][None]
+                            test_voxel = torch.mean(voxel[locs],axis=0)[None]
+                        else:
+                            test_image = torch.vstack((test_image, images[im][None]))
+                            test_voxel = torch.vstack((test_voxel, torch.mean(voxel[locs],axis=0)[None]))
+    
+                # sample of batch_size
+                random_indices = torch.arange(len(test_voxel))[:batch_size] #torch.randperm(len(test_voxel))[:300]
+                voxel = test_voxel[random_indices].to(device)
+                image = test_image[random_indices].to(device)
+
+                current_past_behav = past_behav[random_indices]
+
+                past_15_voxels = voxels[current_past_behav[:,:,5].cpu().long()].to(device) # batch_size, 15, 15279
+                past_15_times = torch.Tensor([i for i in range(15)]).to(device) # 15
+
+                assert len(image) == batch_size
+    
+                blurry_image_enc = autoenc.encode(image).latent_dist.mode()
+        
+                clip_target = clip_model.embed_image(image.float())
+
+                past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+                past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+                past_15_times = past_15_times.reshape(-1)
+                
+                print(past_15_voxels.shape, past_15_times.shape)
+
+                #print(past_15_voxels.shape, past_15_times.shape)
+                time_embeddings = model.time_embedding(past_15_times)
+                past_info_full = torch.cat((past_15_voxels, time_embeddings), dim=-1)
+                embeds_past_voxels = model.ridge(past_info_full)
+                #print(embeds_past_voxels.shape)
+                embeds_past_voxels = embeds_past_voxels.reshape(voxel.shape[0], 15, -1)
+                #print(embeds_past_voxels.shape)
+                information_past_voxels = model.memory_compressor(embeds_past_voxels)
+                positional_current_voxel = torch.zeros((voxel.shape[0], time_embeddings.shape[-1])).to(voxel.device)
+    
+                voxel_ridge = torch.cat([model.ridge(torch.cat((voxel, positional_current_voxel), dim=-1)), information_past_voxels], dim=-1)
+                
+                clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+                
+                clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+                clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+        
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006)
+
+                loss_blurry = mse(blurry_image_enc_, blurry_image_enc)
+                
+                loss = loss_blurry + loss_clip
+                
+                utils.check_loss(loss)
+        
+                test_losses.append(loss.item())
+        
+                # forward and backward top 1 accuracy        
+                labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+                test_fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)
+                test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)
+
+                # halving the batch size because the decoder is computationally heavy
+                blurry_recon_images = autoenc.decode(blurry_image_enc_[:len(voxel)//2]).sample.clamp(0,1)
+                blurry_recon_images = torch.vstack((blurry_recon_images, autoenc.decode(blurry_image_enc_[len(voxel)//2:]).sample.clamp(0,1)))
+                test_blurry_pixcorr += pixcorr(image, blurry_recon_images)
+
+                # transform blurry recon latents to images and plot it
+                #fig, axes = plt.subplots(1, 4, figsize=(8, 4))
+                #axes[0].imshow(utils.torch_to_Image(image[[0]]))
+                #axes[1].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[0]]).sample.clamp(0,1)))
+                #axes[2].imshow(utils.torch_to_Image(image[[1]]))
+                #axes[3].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[1]]).sample.clamp(0,1)))
+                #axes[0].axis('off'); axes[1].axis('off'); axes[2].axis('off'); axes[3].axis('off')
+                #plt.show()
+
+                
+    if local_rank==0:      
+        # if utils.is_interactive(): clear_output(wait=True)
+        with torch.cuda.amp.autocast():
+            with torch.no_grad():   
+                wandb.log({"gt": [wandb.Image(utils.torch_to_Image(image[[0]])), wandb.Image(utils.torch_to_Image(image[[1]])) ]})
+                wandb.log({"preds": [wandb.Image(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[0]]).sample.clamp(0,1))), wandb.Image(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[1]]).sample.clamp(0,1))) ]})
+
+        assert (test_i+1) == 1
+        logs = {"train/loss": np.mean(losses[-(train_i+1):]),
+            "test/loss": np.mean(test_losses[-(test_i+1):]),
+            "train/lr": lrs[-1],
+            "train/num_steps": len(losses),
+            "test/num_steps": len(test_losses),
+            "train/fwd_pct_correct": fwd_percent_correct / (train_i + 1),
+            "train/bwd_pct_correct": bwd_percent_correct / (train_i + 1),
+            "test/test_fwd_pct_correct": test_fwd_percent_correct / (test_i + 1),
+            "test/test_bwd_pct_correct": test_bwd_percent_correct / (test_i + 1),
+            "train/loss_clip_total": loss_clip_total / (train_i + 1),
+            "train/loss_blurry_total": loss_blurry_total / (train_i + 1),
+            "test/loss_clip_total": test_loss_clip_total / (test_i + 1),
+            "test/loss_blurry_total": test_loss_blurry_total / (test_i + 1),
+            "train/blurry_pixcorr": blurry_pixcorr / (train_i + 1),
+            "test/blurry_pixcorr": test_blurry_pixcorr / (test_i + 1),
+            }
+        progress_bar.set_postfix(**logs)
+
+        # Save model checkpoint and reconstruct
+        if epoch % ckpt_interval == 0:
+            if not utils.is_interactive():
+                save_ckpt(f'last')
+                
+        if wandb_log: wandb.log(logs)
+
+    # wait for other GPUs to catch up if needed
+    accelerator.wait_for_everyone()
+    torch.cuda.empty_cache()
+    gc.collect()
+
+print("\n===Finished!===\n")
+if ckpt_saving:
+    save_ckpt(f'last')
+if not utils.is_interactive():
+    sys.exit(0)
+
+
+
+# In[ ]:
+
+
+plt.plot(losses)
+plt.show()
+plt.plot(test_losses)
+plt.show()
+

src/Train-with-memory-rr.py

@@ -0,0 +1,1018 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[ ]:
+
+
+# # Code to convert this notebook to .py if you want to run it via command line or with Slurm
+#from subprocess import call
+#command = "jupyter nbconvert Train-with-memory-rr.ipynb --to python"
+#call(command,shell=True)
+
+
+# In[24]:
+
+
+#get_ipython().system('nvidia-smi')
+
+
+# # Import packages & functions
+
+# In[3]:
+
+
+import os
+import sys
+import json
+import argparse
+import numpy as np
+import math
+from einops import rearrange
+import time
+import random
+import h5py
+from tqdm import tqdm
+
+import webdataset as wds
+import gc
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+
+from accelerate import Accelerator, DeepSpeedPlugin
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+
+global_batch_size = 512 #128
+
+
+# In[4]:
+
+
+### Multi-GPU config ###
+local_rank = os.getenv('RANK')
+if local_rank is None: 
+    local_rank = 0
+else:
+    local_rank = int(local_rank)
+print("LOCAL RANK ", local_rank)  
+
+num_devices = torch.cuda.device_count()
+if num_devices==0: num_devices = 1
+
+accelerator = Accelerator(split_batches=False)
+
+### UNCOMMENT BELOW STUFF TO USE DEEPSPEED (also comment out the immediately above "accelerator = " line) ###
+
+# if num_devices <= 1 and utils.is_interactive():
+#     # can emulate a distributed environment for deepspeed to work in jupyter notebook
+#     os.environ["MASTER_ADDR"] = "localhost"
+#     os.environ["MASTER_PORT"] = str(np.random.randint(10000)+9000)
+#     os.environ["RANK"] = "0"
+#     os.environ["LOCAL_RANK"] = "0"
+#     os.environ["WORLD_SIZE"] = "1"
+#     os.environ["GLOBAL_BATCH_SIZE"] = str(global_batch_size) # set this to your batch size!
+#     global_batch_size = os.environ["GLOBAL_BATCH_SIZE"]
+
+# # alter the deepspeed config according to your global and local batch size
+# if local_rank == 0:
+#     with open('deepspeed_config_stage2.json', 'r') as file:
+#         config = json.load(file)
+#     config['train_batch_size'] = int(os.environ["GLOBAL_BATCH_SIZE"])
+#     config['train_micro_batch_size_per_gpu'] = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices
+#     with open('deepspeed_config_stage2.json', 'w') as file:
+#         json.dump(config, file)
+# else:
+#     # give some time for the local_rank=0 gpu to prep new deepspeed config file
+#     time.sleep(10)
+# deepspeed_plugin = DeepSpeedPlugin("deepspeed_config_stage2.json")
+# accelerator = Accelerator(split_batches=False, deepspeed_plugin=deepspeed_plugin)
+
+
+# In[5]:
+
+
+print("PID of this process =",os.getpid())
+device = accelerator.device
+print("device:",device)
+num_workers = num_devices
+print(accelerator.state)
+world_size = accelerator.state.num_processes
+distributed = not accelerator.state.distributed_type == 'NO'
+print("distributed =",distributed, "num_devices =", num_devices, "local rank =", local_rank, "world size =", world_size)
+print = accelerator.print # only print if local_rank=0
+
+
+# # Configurations
+
+# In[6]:
+
+
+# if running this interactively, can specify jupyter_args here for argparser to use
+if utils.is_interactive():
+    # Example use
+    jupyter_args = f"--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset \
+                    --model_name=test \
+                    --subj=1 --batch_size={global_batch_size} --n_samples_save=0 \
+                    --max_lr=3e-5 --mixup_pct=.66 --num_epochs=12 --ckpt_interval=999 --no-use_image_aug"
+
+    jupyter_args = jupyter_args.split()
+    print(jupyter_args)
+    
+    from IPython.display import clear_output # function to clear print outputs in cell
+    get_ipython().run_line_magic('load_ext', 'autoreload')
+    # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions
+    get_ipython().run_line_magic('autoreload', '2')
+
+
+# In[7]:
+
+
+parser = argparse.ArgumentParser(description="Model Training Configuration")
+parser.add_argument(
+    "--model_name", type=str, default="memory_cat_rr",
+    help="name of model, used for ckpt saving and wandb logging (if enabled)",
+)
+parser.add_argument(
+    "--data_path", type=str, default="/fsx/proj-fmri/shared/natural-scenes-dataset",
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,5,7],
+)
+parser.add_argument(
+    "--batch_size", type=int, default=32,
+    help="Batch size can be increased by 10x if only training v2c and not diffusion diffuser",
+)
+parser.add_argument(
+    "--wandb_log",action=argparse.BooleanOptionalAction,default=False,
+    help="whether to log to wandb",
+)
+parser.add_argument(
+    "--resume_from_ckpt",action=argparse.BooleanOptionalAction,default=False,
+    help="if not using wandb and want to resume from a ckpt",
+)
+parser.add_argument(
+    "--wandb_project",type=str,default="stability",
+    help="wandb project name",
+)
+parser.add_argument(
+    "--mixup_pct",type=float,default=.33,
+    help="proportion of way through training when to switch from BiMixCo to SoftCLIP",
+)
+parser.add_argument(
+    "--use_image_aug",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to use image augmentation",
+)
+parser.add_argument(
+    "--num_epochs",type=int,default=240,
+    help="number of epochs of training",
+)
+parser.add_argument(
+    "--lr_scheduler_type",type=str,default='cycle',choices=['cycle','linear'],
+)
+parser.add_argument(
+    "--ckpt_saving",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--ckpt_interval",type=int,default=5,
+    help="save backup ckpt and reconstruct every x epochs",
+)
+parser.add_argument(
+    "--seed",type=int,default=42,
+)
+parser.add_argument(
+    "--max_lr",type=float,default=3e-4,
+)
+parser.add_argument(
+    "--n_samples_save",type=int,default=0,choices=[0,1],
+    help="Number of reconstructions for monitoring progress, 0 will speed up training",
+)
+
+if utils.is_interactive():
+    args = parser.parse_args(jupyter_args)
+else:
+    args = parser.parse_args()
+
+# create global variables without the args prefix
+for attribute_name in vars(args).keys():
+    globals()[attribute_name] = getattr(args, attribute_name)
+
+print("global batch_size", batch_size)
+batch_size = int(batch_size / num_devices)
+print("batch_size", batch_size)
+
+
+# In[8]:
+
+
+outdir = os.path.abspath(f'../train_mem_logs/{model_name}')
+if not os.path.exists(outdir):
+    os.makedirs(outdir,exist_ok=True)
+if use_image_aug:
+    import kornia
+    from kornia.augmentation.container import AugmentationSequential
+    img_augment = AugmentationSequential(
+        kornia.augmentation.RandomResizedCrop((224,224), (0.6,1), p=0.3),
+        kornia.augmentation.Resize((224, 224)),
+        kornia.augmentation.RandomHorizontalFlip(p=0.3),
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.3),
+        kornia.augmentation.RandomGrayscale(p=0.3),
+        same_on_batch=False,
+        data_keys=["input"],
+    )
+
+
+# # Prep data, models, and dataloaders
+
+# ## Dataloader
+
+# In[9]:
+
+
+if subj==1:
+    num_train = 24958
+    num_test = 2770
+test_batch_size = num_test
+
+def my_split_by_node(urls): return urls
+    
+train_url = f"{data_path}/wds/subj0{subj}/train/" + "{0..36}.tar"
+print(train_url)
+
+train_data = wds.WebDataset(train_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+print(test_url)
+
+test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+test_dl = torch.utils.data.DataLoader(test_data, batch_size=test_batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+
+# ### check dataloaders are working
+
+# In[10]:
+
+
+# test_indices = []
+# test_images = []
+# for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+#     test_indices = np.append(test_indices, behav[:,0,5].numpy())
+#     test_images = np.append(test_images, behav[:,0,0].numpy())
+# test_indices = test_indices.astype(np.int16)
+# print(test_i, (test_i+1) * test_batch_size, len(test_indices))
+# print("---\n")
+
+# train_indices = []
+# train_images = []
+# for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+#     train_indices = np.append(train_indices, behav[:,0,5].long().numpy())
+#     train_images = np.append(train_images, behav[:,0,0].numpy())
+# train_indices = train_indices.astype(np.int16)
+# print(train_i, (train_i+1) * batch_size, len(train_indices))
+
+# # train_images = np.hstack((train_images, test_images))
+# # print("WARNING: ADDED TEST IMAGES TO TRAIN IMAGES")
+
+
+# ## Load data and images
+
+# In[11]:
+
+
+# load betas
+f = h5py.File(f'{data_path}/betas_all_subj0{subj}.hdf5', 'r')
+voxels = f['betas'][:]
+print(f"subj0{subj} betas loaded into memory")
+voxels = torch.Tensor(voxels).to("cpu").half()
+if subj==1:
+    voxels = torch.hstack((voxels, torch.zeros((len(voxels), 5))))
+print("voxels", voxels.shape)
+num_voxels = voxels.shape[-1]
+
+# load orig images
+f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+images = f['images'][:]
+images = torch.Tensor(images).to("cpu").half()
+print("images", images.shape)
+
+
+# ## Load models
+
+# ### CLIP image embeddings  model
+
+# In[12]:
+
+
+from models import Clipper
+clip_model = Clipper("ViT-L/14", device=torch.device(f"cuda:{local_rank}"), hidden_state=True, norm_embs=True)
+
+clip_seq_dim = 257
+clip_emb_dim = 768
+hidden_dim = 4096
+
+
+# ### SD VAE (blurry images)
+
+# In[13]:
+
+
+from diffusers import AutoencoderKL
+autoenc = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, cache_dir="/fsx/proj-fmri/shared/cache")
+# autoenc.load_state_dict(torch.load('../train_logs/sdxl_vae_normed/best.pth')["model_state_dict"])
+autoenc.eval()
+autoenc.requires_grad_(False)
+autoenc.to(device)
+utils.count_params(autoenc)
+
+
+# ### MindEye modules
+
+# In[14]:
+
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+model
+
+
+# In[15]:
+
+
+time_embedding_dim = 512
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer
+    def __init__(self, input_size, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linear = torch.nn.Linear(input_size, out_features)
+    def forward(self, x):
+        return self.linear(x)
+        
+model.ridge = RidgeRegression(voxels.shape[1] + time_embedding_dim, out_features=hidden_dim)
+utils.count_params(model.ridge)
+utils.count_params(model)
+
+b = torch.randn((2,1,voxels.shape[1]))
+time_emb_test = torch.randn((2,1,time_embedding_dim))
+print(b.shape, model.ridge(torch.cat((b,time_emb_test),dim=-1)).shape)
+
+
+# In[16]:
+
+
+from functools import partial
+from diffusers.models.vae import Decoder
+class BrainNetwork(nn.Module):
+    def __init__(self, out_dim=768, in_dim=15724, clip_size=768, h=4096, n_blocks=4, norm_type='ln', act_first=False, drop=.25, blurry_dim=16):
+        super().__init__()
+        self.blurry_dim = blurry_dim
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        self.lin0 = nn.Linear(in_dim, h)
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, out_dim, bias=True)
+        self.blin1 = nn.Linear(out_dim, blurry_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.clip_size = clip_size
+        self.clip_proj = nn.Sequential(
+            nn.LayerNorm(clip_size),
+            nn.GELU(),
+            nn.Linear(clip_size, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, clip_size)
+        )
+        self.upsampler = Decoder(
+                in_channels=64,
+                out_channels=4,
+                up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D"],
+                block_out_channels=[64, 128, 256],
+                layers_per_block=1,
+            )
+        
+    def forward(self, x):
+        x = self.lin0(x)
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        b = self.blin1(x)
+        b = self.upsampler(b.reshape(len(b), -1, 7, 7))
+        c = self.clip_proj(x.reshape(len(x), -1, self.clip_size))
+        return c, b
+
+model.backbone = BrainNetwork(h=2048, in_dim=hidden_dim*2, clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, blurry_dim=64*7*7) 
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+b = torch.randn((2,8192))
+print(b.shape)
+clip_, blur_ = model.backbone(b)
+print(clip_.shape, blur_.shape)
+
+
+# In[17]:
+
+
+# memory model
+
+from timm.layers.mlp import Mlp
+
+class MemoryEncoder(nn.Module):
+    def __init__(self, in_dim=15279, out_dim=768, h=4096, num_past_voxels=15, embedding_time_dim = 512, n_blocks=4, norm_type='ln', act_first=False, drop=.25):
+        super().__init__()
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        self.out_dim = out_dim
+        self.embedding_time = nn.Embedding(num_past_voxels, embedding_time_dim)
+        self.final_input_dim = in_dim + embedding_time_dim
+        self.lin0 = nn.Linear(self.final_input_dim, h)
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, out_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.num_past_voxels = num_past_voxels
+        self.embedding_time_dim = embedding_time_dim
+        self.memory = nn.Parameter(torch.randn((self.num_past_voxels, self.embedding_time_dim)))
+
+
+    def forward(self, x, time):
+        time = time.long()
+        time = self.embedding_time(time)
+        x = torch.cat((x, time), dim=-1)
+        x = self.lin0(x)
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        return x
+    
+# # test the memory encoder
+# memory_encoder = MemoryEncoder(in_dim=voxels.shape[1], out_dim=hidden_dim, num_past_voxels=15, embedding_time_dim=512)
+
+# device = torch.device("cpu")
+# memory_encoder.to(device)
+
+# # count params
+# total_parameters = 0
+# for parameter in memory_encoder.parameters():
+#     total_parameters += parameter.numel()
+
+# rand_input = torch.randn((2, 15279)).to(device)
+# rand_time = torch.randint(0, 15, (2,)).to(device)
+# print(rand_input.shape, rand_time.shape)
+# memory_encoder(rand_input, rand_time).shape
+
+class MemoryCompressor(nn.Module):
+    def __init__(self, in_dim=768, num_past = 15, output_dim=768, h=4096, n_blocks=4, norm_type='ln', act_first=False, drop=.25):
+        super().__init__()
+        self.num_past = num_past
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        self.final_input_dim = in_dim * num_past
+        self.lin0 = nn.Linear(self.final_input_dim, h)
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, output_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.num_past = num_past
+        self.output_dim = output_dim
+
+    def forward(self, x):
+        # x is (batch_size, num_past, in_dim)
+        x = x.reshape(len(x), -1)
+        x = self.lin0(x)
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        return x
+    
+# # test the memory compressor
+# memory_compressor = MemoryCompressor(in_dim=768, num_past=15, output_dim=768)
+
+# device = torch.device("cpu")
+# memory_compressor.to(device)
+
+# # count params
+# total_parameters = 0
+# for parameter in memory_compressor.parameters():
+#     total_parameters += parameter.numel()
+
+# rand_input = torch.randn((2, 15, 768)).to(device)
+# print(rand_input.shape)
+# memory_compressor(rand_input).shape
+
+class TimeEmbedding(nn.Module):
+    def __init__(self, embedding_time_dim=512, num_past_voxels=15):
+        super().__init__()
+        self.embedding_time = nn.Embedding(num_past_voxels, embedding_time_dim)
+        self.num_past_voxels = num_past_voxels
+        self.embedding_time_dim = embedding_time_dim
+
+    def forward(self, time):
+        # time is (batch_size,)
+        time = time.long()
+        time = self.embedding_time(time)
+        return time # (batch_size, embedding_time_dim)
+    
+
+#model.memory_encoder = MemoryEncoder(in_dim=voxels.shape[1], out_dim=4096, num_past_voxels=15, embedding_time_dim=512)
+model.time_embedding = TimeEmbedding(embedding_time_dim=512, num_past_voxels=15)
+model.memory_compressor = MemoryCompressor(in_dim=model.ridge.out_features, num_past=15, output_dim=4096)
+
+#utils.count_params(model.memory_encoder)
+utils.count_params(model.memory_compressor)
+utils.count_params(model)
+
+
+
+# In[18]:
+
+
+no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+opt_grouped_parameters = [
+    {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
+    #{'params': [p for n, p in model.memory_encoder.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.memory_compressor.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.time_embedding.named_parameters()], 'weight_decay': 0.0},
+]
+
+optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr, betas=(0.9, 0.95))
+
+if lr_scheduler_type == 'linear':
+    lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+        optimizer,
+        total_iters=int(num_epochs*(num_train*num_devices//batch_size)),
+        last_epoch=-1
+    )
+elif lr_scheduler_type == 'cycle':
+    total_steps=int(num_epochs*(num_train*num_devices//batch_size))
+    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer, 
+        max_lr=max_lr,
+        total_steps=total_steps,
+        final_div_factor=1000,
+        last_epoch=-1, pct_start=2/num_epochs
+    )
+    
+def save_ckpt(tag):    
+    ckpt_path = outdir+f'/{tag}.pth'
+    print(f'saving {ckpt_path}',flush=True)
+    unwrapped_model = accelerator.unwrap_model(model)
+    try:
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': unwrapped_model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'lr_scheduler': lr_scheduler.state_dict(),
+            'train_losses': losses,
+            'test_losses': test_losses,
+            'lrs': lrs,
+            }, ckpt_path)
+    except:
+        print("Couldn't save... moving on to prevent crashing.")
+    del unwrapped_model
+        
+print("\nDone with model preparations!")
+utils.count_params(model)
+
+
+# In[ ]:
+
+
+
+
+
+# # Weights and Biases
+
+# In[19]:
+
+
+# params for wandb
+wandb_log = True
+if local_rank==0 and wandb_log: # only use main process for wandb logging
+    import wandb
+    
+    wandb_project = 'stability'
+    wandb_run = model_name
+    wandb_notes = ''
+    
+    print(f"wandb {wandb_project} run {wandb_run}")
+    wandb.login(host='https://stability.wandb.io')#, relogin=True)
+    wandb_config = {
+      "model_name": model_name,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "use_image_aug": use_image_aug,
+      "max_lr": max_lr,
+      "lr_scheduler_type": lr_scheduler_type,
+      "mixup_pct": mixup_pct,
+      "num_train": num_train,
+      "num_test": num_test,
+      "seed": seed,
+      "distributed": distributed,
+      "num_devices": num_devices,
+      "world_size": world_size,
+    }
+    print("wandb_config:\n",wandb_config)
+    if False: # wandb_auto_resume
+        print("wandb_id:",model_name)
+        wandb.init(
+            id = model_name,
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+            resume="allow",
+        )
+    else:
+        wandb.init(
+            project=wandb_project,
+            name=model_name,
+            config=wandb_config,
+            notes=wandb_notes,
+        )
+else:
+    wandb_log = False
+
+
+# # More custom functions
+
+# In[20]:
+
+
+# using the same preprocessing as was used in MindEye + BrainDiffuser
+pixcorr_preprocess = transforms.Compose([
+    transforms.Resize(425, interpolation=transforms.InterpolationMode.BILINEAR),
+])
+def pixcorr(images,brains):
+    # Flatten images while keeping the batch dimension
+    all_images_flattened = pixcorr_preprocess(images).reshape(len(images), -1)
+    all_brain_recons_flattened = pixcorr_preprocess(brains).view(len(brains), -1)
+    corrmean = torch.diag(utils.batchwise_pearson_correlation(all_images_flattened, all_brain_recons_flattened)).mean()
+    return corrmean
+
+
+# # Main
+
+# In[21]:
+
+
+epoch = 0
+losses, test_losses, lrs = [], [], []
+best_test_loss = 1e9
+soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))
+
+# Optionally resume from checkpoint #
+if resume_from_ckpt:
+    print("\n---resuming from last.pth ckpt---\n")
+    try:
+        checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+    except:
+        print('last.pth failed... trying last_backup.pth')
+        checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+    epoch = checkpoint['epoch']
+    print("Epoch",epoch)
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])
+    del checkpoint
+elif wandb_log:
+    if wandb.run.resumed:
+        print("\n---resuming from last.pth ckpt---\n")
+        try:
+            checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+        except:
+            print('last.pth failed... trying last_backup.pth')
+            checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+        epoch = checkpoint['epoch']
+        print("Epoch",epoch)
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+        diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])
+        del checkpoint
+torch.cuda.empty_cache()
+
+
+# In[22]:
+
+
+model, optimizer, train_dl, test_dl, lr_scheduler = accelerator.prepare(
+model, optimizer, train_dl, test_dl, lr_scheduler
+)
+
+
+# In[23]:
+
+
+print(f"{model_name} starting with epoch {epoch} / {num_epochs}")
+progress_bar = tqdm(range(0,num_epochs), ncols=1200, disable=(local_rank!=0))
+test_image, test_voxel = None, None
+mse = nn.MSELoss()
+for epoch in progress_bar:
+    model.train()
+    
+    fwd_percent_correct = 0.
+    bwd_percent_correct = 0.
+    test_fwd_percent_correct = 0.
+    test_bwd_percent_correct = 0.
+
+    loss_clip_total = 0.
+    loss_blurry_total = 0.
+    test_loss_clip_total = 0.
+    test_loss_blurry_total = 0.
+
+    blurry_pixcorr = 0.
+    test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1
+    
+    for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+        #if epoch == 0 or epoch == 1:
+        #    break
+        with torch.cuda.amp.autocast():
+            optimizer.zero_grad()
+
+            voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+            
+            image = images[behav[:,0,0].cpu().long()].to(device).float()
+
+            past_15_voxels = voxels[past_behav[:,:,5].cpu().long()].to(device) # batch_size, 15, 15279
+            past_15_times = torch.Tensor([i for i in range(15)]).to(device) # 15
+
+            blurry_image_enc = autoenc.encode(image).latent_dist.mode()
+            
+            if use_image_aug: image = img_augment(image)
+            
+            clip_target = clip_model.embed_image(image)
+            assert not torch.any(torch.isnan(clip_target))
+  
+            if epoch < int(mixup_pct * num_epochs):
+                voxel, perm, betas, select = utils.mixco(voxel)
+
+            # reshape past voxels to be (batch_size * 15, 15279)
+            past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+            past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+            past_15_times = past_15_times.reshape(-1)
+            
+            #print(past_15_voxels.shape, past_15_times.shape)
+            time_embeddings = model.time_embedding(past_15_times)
+            past_info_full = torch.cat((past_15_voxels, time_embeddings), dim=-1)
+            embeds_past_voxels = model.ridge(past_info_full)
+            #print(embeds_past_voxels.shape)
+            embeds_past_voxels = embeds_past_voxels.reshape(voxel.shape[0], 15, -1)
+            #print(embeds_past_voxels.shape)
+            information_past_voxels = model.memory_compressor(embeds_past_voxels)
+
+            positional_current_voxel = torch.zeros((voxel.shape[0], time_embeddings.shape[-1])).to(voxel.device)
+            #print(torch.cat((voxel, positional_current_voxel), dim=-1).shape, positional_current_voxel.shape, voxel.shape)
+            voxel_ridge = torch.cat([model.ridge(torch.cat((voxel, positional_current_voxel), dim=-1)), information_past_voxels], dim=-1)
+            
+            clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+            
+            clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+            clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+            if epoch < int(mixup_pct * num_epochs):                
+                loss_clip = utils.mixco_nce(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006, 
+                    perm=perm, betas=betas, select=select)
+            else:
+                epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=epoch_temp)
+
+            loss_blurry = mse(blurry_image_enc_, blurry_image_enc) 
+
+            loss_clip_total += loss_clip.item()
+            loss_blurry_total += loss_blurry.item()
+
+            loss = loss_blurry + loss_clip
+            
+            utils.check_loss(loss)
+
+            accelerator.backward(loss)
+            optimizer.step()
+    
+            losses.append(loss.item())
+            lrs.append(optimizer.param_groups[0]['lr'])
+    
+            # forward and backward top 1 accuracy        
+            labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+            fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)
+            bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)
+
+            with torch.no_grad():
+                # only doing pixcorr eval on a subset (8) of the samples per batch because its costly & slow to compute autoenc.decode()
+                random_samps = np.random.choice(np.arange(len(voxel)), size=2, replace=False)
+                blurry_recon_images = autoenc.decode(blurry_image_enc_[random_samps]).sample.clamp(0,1)
+                blurry_pixcorr += pixcorr(image[random_samps], blurry_recon_images)
+
+            if lr_scheduler_type is not None:
+                lr_scheduler.step()
+    
+    model.eval()
+    for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+        print('test')
+        with torch.cuda.amp.autocast():
+            with torch.no_grad():   
+                # all test samples should be loaded per batch such that test_i should never exceed 0
+                if len(behav) != num_test: print("!",len(behav),num_test)
+
+                
+                ## Average same-image repeats ##
+                if test_image is None:
+                    voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+                    
+                    image = behav[:,0,0].cpu().long()
+                    
+                    unique_image, sort_indices = torch.unique(image, return_inverse=True)
+                    for im in unique_image:
+                        locs = torch.where(im == image)[0]
+                        if test_image is None:
+                            test_image = images[im][None]
+                            test_voxel = torch.mean(voxel[locs],axis=0)[None]
+                        else:
+                            test_image = torch.vstack((test_image, images[im][None]))
+                            test_voxel = torch.vstack((test_voxel, torch.mean(voxel[locs],axis=0)[None]))
+    
+                # sample of batch_size
+                random_indices = torch.arange(len(test_voxel))[:batch_size] #torch.randperm(len(test_voxel))[:300]
+                voxel = test_voxel[random_indices].to(device)
+                image = test_image[random_indices].to(device)
+
+                current_past_behav = past_behav[random_indices]
+
+                past_15_voxels = voxels[current_past_behav[:,:,5].cpu().long()].to(device) # batch_size, 15, 15279
+                past_15_times = torch.Tensor([i for i in range(15)]).to(device) # 15
+
+                assert len(image) == batch_size
+    
+                blurry_image_enc = autoenc.encode(image).latent_dist.mode()
+        
+                clip_target = clip_model.embed_image(image.float())
+
+                past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+                past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+                past_15_times = past_15_times.reshape(-1)
+                
+                print(past_15_voxels.shape, past_15_times.shape)
+
+                #print(past_15_voxels.shape, past_15_times.shape)
+                time_embeddings = model.time_embedding(past_15_times)
+                past_info_full = torch.cat((past_15_voxels, time_embeddings), dim=-1)
+                embeds_past_voxels = model.ridge(past_info_full)
+                #print(embeds_past_voxels.shape)
+                embeds_past_voxels = embeds_past_voxels.reshape(voxel.shape[0], 15, -1)
+                #print(embeds_past_voxels.shape)
+                information_past_voxels = model.memory_compressor(embeds_past_voxels)
+                positional_current_voxel = torch.zeros((voxel.shape[0], time_embeddings.shape[-1])).to(voxel.device)
+    
+                voxel_ridge = torch.cat([model.ridge(torch.cat((voxel, positional_current_voxel), dim=-1)), information_past_voxels], dim=-1)
+                
+                clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+                
+                clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+                clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+        
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006)
+
+                loss_blurry = mse(blurry_image_enc_, blurry_image_enc)
+                
+                loss = loss_blurry + loss_clip
+                
+                utils.check_loss(loss)
+        
+                test_losses.append(loss.item())
+        
+                # forward and backward top 1 accuracy        
+                labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+                test_fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)
+                test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)
+
+                # halving the batch size because the decoder is computationally heavy
+                blurry_recon_images = autoenc.decode(blurry_image_enc_[:len(voxel)//2]).sample.clamp(0,1)
+                blurry_recon_images = torch.vstack((blurry_recon_images, autoenc.decode(blurry_image_enc_[len(voxel)//2:]).sample.clamp(0,1)))
+                test_blurry_pixcorr += pixcorr(image, blurry_recon_images)
+
+                # transform blurry recon latents to images and plot it
+                #fig, axes = plt.subplots(1, 4, figsize=(8, 4))
+                #axes[0].imshow(utils.torch_to_Image(image[[0]]))
+                #axes[1].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[0]]).sample.clamp(0,1)))
+                #axes[2].imshow(utils.torch_to_Image(image[[1]]))
+                #axes[3].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[1]]).sample.clamp(0,1)))
+                #axes[0].axis('off'); axes[1].axis('off'); axes[2].axis('off'); axes[3].axis('off')
+                #plt.show()
+                wandb.log({"gt": [wandb.Image(utils.torch_to_Image(image[[0]])), wandb.Image(utils.torch_to_Image(image[[1]])) ]}
+                wandb.log({"preds": [utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[0]]).sample.clamp(0,1)), utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[1]]).sample.clamp(0,1)) ]}
+
+    if local_rank==0:      
+        # if utils.is_interactive(): clear_output(wait=True)
+        assert (test_i+1) == 1
+        logs = {"train/loss": np.mean(losses[-(train_i+1):]),
+            "test/loss": np.mean(test_losses[-(test_i+1):]),
+            "train/lr": lrs[-1],
+            "train/num_steps": len(losses),
+            "test/num_steps": len(test_losses),
+            "train/fwd_pct_correct": fwd_percent_correct / (train_i + 1),
+            "train/bwd_pct_correct": bwd_percent_correct / (train_i + 1),
+            "test/test_fwd_pct_correct": test_fwd_percent_correct / (test_i + 1),
+            "test/test_bwd_pct_correct": test_bwd_percent_correct / (test_i + 1),
+            "train/loss_clip_total": loss_clip_total / (train_i + 1),
+            "train/loss_blurry_total": loss_blurry_total / (train_i + 1),
+            "test/loss_clip_total": test_loss_clip_total / (test_i + 1),
+            "test/loss_blurry_total": test_loss_blurry_total / (test_i + 1),
+            "train/blurry_pixcorr": blurry_pixcorr / (train_i + 1),
+            "test/blurry_pixcorr": test_blurry_pixcorr / (test_i + 1),
+            }
+        progress_bar.set_postfix(**logs)
+
+        # Save model checkpoint and reconstruct
+        if epoch % ckpt_interval == 0:
+            if not utils.is_interactive():
+                save_ckpt(f'last')
+                
+        if wandb_log: wandb.log(logs)
+
+    # wait for other GPUs to catch up if needed
+    accelerator.wait_for_everyone()
+    torch.cuda.empty_cache()
+    gc.collect()
+
+print("\n===Finished!===\n")
+if ckpt_saving:
+    save_ckpt(f'last')
+if not utils.is_interactive():
+    sys.exit(0)
+
+
+
+# In[ ]:
+
+
+plt.plot(losses)
+plt.show()
+plt.plot(test_losses)
+plt.show()
+

src/Train-with-memory.py

@@ -0,0 +1,978 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+# # Code to convert this notebook to .py if you want to run it via command line or with Slurm
+# from subprocess import call
+# command = "jupyter nbconvert Train.ipynb --to python"
+# call(command,shell=True)
+
+
+# # Import packages & functions
+
+# In[2]:
+
+
+import os
+import sys
+import json
+import argparse
+import numpy as np
+import math
+#from einops import rearrange
+import time
+import random
+import h5py
+from tqdm import tqdm
+
+import webdataset as wds
+import gc
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+
+from accelerate import Accelerator, DeepSpeedPlugin
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+
+global_batch_size = 128 #128
+
+
+# In[3]:
+
+
+### Multi-GPU config ###
+local_rank = os.getenv('RANK')
+if local_rank is None: 
+    local_rank = 0
+else:
+    local_rank = int(local_rank)
+print("LOCAL RANK ", local_rank)  
+
+num_devices = torch.cuda.device_count()
+if num_devices==0: num_devices = 1
+
+accelerator = Accelerator(split_batches=False)
+
+### UNCOMMENT BELOW STUFF TO USE DEEPSPEED (also comment out the immediately above "accelerator = " line) ###
+
+# if num_devices <= 1 and utils.is_interactive():
+#     # can emulate a distributed environment for deepspeed to work in jupyter notebook
+#     os.environ["MASTER_ADDR"] = "localhost"
+#     os.environ["MASTER_PORT"] = str(np.random.randint(10000)+9000)
+#     os.environ["RANK"] = "0"
+#     os.environ["LOCAL_RANK"] = "0"
+#     os.environ["WORLD_SIZE"] = "1"
+#     os.environ["GLOBAL_BATCH_SIZE"] = str(global_batch_size) # set this to your batch size!
+#     global_batch_size = os.environ["GLOBAL_BATCH_SIZE"]
+
+# # alter the deepspeed config according to your global and local batch size
+# if local_rank == 0:
+#     with open('deepspeed_config_stage2.json', 'r') as file:
+#         config = json.load(file)
+#     config['train_batch_size'] = int(os.environ["GLOBAL_BATCH_SIZE"])
+#     config['train_micro_batch_size_per_gpu'] = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices
+#     with open('deepspeed_config_stage2.json', 'w') as file:
+#         json.dump(config, file)
+# else:
+#     # give some time for the local_rank=0 gpu to prep new deepspeed config file
+#     time.sleep(10)
+# deepspeed_plugin = DeepSpeedPlugin("deepspeed_config_stage2.json")
+# accelerator = Accelerator(split_batches=False, deepspeed_plugin=deepspeed_plugin)
+
+
+# In[4]:
+
+
+print("PID of this process =",os.getpid())
+device = accelerator.device
+print("device:",device)
+num_workers = num_devices
+print(accelerator.state)
+world_size = accelerator.state.num_processes
+distributed = not accelerator.state.distributed_type == 'NO'
+print("distributed =",distributed, "num_devices =", num_devices, "local rank =", local_rank, "world size =", world_size)
+print = accelerator.print # only print if local_rank=0
+
+
+# # Configurations
+
+# In[5]:
+
+
+# if running this interactively, can specify jupyter_args here for argparser to use
+if utils.is_interactive():
+    # Example use
+    jupyter_args = f"--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset \
+                    --model_name=test \
+                    --subj=1 --batch_size={global_batch_size} --n_samples_save=0 \
+                    --max_lr=3e-5 --mixup_pct=.66 --num_epochs=12 --ckpt_interval=999 --no-use_image_aug"
+
+    jupyter_args = jupyter_args.split()
+    print(jupyter_args)
+    
+    from IPython.display import clear_output # function to clear print outputs in cell
+    get_ipython().run_line_magic('load_ext', 'autoreload')
+    # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions
+    get_ipython().run_line_magic('autoreload', '2')
+
+
+# In[6]:
+
+
+parser = argparse.ArgumentParser(description="Model Training Configuration")
+parser.add_argument(
+    "--model_name", type=str, default="testing",
+    help="name of model, used for ckpt saving and wandb logging (if enabled)",
+)
+parser.add_argument(
+    "--data_path", type=str, default="/fsx/proj-fmri/shared/natural-scenes-dataset",
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,5,7],
+)
+parser.add_argument(
+    "--batch_size", type=int, default=32,
+    help="Batch size can be increased by 10x if only training v2c and not diffusion diffuser",
+)
+parser.add_argument(
+    "--wandb_log",action=argparse.BooleanOptionalAction,default=False,
+    help="whether to log to wandb",
+)
+parser.add_argument(
+    "--resume_from_ckpt",action=argparse.BooleanOptionalAction,default=False,
+    help="if not using wandb and want to resume from a ckpt",
+)
+parser.add_argument(
+    "--wandb_project",type=str,default="stability",
+    help="wandb project name",
+)
+parser.add_argument(
+    "--mixup_pct",type=float,default=.33,
+    help="proportion of way through training when to switch from BiMixCo to SoftCLIP",
+)
+parser.add_argument(
+    "--use_image_aug",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to use image augmentation",
+)
+parser.add_argument(
+    "--num_epochs",type=int,default=240,
+    help="number of epochs of training",
+)
+parser.add_argument(
+    "--lr_scheduler_type",type=str,default='cycle',choices=['cycle','linear'],
+)
+parser.add_argument(
+    "--ckpt_saving",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--ckpt_interval",type=int,default=5,
+    help="save backup ckpt and reconstruct every x epochs",
+)
+parser.add_argument(
+    "--seed",type=int,default=42,
+)
+parser.add_argument(
+    "--max_lr",type=float,default=3e-4,
+)
+parser.add_argument(
+    "--n_samples_save",type=int,default=0,choices=[0,1],
+    help="Number of reconstructions for monitoring progress, 0 will speed up training",
+)
+
+if utils.is_interactive():
+    args = parser.parse_args(jupyter_args)
+else:
+    args = parser.parse_args()
+
+# create global variables without the args prefix
+for attribute_name in vars(args).keys():
+    globals()[attribute_name] = getattr(args, attribute_name)
+
+print("global batch_size", batch_size)
+batch_size = int(batch_size / num_devices)
+print("batch_size", batch_size)
+
+
+# In[7]:
+
+
+outdir = os.path.abspath(f'../train_mem_logs/{model_name}')
+if not os.path.exists(outdir):
+    os.makedirs(outdir,exist_ok=True)
+if use_image_aug:
+    import kornia
+    from kornia.augmentation.container import AugmentationSequential
+    img_augment = AugmentationSequential(
+        kornia.augmentation.RandomResizedCrop((224,224), (0.6,1), p=0.3),
+        kornia.augmentation.Resize((224, 224)),
+        kornia.augmentation.RandomHorizontalFlip(p=0.3),
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.3),
+        kornia.augmentation.RandomGrayscale(p=0.3),
+        same_on_batch=False,
+        data_keys=["input"],
+    )
+
+
+# # Prep data, models, and dataloaders
+
+# ## Dataloader
+
+# In[8]:
+
+
+if subj==1:
+    num_train = 24958
+    num_test = 2770
+test_batch_size = num_test
+
+def my_split_by_node(urls): return urls
+    
+train_url = f"{data_path}/wds/subj0{subj}/train/" + "{0..36}.tar"
+print(train_url)
+
+train_data = wds.WebDataset(train_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+print(test_url)
+
+test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+test_dl = torch.utils.data.DataLoader(test_data, batch_size=test_batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+
+# ### check dataloaders are working
+
+# In[9]:
+
+
+# test_indices = []
+# test_images = []
+# for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+#     test_indices = np.append(test_indices, behav[:,0,5].numpy())
+#     test_images = np.append(test_images, behav[:,0,0].numpy())
+# test_indices = test_indices.astype(np.int16)
+# print(test_i, (test_i+1) * test_batch_size, len(test_indices))
+# print("---\n")
+
+# train_indices = []
+# train_images = []
+# for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+#     train_indices = np.append(train_indices, behav[:,0,5].long().numpy())
+#     train_images = np.append(train_images, behav[:,0,0].numpy())
+# train_indices = train_indices.astype(np.int16)
+# print(train_i, (train_i+1) * batch_size, len(train_indices))
+
+# # train_images = np.hstack((train_images, test_images))
+# # print("WARNING: ADDED TEST IMAGES TO TRAIN IMAGES")
+
+
+# ## Load data and images
+
+# In[10]:
+
+
+# load betas
+f = h5py.File(f'{data_path}/betas_all_subj0{subj}.hdf5', 'r')
+voxels = f['betas'][:]
+print(f"subj0{subj} betas loaded into memory")
+voxels = torch.Tensor(voxels).to("cpu").half()
+if subj==1:
+    voxels = torch.hstack((voxels, torch.zeros((len(voxels), 5))))
+print("voxels", voxels.shape)
+num_voxels = voxels.shape[-1]
+
+# load orig images
+f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+images = f['images'][:]
+images = torch.Tensor(images).to("cpu").half()
+print("images", images.shape)
+
+
+# ## Load models
+
+# ### CLIP image embeddings  model
+
+# In[11]:
+
+
+from models import Clipper
+clip_model = Clipper("ViT-L/14", device=torch.device(f"cuda:{local_rank}"), hidden_state=True, norm_embs=True)
+
+clip_seq_dim = 257
+clip_emb_dim = 768
+hidden_dim = 4096
+
+
+# ### SD VAE (blurry images)
+
+# In[12]:
+
+
+from diffusers import AutoencoderKL
+autoenc = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, cache_dir="/fsx/proj-fmri/shared/cache")
+# autoenc.load_state_dict(torch.load('../train_logs/sdxl_vae_normed/best.pth')["model_state_dict"])
+autoenc.eval()
+autoenc.requires_grad_(False)
+autoenc.to(device)
+utils.count_params(autoenc)
+
+
+# ### MindEye modules
+
+# In[13]:
+
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+model
+
+
+# In[14]:
+
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer
+    def __init__(self, input_size, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linear = torch.nn.Linear(input_size, out_features)
+    def forward(self, x):
+        return self.linear(x)
+        
+model.ridge = RidgeRegression(voxels.shape[1], out_features=hidden_dim)
+utils.count_params(model.ridge)
+utils.count_params(model)
+
+b = torch.randn((2,1,voxels.shape[1]))
+print(b.shape, model.ridge(b).shape)
+
+
+# In[15]:
+
+
+from functools import partial
+from diffusers.models.vae import Decoder
+class BrainNetwork(nn.Module):
+    def __init__(self, out_dim=768, in_dim=15724, clip_size=768, h=4096, n_blocks=4, norm_type='ln', act_first=False, drop=.15, blurry_dim=16):
+        super().__init__()
+        self.blurry_dim = blurry_dim
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        self.lin0 = nn.Linear(in_dim, h)
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, out_dim, bias=True)
+        self.blin1 = nn.Linear(out_dim, blurry_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.clip_size = clip_size
+        self.clip_proj = nn.Sequential(
+            nn.LayerNorm(clip_size),
+            nn.GELU(),
+            nn.Linear(clip_size, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, clip_size)
+        )
+        self.upsampler = Decoder(
+                in_channels=64,
+                out_channels=4,
+                up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D"],
+                block_out_channels=[64, 128, 256],
+                layers_per_block=1,
+            )
+        
+    def forward(self, x):
+        x = self.lin0(x)
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        b = self.blin1(x)
+        b = self.upsampler(b.reshape(len(b), -1, 7, 7))
+        c = self.clip_proj(x.reshape(len(x), -1, self.clip_size))
+        return c, b
+
+model.backbone = BrainNetwork(h=2048, in_dim=hidden_dim, clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, blurry_dim=64*7*7) 
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+b = torch.randn((2,hidden_dim))
+print(b.shape)
+clip_, blur_ = model.backbone(b)
+print(clip_.shape, blur_.shape)
+
+
+# In[19]:
+
+
+# memory model
+
+from timm.layers.mlp import Mlp
+
+class MemoryEncoder(nn.Module):
+    def __init__(self, in_dim=15279, out_dim=768, h=4096, num_past_voxels=15, embedding_time_dim = 512, n_blocks=4, norm_type='ln', act_first=False, drop=.15):
+        super().__init__()
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        self.out_dim = out_dim
+        self.embedding_time = nn.Embedding(num_past_voxels, embedding_time_dim)
+        self.final_input_dim = in_dim + embedding_time_dim
+        self.lin0 = nn.Linear(self.final_input_dim, h)
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, out_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.num_past_voxels = num_past_voxels
+        self.embedding_time_dim = embedding_time_dim
+        self.memory = nn.Parameter(torch.randn((self.num_past_voxels, self.embedding_time_dim)))
+
+
+    def forward(self, x, time):
+        time = time.long()
+        time = self.embedding_time(time)
+        x = torch.cat((x, time), dim=-1)
+        x = self.lin0(x)
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        return x
+    
+# # test the memory encoder
+# memory_encoder = MemoryEncoder(in_dim=voxels.shape[1], out_dim=hidden_dim, num_past_voxels=15, embedding_time_dim=512)
+
+# device = torch.device("cpu")
+# memory_encoder.to(device)
+
+# # count params
+# total_parameters = 0
+# for parameter in memory_encoder.parameters():
+#     total_parameters += parameter.numel()
+
+# rand_input = torch.randn((2, 15279)).to(device)
+# rand_time = torch.randint(0, 15, (2,)).to(device)
+# print(rand_input.shape, rand_time.shape)
+# memory_encoder(rand_input, rand_time).shape
+
+class MemoryCompressor(nn.Module):
+    def __init__(self, in_dim=768, num_past = 15, output_dim=768, h=4096, n_blocks=4, norm_type='ln', act_first=False, drop=.15):
+        super().__init__()
+        self.num_past = num_past
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        self.final_input_dim = in_dim * num_past
+        self.lin0 = nn.Linear(self.final_input_dim, h)
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, output_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.num_past = num_past
+        self.output_dim = output_dim
+
+    def forward(self, x):
+        # x is (batch_size, num_past, in_dim)
+        x = x.reshape(len(x), -1)
+        x = self.lin0(x)
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        return x
+    
+# # test the memory compressor
+# memory_compressor = MemoryCompressor(in_dim=768, num_past=15, output_dim=768)
+
+# device = torch.device("cpu")
+# memory_compressor.to(device)
+
+# # count params
+# total_parameters = 0
+# for parameter in memory_compressor.parameters():
+#     total_parameters += parameter.numel()
+
+# rand_input = torch.randn((2, 15, 768)).to(device)
+# print(rand_input.shape)
+# memory_compressor(rand_input).shape
+
+model.memory_encoder = MemoryEncoder(in_dim=voxels.shape[1], out_dim=4096, num_past_voxels=15, embedding_time_dim=512)
+model.memory_compressor = MemoryCompressor(in_dim=model.memory_encoder.out_dim, num_past=15, output_dim=4096)
+
+utils.count_params(model.memory_encoder)
+utils.count_params(model.memory_compressor)
+utils.count_params(model)
+
+
+
+# In[17]:
+
+
+no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+opt_grouped_parameters = [
+    {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
+    {'params': [p for n, p in model.memory_encoder.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.memory_compressor.named_parameters()], 'weight_decay': 1e-2},
+]
+
+optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr, betas=(0.9, 0.95))
+
+if lr_scheduler_type == 'linear':
+    lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+        optimizer,
+        total_iters=int(num_epochs*(num_train*num_devices//batch_size)),
+        last_epoch=-1
+    )
+elif lr_scheduler_type == 'cycle':
+    total_steps=int(num_epochs*(num_train*num_devices//batch_size))
+    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer, 
+        max_lr=max_lr,
+        total_steps=total_steps,
+        final_div_factor=1000,
+        last_epoch=-1, pct_start=2/num_epochs
+    )
+    
+def save_ckpt(tag):    
+    ckpt_path = outdir+f'/{tag}.pth'
+    print(f'saving {ckpt_path}',flush=True)
+    unwrapped_model = accelerator.unwrap_model(model)
+    try:
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': unwrapped_model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'lr_scheduler': lr_scheduler.state_dict(),
+            'train_losses': losses,
+            'test_losses': test_losses,
+            'lrs': lrs,
+            }, ckpt_path)
+    except:
+        print("Couldn't save... moving on to prevent crashing.")
+    del unwrapped_model
+        
+print("\nDone with model preparations!")
+utils.count_params(model)
+
+
+
+# # Weights and Biases
+
+# In[ ]:
+
+
+# params for wandb
+wandb_log = True
+if local_rank==0 and wandb_log: # only use main process for wandb logging
+    import wandb
+    
+    wandb_project = 'stability'
+    wandb_run = model_name
+    wandb_notes = ''
+    
+    print(f"wandb {wandb_project} run {wandb_run}")
+    wandb.login(host='https://stability.wandb.io')#, relogin=True)
+    wandb_config = {
+      "model_name": model_name,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "use_image_aug": use_image_aug,
+      "max_lr": max_lr,
+      "lr_scheduler_type": lr_scheduler_type,
+      "mixup_pct": mixup_pct,
+      "num_train": num_train,
+      "num_test": num_test,
+      "seed": seed,
+      "distributed": distributed,
+      "num_devices": num_devices,
+      "world_size": world_size,
+    }
+    print("wandb_config:\n",wandb_config)
+    if False: # wandb_auto_resume
+        print("wandb_id:",model_name)
+        wandb.init(
+            id = model_name,
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+            resume="allow",
+        )
+    else:
+        wandb.init(
+            project=wandb_project,
+            name=model_name,
+            config=wandb_config,
+            notes=wandb_notes,
+        )
+else:
+    wandb_log = False
+
+
+# # More custom functions
+
+# In[ ]:
+
+
+# using the same preprocessing as was used in MindEye + BrainDiffuser
+pixcorr_preprocess = transforms.Compose([
+    transforms.Resize(425, interpolation=transforms.InterpolationMode.BILINEAR),
+])
+def pixcorr(images,brains):
+    # Flatten images while keeping the batch dimension
+    all_images_flattened = pixcorr_preprocess(images).reshape(len(images), -1)
+    all_brain_recons_flattened = pixcorr_preprocess(brains).view(len(brains), -1)
+    corrmean = torch.diag(utils.batchwise_pearson_correlation(all_images_flattened, all_brain_recons_flattened)).mean()
+    return corrmean
+
+
+# # Main
+
+# In[ ]:
+
+
+epoch = 0
+losses, test_losses, lrs = [], [], []
+best_test_loss = 1e9
+soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))
+
+# Optionally resume from checkpoint #
+if resume_from_ckpt:
+    print("\n---resuming from last.pth ckpt---\n")
+    try:
+        checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+    except:
+        print('last.pth failed... trying last_backup.pth')
+        checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+    epoch = checkpoint['epoch']
+    print("Epoch",epoch)
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])
+    del checkpoint
+elif wandb_log:
+    if wandb.run.resumed:
+        print("\n---resuming from last.pth ckpt---\n")
+        try:
+            checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+        except:
+            print('last.pth failed... trying last_backup.pth')
+            checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+        epoch = checkpoint['epoch']
+        print("Epoch",epoch)
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+        diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])
+        del checkpoint
+torch.cuda.empty_cache()
+
+
+# In[ ]:
+
+
+model, optimizer, train_dl, test_dl, lr_scheduler = accelerator.prepare(
+model, optimizer, train_dl, test_dl, lr_scheduler
+)
+
+
+# In[ ]:
+
+
+print(f"{model_name} starting with epoch {epoch} / {num_epochs}")
+progress_bar = tqdm(range(0,num_epochs), ncols=1200, disable=(local_rank!=0))
+test_image, test_voxel = None, None
+mse = nn.MSELoss()
+for epoch in progress_bar:
+    model.train()
+    
+    fwd_percent_correct = 0.
+    bwd_percent_correct = 0.
+    test_fwd_percent_correct = 0.
+    test_bwd_percent_correct = 0.
+
+    loss_clip_total = 0.
+    loss_blurry_total = 0.
+    test_loss_clip_total = 0.
+    test_loss_blurry_total = 0.
+
+    blurry_pixcorr = 0.
+    test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1
+    
+    for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+        #if epoch == 0 or epoch == 1:
+        #    break
+        with torch.cuda.amp.autocast():
+            optimizer.zero_grad()
+
+            voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+            
+            image = images[behav[:,0,0].cpu().long()].to(device).float()
+
+            past_15_voxels = voxels[past_behav[:,:,5].cpu().long()].to(device) # batch_size, 15, 15279
+            past_15_times = torch.Tensor([i for i in range(15)]).to(device) # 15
+
+            blurry_image_enc = autoenc.encode(image).latent_dist.mode()
+            
+            if use_image_aug: image = img_augment(image)
+            
+            clip_target = clip_model.embed_image(image)
+            assert not torch.any(torch.isnan(clip_target))
+  
+            if epoch < int(mixup_pct * num_epochs):
+                voxel, perm, betas, select = utils.mixco(voxel)
+
+            # reshape past voxels to be (batch_size * 15, 15279)
+            past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+            past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+            past_15_times = past_15_times.reshape(-1)
+            
+            #print(past_15_voxels.shape, past_15_times.shape)
+
+            embeds_past_voxels = model.memory_encoder(past_15_voxels, past_15_times)
+            #print(embeds_past_voxels.shape)
+            embeds_past_voxels = embeds_past_voxels.reshape(voxel.shape[0], 15, -1)
+            #print(embeds_past_voxels.shape)
+            information_past_voxels = model.memory_compressor(embeds_past_voxels)
+
+
+            voxel_ridge = model.ridge(voxel) + information_past_voxels
+            
+            clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+            
+            clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+            clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+            if epoch < int(mixup_pct * num_epochs):                
+                loss_clip = utils.mixco_nce(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006, 
+                    perm=perm, betas=betas, select=select)
+            else:
+                epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=epoch_temp)
+
+            loss_blurry = mse(blurry_image_enc_, blurry_image_enc) 
+
+            loss_clip_total += loss_clip.item()
+            loss_blurry_total += loss_blurry.item()
+
+            loss = loss_blurry + loss_clip
+            
+            utils.check_loss(loss)
+
+            accelerator.backward(loss)
+            optimizer.step()
+    
+            losses.append(loss.item())
+            lrs.append(optimizer.param_groups[0]['lr'])
+    
+            # forward and backward top 1 accuracy        
+            labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+            fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)
+            bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)
+
+            with torch.no_grad():
+                # only doing pixcorr eval on a subset (8) of the samples per batch because its costly & slow to compute autoenc.decode()
+                random_samps = np.random.choice(np.arange(len(voxel)), size=2, replace=False)
+                blurry_recon_images = autoenc.decode(blurry_image_enc_[random_samps]).sample.clamp(0,1)
+                blurry_pixcorr += pixcorr(image[random_samps], blurry_recon_images)
+
+            if lr_scheduler_type is not None:
+                lr_scheduler.step()
+    
+    model.eval()
+    for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+        print('test')
+        with torch.cuda.amp.autocast():
+            with torch.no_grad():   
+                # all test samples should be loaded per batch such that test_i should never exceed 0
+                if len(behav) != num_test: print("!",len(behav),num_test)
+
+                
+                ## Average same-image repeats ##
+                if test_image is None:
+                    voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+                    
+                    image = behav[:,0,0].cpu().long()
+                    
+                    unique_image, sort_indices = torch.unique(image, return_inverse=True)
+                    for im in unique_image:
+                        locs = torch.where(im == image)[0]
+                        if test_image is None:
+                            test_image = images[im][None]
+                            test_voxel = torch.mean(voxel[locs],axis=0)[None]
+                        else:
+                            test_image = torch.vstack((test_image, images[im][None]))
+                            test_voxel = torch.vstack((test_voxel, torch.mean(voxel[locs],axis=0)[None]))
+    
+                # sample of batch_size
+                random_indices = torch.arange(len(test_voxel))[:batch_size] #torch.randperm(len(test_voxel))[:300]
+                voxel = test_voxel[random_indices].to(device)
+                image = test_image[random_indices].to(device)
+
+                current_past_behav = past_behav[random_indices]
+
+                past_15_voxels = voxels[current_past_behav[:,:,5].cpu().long()].to(device) # batch_size, 15, 15279
+                past_15_times = torch.Tensor([i for i in range(15)]).to(device) # 15
+
+                assert len(image) == batch_size
+    
+                blurry_image_enc = autoenc.encode(image).latent_dist.mode()
+        
+                clip_target = clip_model.embed_image(image.float())
+
+                past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+                past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+                past_15_times = past_15_times.reshape(-1)
+                
+                print(past_15_voxels.shape, past_15_times.shape)
+
+                embeds_past_voxels = model.memory_encoder(past_15_voxels, past_15_times)
+                embeds_past_voxels = embeds_past_voxels.reshape(batch_size, 15, -1)
+                information_past_voxels = model.memory_compressor(embeds_past_voxels)
+
+    
+                voxel_ridge = model.ridge(voxel) + information_past_voxels
+                
+                clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+                
+                clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+                clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+        
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006)
+
+                loss_blurry = mse(blurry_image_enc_, blurry_image_enc)
+                
+                loss = loss_blurry + loss_clip
+                
+                utils.check_loss(loss)
+        
+                test_losses.append(loss.item())
+        
+                # forward and backward top 1 accuracy        
+                labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+                test_fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)
+                test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)
+
+                # halving the batch size because the decoder is computationally heavy
+                blurry_recon_images = autoenc.decode(blurry_image_enc_[:len(voxel)//2]).sample.clamp(0,1)
+                blurry_recon_images = torch.vstack((blurry_recon_images, autoenc.decode(blurry_image_enc_[len(voxel)//2:]).sample.clamp(0,1)))
+                test_blurry_pixcorr += pixcorr(image, blurry_recon_images)
+
+                # transform blurry recon latents to images and plot it
+                fig, axes = plt.subplots(1, 4, figsize=(8, 4))
+                axes[0].imshow(utils.torch_to_Image(image[[0]]))
+                axes[1].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[0]]).sample.clamp(0,1)))
+                axes[2].imshow(utils.torch_to_Image(image[[1]]))
+                axes[3].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[1]]).sample.clamp(0,1)))
+                axes[0].axis('off'); axes[1].axis('off'); axes[2].axis('off'); axes[3].axis('off')
+                plt.show()
+
+    if local_rank==0:      
+        # if utils.is_interactive(): clear_output(wait=True)
+        assert (test_i+1) == 1
+        logs = {"train/loss": np.mean(losses[-(train_i+1):]),
+            "test/loss": np.mean(test_losses[-(test_i+1):]),
+            "train/lr": lrs[-1],
+            "train/num_steps": len(losses),
+            "test/num_steps": len(test_losses),
+            "train/fwd_pct_correct": fwd_percent_correct / (train_i + 1),
+            "train/bwd_pct_correct": bwd_percent_correct / (train_i + 1),
+            "test/test_fwd_pct_correct": test_fwd_percent_correct / (test_i + 1),
+            "test/test_bwd_pct_correct": test_bwd_percent_correct / (test_i + 1),
+            "train/loss_clip_total": loss_clip_total / (train_i + 1),
+            "train/loss_blurry_total": loss_blurry_total / (train_i + 1),
+            "test/loss_clip_total": test_loss_clip_total / (test_i + 1),
+            "test/loss_blurry_total": test_loss_blurry_total / (test_i + 1),
+            "train/blurry_pixcorr": blurry_pixcorr / (train_i + 1),
+            "test/blurry_pixcorr": test_blurry_pixcorr / (test_i + 1),
+            }
+        progress_bar.set_postfix(**logs)
+
+        # Save model checkpoint and reconstruct
+        if epoch % ckpt_interval == 0:
+            if not utils.is_interactive():
+                save_ckpt(f'last')
+                
+        if wandb_log: wandb.log(logs)
+
+    # wait for other GPUs to catch up if needed
+    accelerator.wait_for_everyone()
+    torch.cuda.empty_cache()
+    gc.collect()
+
+print("\n===Finished!===\n")
+if ckpt_saving:
+    save_ckpt(f'last')
+if not utils.is_interactive():
+    sys.exit(0)
+
+
+# In[ ]:
+
+
+plt.plot(losses)
+plt.show()
+plt.plot(test_losses)
+plt.show()
+

src/Train.py

@@ -0,0 +1,761 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+# # Code to convert this notebook to .py if you want to run it via command line or with Slurm
+# from subprocess import call
+# command = "jupyter nbconvert Train.ipynb --to python"
+# call(command,shell=True)
+
+
+# # Import packages & functions
+
+# In[2]:
+
+
+import os
+import sys
+import json
+import argparse
+import numpy as np
+import math
+from einops import rearrange
+import time
+import random
+import h5py
+from tqdm import tqdm
+
+import webdataset as wds
+import gc
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+
+from accelerate import Accelerator, DeepSpeedPlugin
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+
+global_batch_size = 128 #128
+
+
+# In[ ]:
+
+
+### Multi-GPU config ###
+local_rank = os.getenv('RANK')
+if local_rank is None: 
+    local_rank = 0
+else:
+    local_rank = int(local_rank)
+print("LOCAL RANK ", local_rank)  
+
+num_devices = torch.cuda.device_count()
+if num_devices==0: num_devices = 1
+
+accelerator = Accelerator(split_batches=False)
+
+### UNCOMMENT BELOW STUFF TO USE DEEPSPEED (also comment out the immediately above "accelerator = " line) ###
+
+# if num_devices <= 1 and utils.is_interactive():
+#     # can emulate a distributed environment for deepspeed to work in jupyter notebook
+#     os.environ["MASTER_ADDR"] = "localhost"
+#     os.environ["MASTER_PORT"] = str(np.random.randint(10000)+9000)
+#     os.environ["RANK"] = "0"
+#     os.environ["LOCAL_RANK"] = "0"
+#     os.environ["WORLD_SIZE"] = "1"
+#     os.environ["GLOBAL_BATCH_SIZE"] = str(global_batch_size) # set this to your batch size!
+#     global_batch_size = os.environ["GLOBAL_BATCH_SIZE"]
+
+# # alter the deepspeed config according to your global and local batch size
+# if local_rank == 0:
+#     with open('deepspeed_config_stage2.json', 'r') as file:
+#         config = json.load(file)
+#     config['train_batch_size'] = int(os.environ["GLOBAL_BATCH_SIZE"])
+#     config['train_micro_batch_size_per_gpu'] = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices
+#     with open('deepspeed_config_stage2.json', 'w') as file:
+#         json.dump(config, file)
+# else:
+#     # give some time for the local_rank=0 gpu to prep new deepspeed config file
+#     time.sleep(10)
+# deepspeed_plugin = DeepSpeedPlugin("deepspeed_config_stage2.json")
+# accelerator = Accelerator(split_batches=False, deepspeed_plugin=deepspeed_plugin)
+                          
+### Multi-GPU config ###
+print("PID of this process =",os.getpid())
+device = accelerator.device
+print("device:",device)
+num_workers = num_devices
+print(accelerator.state)
+world_size = accelerator.state.num_processes
+distributed = not accelerator.state.distributed_type == 'NO'
+print("distributed =",distributed, "num_devices =", num_devices, "local rank =", local_rank, "world size =", world_size)
+print = accelerator.print # only print if local_rank=0
+
+# In[ ]:
+
+
+
+
+
+# # Configurations
+
+# In[3]:
+
+
+
+# In[4]:
+
+parser = argparse.ArgumentParser(description="Model Training Configuration")
+parser.add_argument(
+    "--model_name", type=str, default="testing",
+    help="name of model, used for ckpt saving and wandb logging (if enabled)",
+)
+parser.add_argument(
+    "--data_path", type=str, default="/fsx/proj-fmri/shared/natural-scenes-dataset",
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,5,7],
+)
+parser.add_argument(
+    "--batch_size", type=int, default=32,
+    help="Batch size can be increased by 10x if only training v2c and not diffusion diffuser",
+)
+parser.add_argument(
+    "--wandb_log",action=argparse.BooleanOptionalAction,default=False,
+    help="whether to log to wandb",
+)
+parser.add_argument(
+    "--resume_from_ckpt",action=argparse.BooleanOptionalAction,default=False,
+    help="if not using wandb and want to resume from a ckpt",
+)
+parser.add_argument(
+    "--wandb_project",type=str,default="stability",
+    help="wandb project name",
+)
+parser.add_argument(
+    "--mixup_pct",type=float,default=.33,
+    help="proportion of way through training when to switch from BiMixCo to SoftCLIP",
+)
+parser.add_argument(
+    "--use_image_aug",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to use image augmentation",
+)
+parser.add_argument(
+    "--num_epochs",type=int,default=240,
+    help="number of epochs of training",
+)
+parser.add_argument(
+    "--lr_scheduler_type",type=str,default='cycle',choices=['cycle','linear'],
+)
+parser.add_argument(
+    "--ckpt_saving",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--ckpt_interval",type=int,default=5,
+    help="save backup ckpt and reconstruct every x epochs",
+)
+parser.add_argument(
+    "--seed",type=int,default=42,
+)
+parser.add_argument(
+    "--max_lr",type=float,default=3e-4,
+)
+parser.add_argument(
+    "--n_samples_save",type=int,default=0,choices=[0,1],
+    help="Number of reconstructions for monitoring progress, 0 will speed up training",
+)
+
+if utils.is_interactive():
+    args = parser.parse_args(jupyter_args)
+else:
+    args = parser.parse_args()
+
+# create global variables without the args prefix
+for attribute_name in vars(args).keys():
+    globals()[attribute_name] = getattr(args, attribute_name)
+
+print("global batch_size", batch_size)
+batch_size = int(batch_size / num_devices)
+print("batch_size", batch_size)
+
+# In[5]:
+outdir = os.path.abspath(f'../train_logs/{model_name}')
+if not os.path.exists(outdir):
+    os.makedirs(outdir,exist_ok=True)
+if use_image_aug:
+    import kornia
+    from kornia.augmentation.container import AugmentationSequential
+    img_augment = AugmentationSequential(
+        kornia.augmentation.RandomResizedCrop((224,224), (0.6,1), p=0.3),
+        kornia.augmentation.Resize((224, 224)),
+        kornia.augmentation.RandomHorizontalFlip(p=0.3),
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.3),
+        kornia.augmentation.RandomGrayscale(p=0.3),
+        same_on_batch=False,
+        data_keys=["input"],
+    )
+
+
+# # Prep data, models, and dataloaders
+
+# ## Dataloader
+
+# In[6]:
+
+if subj==1:
+    num_train = 24958
+    num_test = 2770
+test_batch_size = num_test
+
+def my_split_by_node(urls): return urls
+    
+train_url = f"{data_path}/wds/subj0{subj}/train/" + "{0..36}.tar"
+print(train_url)
+
+train_data = wds.WebDataset(train_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+print(test_url)
+
+test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+test_dl = torch.utils.data.DataLoader(test_data, batch_size=test_batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+# ### check dataloaders are working
+
+# In[7]:
+
+
+# test_indices = []
+# test_images = []
+# for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+#     test_indices = np.append(test_indices, behav[:,0,5].numpy())
+#     test_images = np.append(test_images, behav[:,0,0].numpy())
+# test_indices = test_indices.astype(np.int16)
+# print(test_i, (test_i+1) * test_batch_size, len(test_indices))
+# print("---\n")
+
+# train_indices = []
+# train_images = []
+# for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+#     train_indices = np.append(train_indices, behav[:,0,5].long().numpy())
+#     train_images = np.append(train_images, behav[:,0,0].numpy())
+# train_indices = train_indices.astype(np.int16)
+# print(train_i, (train_i+1) * batch_size, len(train_indices))
+
+
+# ## Load voxel betas, K-means clustering model, and images
+
+# In[8]:
+
+
+# load betas
+f = h5py.File(f'{data_path}/betas_all_subj0{subj}.hdf5', 'r')
+voxels = f['betas'][:]
+print(f"subj0{subj} betas loaded into memory")
+voxels = torch.Tensor(voxels).to("cpu").half()
+if subj==1:
+    voxels = torch.hstack((voxels, torch.zeros((len(voxels), 5))))
+print("voxels", voxels.shape)
+num_voxels = voxels.shape[-1]
+
+# load orig images
+f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+images = f['images'][:]
+images = torch.Tensor(images).to("cpu").half()
+print("images", images.shape)
+
+
+# In[9]:
+
+
+from models import Clipper
+clip_model = Clipper("ViT-L/14", device=torch.device(f"cuda:{local_rank}"), hidden_state=True, norm_embs=True)
+
+clip_seq_dim = 257
+clip_emb_dim = 768
+hidden_dim = 4096
+
+from diffusers import AutoencoderKL
+autoenc = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, cache_dir="/fsx/proj-fmri/shared/cache")
+# autoenc.load_state_dict(torch.load('../train_logs/sdxl_vae_normed/best.pth')["model_state_dict"])
+autoenc.eval()
+autoenc.requires_grad_(False)
+autoenc.to(device)
+utils.count_params(autoenc)
+
+
+# In[10]:
+
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+model
+
+
+# In[11]:
+
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer
+    def __init__(self, input_size, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linear = torch.nn.Linear(input_size, out_features)
+    def forward(self, x):
+        return self.linear(x)
+        
+model.ridge = RidgeRegression(voxels.shape[1], out_features=hidden_dim)
+utils.count_params(model.ridge)
+utils.count_params(model)
+
+b = torch.randn((2,1,voxels.shape[1]))
+print(b.shape, model.ridge(b).shape)
+
+# In[12]:
+from functools import partial
+from diffusers.models.vae import Decoder
+class BrainNetwork(nn.Module):
+    def __init__(self, out_dim=768, in_dim=15724, clip_size=768, h=4096, n_blocks=4, norm_type='ln', act_first=False, drop=.15, blurry_dim=16):
+        super().__init__()
+        self.blurry_dim = blurry_dim
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        self.lin0 = nn.Linear(in_dim, h)
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, out_dim, bias=True)
+        self.blin1 = nn.Linear(out_dim, blurry_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.clip_size = clip_size
+        self.clip_proj = nn.Sequential(
+            nn.LayerNorm(clip_size),
+            nn.GELU(),
+            nn.Linear(clip_size, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, clip_size)
+        )
+        self.upsampler = Decoder(
+                in_channels=64,
+                out_channels=4,
+                up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D"],
+                block_out_channels=[64, 128, 256],
+                layers_per_block=1,
+            )
+        
+    def forward(self, x):
+        x = self.lin0(x)
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        b = self.blin1(x)
+        b = self.upsampler(b.reshape(len(b), -1, 7, 7))
+        c = self.clip_proj(x.reshape(len(x), -1, self.clip_size))
+        return c, b
+
+model.backbone = BrainNetwork(h=2048, in_dim=hidden_dim, clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, blurry_dim=64*7*7) 
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+b = torch.randn((2,hidden_dim))
+print(b.shape)
+clip_, blur_ = model.backbone(b)
+print(clip_.shape, blur_.shape)
+
+
+# In[13]:
+
+no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+opt_grouped_parameters = [
+    {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
+]
+
+optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr, betas=(0.9, 0.95))
+
+if lr_scheduler_type == 'linear':
+    lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+        optimizer,
+        total_iters=int(num_epochs*(num_train*num_devices//batch_size)),
+        last_epoch=-1
+    )
+elif lr_scheduler_type == 'cycle':
+    total_steps=int(num_epochs*(num_train*num_devices//batch_size))
+    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer, 
+        max_lr=max_lr,
+        total_steps=total_steps,
+        final_div_factor=1000,
+        last_epoch=-1, pct_start=2/num_epochs
+    )
+    
+def save_ckpt(tag):    
+    ckpt_path = outdir+f'/{tag}.pth'
+    print(f'saving {ckpt_path}',flush=True)
+    unwrapped_model = accelerator.unwrap_model(model)
+    try:
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': unwrapped_model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'lr_scheduler': lr_scheduler.state_dict(),
+            'train_losses': losses,
+            'test_losses': test_losses,
+            'lrs': lrs,
+            }, ckpt_path)
+    except:
+        print("Couldn't save... moving on to prevent crashing.")
+    del unwrapped_model
+        
+print("\nDone with model preparations!")
+utils.count_params(model)
+
+# # Weights and Biases
+
+# In[14]:
+
+
+# params for wandb
+# params for wandb
+wandb_log = True
+if local_rank==0 and wandb_log: # only use main process for wandb logging
+    import wandb
+    
+    wandb_project = 'stability'
+    wandb_run = model_name
+    wandb_notes = ''
+    
+    print(f"wandb {wandb_project} run {wandb_run}")
+    wandb.login(host='https://stability.wandb.io')#, relogin=True)
+    wandb_config = {
+      "model_name": model_name,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "use_image_aug": use_image_aug,
+      "max_lr": max_lr,
+      "lr_scheduler_type": lr_scheduler_type,
+      "mixup_pct": mixup_pct,
+      "num_train": num_train,
+      "num_test": num_test,
+      "seed": seed,
+      "distributed": distributed,
+      "num_devices": num_devices,
+      "world_size": world_size,
+    }
+    print("wandb_config:\n",wandb_config)
+    if False: # wandb_auto_resume
+        print("wandb_id:",model_name)
+        wandb.init(
+            id = model_name,
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+            resume="allow",
+        )
+    else:
+        wandb.init(
+            project=wandb_project,
+            name=model_name,
+            config=wandb_config,
+            notes=wandb_notes,
+        )
+else:
+    wandb_log = False
+
+
+# using the same preprocessing as was used in MindEye + BrainDiffuser
+pixcorr_preprocess = transforms.Compose([
+    transforms.Resize(425, interpolation=transforms.InterpolationMode.BILINEAR),
+])
+def pixcorr(images,brains):
+    # Flatten images while keeping the batch dimension
+    all_images_flattened = pixcorr_preprocess(images).reshape(len(images), -1)
+    all_brain_recons_flattened = pixcorr_preprocess(brains).view(len(brains), -1)
+    corrmean = torch.diag(utils.batchwise_pearson_correlation(all_images_flattened, all_brain_recons_flattened)).mean()
+    return corrmean
+
+# # Main
+
+# In[15]:
+
+epoch = 0
+losses, test_losses, lrs = [], [], []
+best_test_loss = 1e9
+soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))
+
+# Optionally resume from checkpoint #
+resume_from_ckpt = False
+if resume_from_ckpt:
+    print("\n---resuming from last.pth ckpt---\n")
+    try:
+        checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+    except:
+        print('last.pth failed... trying last_backup.pth')
+        checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+    epoch = checkpoint['epoch']
+    print("Epoch",epoch)
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])
+    del checkpoint
+elif False:
+    if wandb.run.resumed:
+        print("\n---resuming from last.pth ckpt---\n")
+        try:
+            checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+        except:
+            print('last.pth failed... trying last_backup.pth')
+            checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+        epoch = checkpoint['epoch']
+        print("Epoch",epoch)
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+        diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])
+        del checkpoint
+torch.cuda.empty_cache()
+
+
+# In[16]:
+
+
+model, optimizer, train_dl, test_dl, lr_scheduler = accelerator.prepare(
+model, optimizer, train_dl, test_dl, lr_scheduler
+)
+
+
+# In[17]:
+print(f"{model_name} starting with epoch {epoch} / {num_epochs}")
+progress_bar = tqdm(range(epoch,num_epochs), ncols=1200, disable=(local_rank!=0))
+test_image, test_voxel = None, None
+mse = nn.MSELoss()
+for epoch in progress_bar:
+    model.train()
+    
+    fwd_percent_correct = 0.
+    bwd_percent_correct = 0.
+    test_fwd_percent_correct = 0.
+    test_bwd_percent_correct = 0.
+
+    loss_clip_total = 0.
+    loss_blurry_total = 0.
+    test_loss_clip_total = 0.
+    test_loss_blurry_total = 0.
+
+    blurry_pixcorr = 0.
+    test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1
+    
+    for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+        with torch.cuda.amp.autocast():
+            optimizer.zero_grad()
+
+            voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+            
+            image = images[behav[:,0,0].cpu().long()].to(device).float()
+
+            blurry_image_enc = autoenc.encode(image).latent_dist.mode()
+            
+            if use_image_aug: image = img_augment(image)
+            
+            clip_target = clip_model.embed_image(image)
+            assert not torch.any(torch.isnan(clip_target))
+  
+            if epoch < int(mixup_pct * num_epochs):
+                voxel, perm, betas, select = utils.mixco(voxel)
+
+            voxel_ridge = model.ridge(voxel)
+            
+            clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+            
+            clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+            clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+            if epoch < int(mixup_pct * num_epochs):                
+                loss_clip = utils.mixco_nce(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006, 
+                    perm=perm, betas=betas, select=select)
+            else:
+                epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=epoch_temp)
+
+            loss_blurry = mse(blurry_image_enc_, blurry_image_enc) 
+
+            loss_clip_total += loss_clip.item()
+            loss_blurry_total += loss_blurry.item()
+
+            loss = loss_blurry + loss_clip
+            
+            utils.check_loss(loss)
+
+            accelerator.backward(loss)
+            optimizer.step()
+    
+            losses.append(loss.item())
+            lrs.append(optimizer.param_groups[0]['lr'])
+    
+            # forward and backward top 1 accuracy        
+            labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+            fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)
+            bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)
+
+            with torch.no_grad():
+                # only doing pixcorr eval on a subset (8) of the samples per batch because its costly & slow to compute autoenc.decode()
+                random_samps = np.random.choice(np.arange(len(voxel)), size=2, replace=False)
+                blurry_recon_images = autoenc.decode(blurry_image_enc_[random_samps]).sample.clamp(0,1)
+                blurry_pixcorr += pixcorr(image[random_samps], blurry_recon_images)
+
+            if lr_scheduler_type is not None:
+                lr_scheduler.step()
+    
+    model.eval()
+    for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+        with torch.cuda.amp.autocast():
+            with torch.no_grad():   
+                # all test samples should be loaded per batch such that test_i should never exceed 0
+                if len(behav) != num_test: print("!",len(behav),num_test)
+                
+                ## Average same-image repeats ##
+                if test_image is None:
+                    voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+                    
+                    image = behav[:,0,0].cpu().long()
+                    
+                    unique_image, sort_indices = torch.unique(image, return_inverse=True)
+                    for im in unique_image:
+                        locs = torch.where(im == image)[0]
+                        if test_image is None:
+                            test_image = images[im][None]
+                            test_voxel = torch.mean(voxel[locs],axis=0)[None]
+                        else:
+                            test_image = torch.vstack((test_image, images[im][None]))
+                            test_voxel = torch.vstack((test_voxel, torch.mean(voxel[locs],axis=0)[None]))
+    
+                # sample of batch_size
+                random_indices = torch.arange(len(test_voxel))[:batch_size] #torch.randperm(len(test_voxel))[:300]
+                voxel = test_voxel[random_indices].to(device)
+                image = test_image[random_indices].to(device)
+                assert len(image) == batch_size
+    
+                blurry_image_enc = autoenc.encode(image).latent_dist.mode()
+        
+                clip_target = clip_model.embed_image(image.float())
+    
+                voxel_ridge = model.ridge(voxel)
+                
+                clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+                
+                clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+                clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+        
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006)
+
+                loss_blurry = mse(blurry_image_enc_, blurry_image_enc)
+                
+                loss = loss_blurry + loss_clip
+                
+                utils.check_loss(loss)
+        
+                test_losses.append(loss.item())
+        
+                # forward and backward top 1 accuracy        
+                labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+                test_fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)
+                test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)
+
+                # halving the batch size because the decoder is computationally heavy
+                blurry_recon_images = autoenc.decode(blurry_image_enc_[:len(voxel)//2]).sample.clamp(0,1)
+                blurry_recon_images = torch.vstack((blurry_recon_images, autoenc.decode(blurry_image_enc_[len(voxel)//2:]).sample.clamp(0,1)))
+                test_blurry_pixcorr += pixcorr(image, blurry_recon_images)
+
+                # transform blurry recon latents to images and plot it
+                fig, axes = plt.subplots(1, 4, figsize=(8, 4))
+                axes[0].imshow(utils.torch_to_Image(image[[0]]))
+                axes[1].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[0]]).sample.clamp(0,1)))
+                axes[2].imshow(utils.torch_to_Image(image[[1]]))
+                axes[3].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[1]]).sample.clamp(0,1)))
+                axes[0].axis('off'); axes[1].axis('off'); axes[2].axis('off'); axes[3].axis('off')
+                plt.show()
+
+    if local_rank==0:      
+        # if utils.is_interactive(): clear_output(wait=True)
+        assert (test_i+1) == 1
+        logs = {"train/loss": np.mean(losses[-(train_i+1):]),
+            "test/loss": np.mean(test_losses[-(test_i+1):]),
+            "train/lr": lrs[-1],
+            "train/num_steps": len(losses),
+            "test/num_steps": len(test_losses),
+            "train/fwd_pct_correct": fwd_percent_correct / (train_i + 1),
+            "train/bwd_pct_correct": bwd_percent_correct / (train_i + 1),
+            "test/test_fwd_pct_correct": test_fwd_percent_correct / (test_i + 1),
+            "test/test_bwd_pct_correct": test_bwd_percent_correct / (test_i + 1),
+            "train/loss_clip_total": loss_clip_total / (train_i + 1),
+            "train/loss_blurry_total": loss_blurry_total / (train_i + 1),
+            "test/loss_clip_total": test_loss_clip_total / (test_i + 1),
+            "test/loss_blurry_total": test_loss_blurry_total / (test_i + 1),
+            "train/blurry_pixcorr": blurry_pixcorr / (train_i + 1),
+            "test/blurry_pixcorr": test_blurry_pixcorr / (test_i + 1),
+            }
+        progress_bar.set_postfix(**logs)
+
+        # Save model checkpoint and reconstruct
+        if epoch % ckpt_interval == 0:
+            if not utils.is_interactive():
+                save_ckpt(f'last')
+                
+        if wandb_log: wandb.log(logs)
+
+    # wait for other GPUs to catch up if needed
+    accelerator.wait_for_everyone()
+    torch.cuda.empty_cache()
+    gc.collect()
+
+print("\n===Finished!===\n")
+if ckpt_saving:
+    save_ckpt(f'last')
+if not utils.is_interactive():
+    sys.exit(0)

src/Train_MLPMixer-Copy1.py

@@ -0,0 +1,1352 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+# # Code to convert this notebook to .py if you want to run it via command line or with Slurm
+#from subprocess import call
+#command = "jupyter nbconvert Train_MLPMixer-Copy1.ipynb --to python"
+#call(command,shell=True)
+
+
+# # Import packages & functions
+
+# In[2]:
+
+
+import os
+import sys
+import json
+import argparse
+import numpy as np
+import math
+from einops import rearrange
+import time
+import random
+import string
+import h5py
+from tqdm import tqdm
+
+import webdataset as wds
+import gc
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+
+from accelerate import Accelerator, DeepSpeedPlugin
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+
+
+# In[3]:
+
+
+### Multi-GPU config ###
+local_rank = os.getenv('RANK')
+if local_rank is None: 
+    local_rank = 0
+else:
+    local_rank = int(local_rank)
+print("LOCAL RANK ", local_rank)  
+
+num_devices = torch.cuda.device_count()
+if num_devices==0: num_devices = 1
+
+# ## UNCOMMENT BELOW SECTION AND COMMENT OUT DEEPSPEED SECTION TO AVOID USING DEEPSPEED ###
+# accelerator = Accelerator(split_batches=False, mixed_precision="fp16")
+# global_batch_size = batch_size = 128
+# data_type = torch.float16 # change depending on your mixed_precision
+
+### DEEPSPEED INITIALIZATION ###
+if num_devices <= 1 and utils.is_interactive():
+    global_batch_size = batch_size = 128
+    print(f"Setting batch_size to {batch_size}")
+    # can emulate a distributed environment for deepspeed to work in jupyter notebook
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = str(np.random.randint(10000)+9000)
+    os.environ["RANK"] = "0"
+    os.environ["LOCAL_RANK"] = "0"
+    os.environ["WORLD_SIZE"] = "1"
+    os.environ["GLOBAL_BATCH_SIZE"] = str(global_batch_size) # set this to your batch size!
+else:
+    global_batch_size = os.environ["GLOBAL_BATCH_SIZE"]    
+    batch_size = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices
+
+# alter the deepspeed config according to your global and local batch size
+if local_rank == 0:
+    with open('deepspeed_config_stage2.json', 'r') as file:
+        config = json.load(file)
+    config['train_batch_size'] = int(os.environ["GLOBAL_BATCH_SIZE"])
+    config['train_micro_batch_size_per_gpu'] = batch_size
+    config['bf16'] = {'enabled': False}
+    config['fp16'] = {'enabled': True}
+    with open('deepspeed_config_stage2.json', 'w') as file:
+        json.dump(config, file)
+else:
+    # give some time for the local_rank=0 gpu to prep new deepspeed config file
+    time.sleep(10)
+deepspeed_plugin = DeepSpeedPlugin("deepspeed_config_stage2_cpuoffload.json")
+accelerator = Accelerator(split_batches=False, deepspeed_plugin=deepspeed_plugin)
+
+
+# In[4]:
+
+
+print("PID of this process =",os.getpid())
+device = accelerator.device
+print("device:",device)
+num_workers = num_devices
+print(accelerator.state)
+world_size = accelerator.state.num_processes
+distributed = not accelerator.state.distributed_type == 'NO'
+
+# set data_type to match your mixed precision (automatically set based on deepspeed config)
+if accelerator.mixed_precision == "bf16":
+    data_type = torch.bfloat16
+elif accelerator.mixed_precision == "fp16":
+    data_type = torch.float16
+else:
+    data_type = torch.float32
+
+print("distributed =",distributed, "num_devices =", num_devices, "local rank =", local_rank, "world size =", world_size, "data_type =", data_type)
+print = accelerator.print # only print if local_rank=0
+
+
+# # Configurations
+
+# In[5]:
+
+
+# if running this interactively, can specify jupyter_args here for argparser to use
+if utils.is_interactive():
+    # create random model_name
+    model_name = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
+    model_name = model_name + "_interactive"
+    print("model_name:", model_name)
+
+    # global_batch_size and batch_size should already be defined in the above cells
+    # other variables can be specified in the following string:
+    jupyter_args = f"--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset \
+                    --model_name={model_name} \
+                    --subj=1 --batch_size={batch_size} --no-blurry_recon --no-depth_recon --hidden_dim=4096 \
+                    --clip_scale=1. --blur_scale=100. --depth_scale=100. \
+                    --max_lr=3e-4 --mixup_pct=.66 --num_epochs=12 --ckpt_interval=999 --no-use_image_aug --no-ckpt_saving"
+
+    jupyter_args = jupyter_args.split()
+    print(jupyter_args)
+    
+    from IPython.display import clear_output # function to clear print outputs in cell
+    get_ipython().run_line_magic('load_ext', 'autoreload')
+    # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions
+    get_ipython().run_line_magic('autoreload', '2')
+
+
+# In[6]:
+
+
+parser = argparse.ArgumentParser(description="Model Training Configuration")
+parser.add_argument(
+    "--model_name", type=str, default="testing",
+    help="name of model, used for ckpt saving and wandb logging (if enabled)",
+)
+parser.add_argument(
+    "--data_path", type=str, default="/fsx/proj-fmri/shared/natural-scenes-dataset",
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,5,7],
+)
+parser.add_argument(
+    "--batch_size", type=int, default=32,
+    help="Batch size can be increased by 10x if only training v2c and not diffusion diffuser",
+)
+parser.add_argument(
+    "--wandb_log",action=argparse.BooleanOptionalAction,default=False,
+    help="whether to log to wandb",
+)
+parser.add_argument(
+    "--resume_from_ckpt",action=argparse.BooleanOptionalAction,default=False,
+    help="if not using wandb and want to resume from a ckpt",
+)
+parser.add_argument(
+    "--wandb_project",type=str,default="stability",
+    help="wandb project name",
+)
+parser.add_argument(
+    "--mixup_pct",type=float,default=.33,
+    help="proportion of way through training when to switch from BiMixCo to SoftCLIP",
+)
+parser.add_argument(
+    "--blurry_recon",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to output blurry reconstructions",
+)
+parser.add_argument(
+    "--depth_recon",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to output depth reconstructions",
+)
+parser.add_argument(
+    "--blur_scale",type=float,default=100.,
+    help="multiply loss from blurry recons by this number",
+)
+parser.add_argument(
+    "--depth_scale",type=float,default=100.,
+    help="multiply loss from depth recons by this number",
+)
+parser.add_argument(
+    "--clip_scale",type=float,default=1.,
+    help="multiply contrastive loss by this number",
+)
+parser.add_argument(
+    "--use_image_aug",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to use image augmentation",
+)
+parser.add_argument(
+    "--num_epochs",type=int,default=120,
+    help="number of epochs of training",
+)
+parser.add_argument(
+    "--hidden_dim",type=int,default=4096,
+)
+parser.add_argument(
+    "--lr_scheduler_type",type=str,default='cycle',choices=['cycle','linear'],
+)
+parser.add_argument(
+    "--ckpt_saving",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--ckpt_interval",type=int,default=5,
+    help="save backup ckpt and reconstruct every x epochs",
+)
+parser.add_argument(
+    "--seed",type=int,default=42,
+)
+parser.add_argument(
+    "--max_lr",type=float,default=3e-4,
+)
+
+if utils.is_interactive():
+    args = parser.parse_args(jupyter_args)
+else:
+    args = parser.parse_args()
+
+# create global variables without the args prefix
+for attribute_name in vars(args).keys():
+    globals()[attribute_name] = getattr(args, attribute_name)
+
+
+# In[7]:
+
+
+outdir = os.path.abspath(f'../train_logs/{model_name}')
+if not os.path.exists(outdir) and ckpt_saving:
+    os.makedirs(outdir,exist_ok=True)
+if use_image_aug:
+    import kornia
+    from kornia.augmentation.container import AugmentationSequential
+    img_augment = AugmentationSequential(
+        kornia.augmentation.RandomResizedCrop((224,224), (0.6,1), p=0.3),
+        kornia.augmentation.Resize((224, 224)),
+        kornia.augmentation.RandomHorizontalFlip(p=0.3),
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.3),
+        kornia.augmentation.RandomGrayscale(p=0.3),
+        same_on_batch=False,
+        data_keys=["input"],
+    )
+
+
+# # Prep data, models, and dataloaders
+
+# ## Dataloader
+
+# In[8]:
+
+
+if subj==1:
+    num_train = 24958
+    num_test = 2770
+test_batch_size = num_test
+
+def my_split_by_node(urls): return urls
+    
+train_url = f"{data_path}/wds/subj0{subj}/train/" + "{0..36}.tar"
+# train_url = f"{data_path}/wds/subj0{subj}/train/" + "{0..1}.tar"
+print(train_url)
+
+train_data = wds.WebDataset(train_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=True, pin_memory=True)
+
+test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+print(test_url)
+
+test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+test_dl = torch.utils.data.DataLoader(test_data, batch_size=test_batch_size, shuffle=False, drop_last=True, pin_memory=True)
+
+
+# ### check dataloaders are working
+
+# In[9]:
+
+
+test_vox_indices = []
+test_73k_images = []
+for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+    test_vox_indices = np.append(test_vox_indices, behav[:,0,5].cpu().numpy())
+    test_73k_images = np.append(test_73k_images, behav[:,0,0].cpu().numpy())
+test_vox_indices = test_vox_indices.astype(np.int16)
+print(test_i, (test_i+1) * test_batch_size, len(test_vox_indices))
+print("---\n")
+
+train_vox_indices = []
+train_73k_images = []
+for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+    train_vox_indices = np.append(train_vox_indices, behav[:,0,5].long().cpu().numpy())
+    train_73k_images = np.append(train_73k_images, behav[:,0,0].cpu().numpy())
+train_vox_indices = train_vox_indices.astype(np.int16)
+print(train_i, (train_i+1) * batch_size, len(train_vox_indices))
+
+
+# ## Load data and images
+
+# In[10]:
+
+
+# load betas
+f = h5py.File(f'{data_path}/betas_all_subj0{subj}.hdf5', 'r')
+# f = h5py.File(f'{data_path}/betas_subj0{subj}_thresholded_wholebrain.hdf5', 'r')
+
+voxels = f['betas'][:]
+print(f"subj0{subj} betas loaded into memory")
+voxels = torch.Tensor(voxels).to("cpu").to(data_type)
+print("voxels", voxels.shape)
+num_voxels = voxels.shape[-1]
+
+# load orig images
+f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+images = f['images'][:]
+images = torch.Tensor(images).to("cpu").to(data_type)
+print("images", images.shape)
+
+
+# ## Load models
+
+# ### CLIP image embeddings  model
+
+# In[11]:
+
+
+from models import Clipper
+clip_model = Clipper("ViT-L/14", device=torch.device(f"cuda:{local_rank}"), hidden_state=True, norm_embs=True)
+clip_seq_dim = 257
+clip_emb_dim = 768 #1024
+# hidden_dim = 4096
+seq_len = 1 #2 #32 
+
+
+# ### SD VAE
+
+# In[12]:
+
+
+# if blurry_recon:
+#     from diffusers import AutoencoderKL
+#     autoenc = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, cache_dir="/fsx/proj-fmri/shared/cache")
+#     # autoenc.load_state_dict(torch.load('../train_logs/sdxl_vae_normed/best.pth')["model_state_dict"])
+#     autoenc.eval()
+#     autoenc.requires_grad_(False)
+#     autoenc.to(device)
+#     utils.count_params(autoenc)
+
+if blurry_recon:# or depth_recon:
+    from diffusers import VQModel
+    autoenc = VQModel.from_pretrained("/fsx/proj-fmri/shared/cache/models--microsoft--vq-diffusion-ithq/snapshots/3f796fb49ee559370dc638dea1d8116af131d993/vqvae", torch_dtype=data_type)
+    autoenc.eval()
+    autoenc.requires_grad_(False)
+    autoenc.to(device)
+    utils.count_params(autoenc)
+
+
+# #### downsampled images
+
+# In[13]:
+
+
+if blurry_recon:
+    if utils.is_interactive(): display(utils.torch_to_Image(images[[30]]))
+
+    input_batch = images[[30]].to(device)
+    print(input_batch.shape)
+
+    downsampled_image = nn.functional.interpolate(input_batch, size=(8, 8), mode='bilinear', align_corners=False)
+    re_upsampled_image = nn.functional.interpolate(downsampled_image, size=(128, 128), mode='nearest')
+    re_upsampled_enc = autoenc.encode(2*re_upsampled_image-1).latents * 0.18215
+    print(re_upsampled_enc.shape)
+    
+    if utils.is_interactive(): display(utils.torch_to_Image((autoenc.decode(re_upsampled_enc/0.18215).sample / 2 + 0.5).clamp(0,1)))
+
+
+# #### MiDaS depth
+
+# In[14]:
+
+
+if depth_recon:
+    from controlnet_aux.midas import MidasDetector
+    
+    midas_depth = MidasDetector.from_pretrained(
+      "valhalla/t2iadapter-aux-models", filename="dpt_large_384.pt", model_type="dpt_large", cache_dir="/fsx/proj-fmri/shared/cache").to(device)
+    midas_depth.model.eval()
+    midas_depth.model.requires_grad_(False)
+    midas_depth.model.to(device)
+    pass
+
+
+# In[15]:
+
+
+if depth_recon:
+    if utils.is_interactive(): display(utils.torch_to_Image(images[[30]]))
+
+    input_batch = images[[30,31]].float().to(device)
+    print(input_batch.shape)
+    
+    midas_emb = midas_depth.model(input_batch).unsqueeze(1)
+    print(midas_emb.shape)
+
+    prediction = utils.resize(midas_emb, 32) #/30).clamp(0,1).half() # 30 is roughly prediction.max()
+    print(prediction.shape)
+    
+    prediction = (prediction / prediction.view(prediction.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(prediction)).half()
+    midas_emb_size = prediction.flatten(1).shape[1]
+    print("midas_emb", prediction.shape, prediction.min(), prediction.max())
+    print("midas_emb_size", midas_emb_size)
+    
+    if utils.is_interactive(): display(utils.torch_to_Image(utils.resize(prediction, 224))) 
+
+    if blurry_recon:
+        prediction = utils.resize(midas_emb, 128).half().repeat(1,3,1,1)
+        prediction = (prediction / prediction.view(prediction.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(prediction)).half()
+        prediction_enc = autoenc.encode(2*prediction-1).latents * 0.18215
+        print("vae midas_emb", prediction_enc.shape, prediction_enc.min(), prediction_enc.max())
+    
+        if utils.is_interactive(): display(utils.torch_to_Image((autoenc.decode(prediction_enc/0.18215).sample / 2 + 0.5).clamp(0,1)))
+
+
+# ### MindEye modules
+
+# In[16]:
+
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+model
+
+
+# In[17]:
+
+
+time_embedding_dim = 512
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer
+    def __init__(self, input_size, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linear = torch.nn.Linear(input_size, out_features)
+    def forward(self, x):
+        return self.linear(x)
+        
+model.ridge = RidgeRegression(voxels.shape[1] + time_embedding_dim, out_features=hidden_dim)
+utils.count_params(model.ridge)
+utils.count_params(model)
+
+b = torch.randn((2,1,voxels.shape[1]))
+time_emb_test = torch.randn((2,1,time_embedding_dim))
+print(b.shape, model.ridge(torch.cat((b,time_emb_test),dim=-1)).shape)
+
+
+# In[59]:
+
+
+num_past_voxels = 15
+seq_len = 1 + 1
+
+
+# In[73]:
+
+
+class BrainNetwork(nn.Module):
+    def __init__(self, out_dim=768, in_dim=15724, seq_len=2, h=4096, n_blocks=4, drop=.15, clip_size=768):
+        super().__init__()
+        self.seq_len = seq_len
+        self.h = h
+        self.clip_size = clip_size
+        
+        # Initial linear layer to match the input dimensions to hidden dimensions
+        # self.lin0 = nn.Linear(in_dim, seq_len * h)
+        
+        # Mixer Blocks
+        self.mixer_blocks1 = nn.ModuleList([
+            self.mixer_block1(h, drop) for _ in range(n_blocks)
+        ])
+        self.mixer_blocks2 = nn.ModuleList([
+            self.mixer_block2(seq_len, drop) for _ in range(n_blocks)
+        ])
+        
+        # Output linear layer
+        self.clin1 = nn.Linear(h * seq_len, out_dim, bias=True)
+
+        # low-rank matrices
+        # self.rank = 500
+        # self.U = nn.Parameter(torch.randn(self.rank, out_dim))
+        # self.V = nn.Parameter(torch.randn(h * seq_len, self.rank))
+        # self.S = nn.Parameter(torch.randn(out_dim))
+
+        self.clip_proj = nn.Sequential(
+            nn.LayerNorm(clip_size),
+            nn.GELU(),
+            nn.Linear(clip_size, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, clip_size)
+        )
+
+        if blurry_recon:
+            # self.blin1 = nn.Sequential(
+            #     nn.Linear(out_dim, 4096, bias=True),
+            #     nn.LayerNorm(4096),
+            #     nn.GELU(),
+            #     nn.Linear(4096, 4096))
+            self.blin1 = nn.Linear(h*seq_len, 4096)
+            self.bgroupnorm = nn.GroupNorm(1, 256)
+            self.bupsampler = Decoder(
+                in_channels=256,
+                out_channels=128,
+                up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D"],
+                block_out_channels=[32, 64, 128],
+                layers_per_block=1,
+            )
+
+        if depth_recon:
+            # self.dlin1 = nn.Sequential(
+            #         nn.Linear(h, midas_emb_size),
+            #         nn.Sigmoid(),
+            #     )
+            self.dlin1 = nn.Linear(h*seq_len, 4096)
+            self.dgroupnorm = nn.GroupNorm(1, 256)
+            self.dupsampler = Decoder(
+                in_channels=256,
+                out_channels=1,#128,
+                up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D"],
+                block_out_channels=[32, 64, 128, 256],
+                layers_per_block=1,
+            )
+        
+    def mixer_block1(self, h, drop):
+        return nn.Sequential(
+            nn.LayerNorm(h),
+            self.mlp(h, h, drop),  # Token mixing
+        )
+
+    def mixer_block2(self, seq_len, drop):
+        return nn.Sequential(
+            nn.LayerNorm(seq_len),
+            self.mlp(seq_len, seq_len, drop)  # Channel mixing
+        )
+    
+    def mlp(self, in_dim, out_dim, drop):
+        return nn.Sequential(
+            nn.Linear(in_dim, out_dim),
+            nn.GELU(),
+            nn.Dropout(drop),
+            nn.Linear(out_dim, out_dim),
+        )
+        
+    def forward(self, x):
+        # make empty tensors for blur and depth outputs
+        b,d = torch.Tensor([0.]), torch.Tensor([0.])
+        
+        # Initial linear layer
+        # x = self.lin0(x)
+        
+        # Reshape to seq_len by dim
+        # x = x.reshape(-1, self.seq_len, self.h)
+        
+        # Mixer blocks
+        residual1 = x
+        residual2 = x.permute(0,2,1)
+        for block1, block2 in zip(self.mixer_blocks1,self.mixer_blocks2):
+            x = block1(x) + residual1
+            residual1 = x
+            x = x.permute(0,2,1)
+            
+            x = block2(x) + residual2
+            residual2 = x
+            x = x.permute(0,2,1)
+        
+        # Flatten
+        x = x.reshape(x.size(0), -1)
+        
+        c = self.clin1(x)
+
+        # low rank linear to out dim cuts # params by nearly half compared to full linear mapping
+        # c = (x @ (self.V/100) @ (self.U/100)) + self.S
+        
+        c = self.clip_proj(c.reshape(len(c), -1, self.clip_size))
+
+        if blurry_recon:
+            b = self.blin1(x)
+            b = b.reshape(len(b), 256, 4, 4)
+            b = self.bgroupnorm(b)
+            b = self.bupsampler(b)
+            
+        if depth_recon:
+            d = self.dlin1(x)#.reshape(len(x), 1, 32, 32)
+            d = d.reshape(len(d), 256, 4, 4)
+            d = self.dgroupnorm(d)
+            d = self.dupsampler(d)
+        
+        return c, b, d
+
+
+class TimeEmbedding(nn.Module):
+    def __init__(self, embedding_time_dim=512, num_past_voxels=15):
+        super().__init__()
+        self.embedding_time = nn.Embedding(num_past_voxels, embedding_time_dim)
+        self.num_past_voxels = num_past_voxels
+        self.embedding_time_dim = embedding_time_dim
+
+    def forward(self, time):
+        # time is (batch_size,)
+        time = time.long()
+        time = self.embedding_time(time)
+        return time # (batch_size, embedding_time_dim)
+    
+
+#model.memory_encoder = MemoryEncoder(in_dim=voxels.shape[1], out_dim=4096, num_past_voxels=15, embedding_time_dim=512)
+model.time_embedding = TimeEmbedding(embedding_time_dim=512, num_past_voxels=15)
+
+model.backbone = BrainNetwork(h=1024, in_dim=1024, seq_len=4, clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim) 
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+
+# test that the model works on some fake data
+b = torch.randn((256,4,1024))
+print("b.shape",b.shape)
+with torch.no_grad():
+    clip_, blur_, depth_ = model.backbone(b)
+print(clip_.shape, blur_.shape, depth_.shape)
+
+
+# In[70]:
+
+
+voxel_ridge = torch.randn(512,4096)
+voxel_ridge = voxel_ridge.view(int(voxel_ridge.shape[0]/seq_len), seq_len, hidden_dim)
+print("b.shape",voxel_ridge.shape)
+with torch.no_grad():
+    clip_, blur_, depth_ = model.backbone(voxel_ridge)
+print(clip_.shape, blur_.shape, depth_.shape)
+
+
+# In[64]:
+
+
+no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+opt_grouped_parameters = [
+    {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
+]
+
+optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr)
+
+if lr_scheduler_type == 'linear':
+    lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+        optimizer,
+        total_iters=int(np.floor(num_epochs*(num_train/num_devices/batch_size))),
+        last_epoch=-1
+    )
+elif lr_scheduler_type == 'cycle':
+    total_steps=int(np.floor(num_epochs*(num_train/num_devices/batch_size)))
+    print("total_steps", total_steps)
+    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer, 
+        max_lr=max_lr,
+        total_steps=total_steps,
+        final_div_factor=1000,
+        last_epoch=-1, pct_start=2/num_epochs
+    )
+    
+def save_ckpt(tag):    
+    ckpt_path = outdir+f'/{tag}.pth'
+    print(f'saving {ckpt_path}',flush=True)
+    unwrapped_model = accelerator.unwrap_model(model)
+    try:
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': unwrapped_model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'lr_scheduler': lr_scheduler.state_dict(),
+            'train_losses': losses,
+            'test_losses': test_losses,
+            'lrs': lrs,
+            }, ckpt_path)
+    except:
+        print("Couldn't save... moving on to prevent crashing.")
+    del unwrapped_model
+        
+print("\nDone with model preparations!")
+utils.count_params(model)
+
+
+# In[49]:
+
+
+seq_len = 4
+
+
+# In[57]:
+
+
+voxel_ridge = torch.randn(512,4096)
+voxel_ridge = voxel_ridge.view(int(voxel_ridge.shape[0]/seq_len), seq_len, hidden_dim)
+
+
+# In[58]:
+
+
+voxel_ridge.shape
+
+
+# In[55]:
+
+
+pp = None
+for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+    with torch.cuda.amp.autocast(dtype=data_type):
+        #optimizer.zero_grad()
+
+        voxel = voxels[behav[:,0,5].cpu().long()]#.to(device)
+        image = images[behav[:,0,0].cpu().long()].float()#.to(device).float()
+
+        past_15_voxels = voxels[past_behav[:,:seq_len-1,5].cpu().long()]#.to(device) # batch_size, 15, 15279
+        past_15_times = torch.Tensor([i for i in range(seq_len)])#.to(device) # 15
+    
+        print(past_15_times)
+        #for past in range(1):
+        #    past_voxel = voxels[past_behav[:,past,5].cpu().long()].to(device)
+        """
+        #if blurry_recon:
+        #    blurry_image_enc = autoenc.encode(2*utils.resize(image,128)-1).latent_dist.mode() * 0.18215
+        blurry_image_enc = autoenc.encode(2*utils.resize(add_saturation(image),128)-1).latents * 0.18215
+
+        if depth_recon:
+            # depth_images = utils.resize(midas_depth.model(image).unsqueeze(1).repeat(1,3,1,1), 128)
+            depth_images = utils.resize(midas_depth.model(image).unsqueeze(1), 32)
+            depth_images = (depth_images / depth_images.view(depth_images.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(depth_images)).half()
+            depth_image_enc = depth_images # autoenc.encode(2*depth_images-1).latents * 0.18215
+
+        if use_image_aug: 
+            image = img_augment(image)
+
+        clip_target = clip_model.embed_image(image)
+        assert not torch.any(torch.isnan(clip_target))
+
+        if epoch < int(mixup_pct * num_epochs):
+            voxel, perm, betas, select = utils.mixco(voxel)
+            past_voxel, _, _, _ = utils.mixco(voxel, perm=perm, betas=betas, select=select)
+        """
+        for p in range(seq_len-1):
+            print(past_behav.shape) #128, 15, 17
+            print(past_behav[:,p,-1])
+            print(past_15_voxels.shape) # 128, 1, 15724
+            mask = past_behav[:,p,-1] == torch.ones_like(past_behav[:,p,-1])
+            print(mask) # 128
+            past_15_voxels[mask, p, :] = torch.zeros_like(past_15_voxels[0, p, :])
+            print(past_15_voxels)
+            pp = past_15_voxels
+            
+        break
+
+
+# In[54]:
+
+
+pp[20, 0, :]
+
+
+# # Weights and Biases
+
+# In[66]:
+
+
+if local_rank==0 and wandb_log: # only use main process for wandb logging
+    import wandb
+    wandb_project = 'mindeyev2'
+    wandb_run = model_name
+    wandb_notes = ''
+    
+    print(f"wandb {wandb_project} run {wandb_run}")
+    wandb.login(host='https://stability.wandb.io')#, relogin=True)
+    wandb_config = {
+      "model_name": model_name,
+      "global_batch_size": global_batch_size,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "clip_scale": clip_scale,
+      "blur_scale": blur_scale,
+      "use_image_aug": use_image_aug,
+      "max_lr": max_lr,
+      "mixup_pct": mixup_pct,
+      "num_train": num_train,
+      "num_test": num_test,
+      "ckpt_interval": ckpt_interval,
+      "ckpt_saving": ckpt_saving,
+      "seed": seed,
+      "distributed": distributed,
+      "num_devices": num_devices,
+      "world_size": world_size,
+      "train_url": train_url,
+      "test_url": test_url,
+    }
+    print("wandb_config:\n",wandb_config)
+    if True: # wandb_auto_resume
+        print("wandb_id:",model_name)
+        wandb.init(
+            id = model_name,
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+            resume="allow",
+        )
+    else:
+        wandb.init(
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+        )
+else:
+    wandb_log = False
+
+
+# # Main
+
+# In[67]:
+
+
+epoch = 0
+losses, test_losses, lrs = [], [], []
+best_test_loss = 1e9
+soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))
+
+# Optionally resume from checkpoint #
+if resume_from_ckpt:
+    print("\n---resuming from last.pth ckpt---\n")
+    try:
+        checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+    except:
+        print('last.pth failed... trying last_backup.pth')
+        checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+    epoch = checkpoint['epoch']
+    print("Epoch",epoch)
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    model.load_state_dict(checkpoint['model_state_dict'])
+    del checkpoint
+elif wandb_log:
+    if wandb.run.resumed:
+        print("\n---resuming from last.pth ckpt---\n")
+        try:
+            checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+        except:
+            print('last.pth failed... trying last_backup.pth')
+            checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+        epoch = checkpoint['epoch']
+        print("Epoch",epoch)
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+        model.load_state_dict(checkpoint['model_state_dict'])
+        del checkpoint
+torch.cuda.empty_cache()
+
+
+# In[68]:
+
+
+model, optimizer, train_dl, lr_scheduler = accelerator.prepare(
+model, optimizer, train_dl, lr_scheduler
+)
+# leaving out test_dl since we will only have local_rank 0 device do evals
+
+
+# In[ ]:
+
+
+def add_saturation(image, alpha=2):
+    gray_image = 0.2989 * image[:, 0, :, :] + 0.5870 * image[:, 1, :, :] + 0.1140 * image[:, 2, :, :]
+    gray_image = gray_image.unsqueeze(1).expand_as(image)
+    saturated_image = alpha * image + (1 - alpha) * gray_image
+    return torch.clamp(saturated_image, 0, 1)
+
+
+# In[65]:
+
+
+print(f"{model_name} starting with epoch {epoch} / {num_epochs}")
+progress_bar = tqdm(range(epoch,num_epochs), ncols=1200, disable=(local_rank!=0))
+test_image, test_voxel = None, None
+mse = nn.MSELoss()
+l1 = nn.L1Loss()
+
+for epoch in progress_bar:
+    model.train()
+    
+    fwd_percent_correct = 0.
+    bwd_percent_correct = 0.
+    test_fwd_percent_correct = 0.
+    test_bwd_percent_correct = 0.
+
+    loss_clip_total = 0.
+    loss_blurry_total = 0.
+    loss_depth_total = 0.
+    test_loss_clip_total = 0.
+    test_loss_blurry_total = 0.
+    test_loss_depth_total = 0.
+
+    blurry_pixcorr = 0.
+    test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1
+    
+    for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+        with torch.cuda.amp.autocast(dtype=data_type):
+            optimizer.zero_grad()
+    
+            #voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+            #image = images[behav[:,0,0].cpu().long()].to(device).float()
+            
+            #past_15_voxels = voxels[past_behav[:,:,5].cpu().long()].to(device) # batch_size, 15, 15279
+            #past_15_times = torch.Tensor([i for i in range(seq_len - 1)]).to(device) # 15
+            
+            voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+            image = images[behav[:,0,0].cpu().long()].to(device).float()
+
+            past_15_voxels = voxels[past_behav[:,:seq_len-1,5].cpu().long()].to(device) # batch_size, 15, 15279
+            past_15_times = torch.Tensor([i for i in range(seq_len-1)]).to(device) # 15
+            #for past in range(1):
+            #    past_voxel = voxels[past_behav[:,past,5].cpu().long()].to(device)
+    
+            if blurry_recon:
+                # blurry_image_enc = autoenc.encode(2*utils.resize(image,128)-1).latent_dist.mode() * 0.18215
+                blurry_image_enc = autoenc.encode(2*utils.resize(add_saturation(image),128)-1).latents * 0.18215
+
+            if depth_recon:
+                # depth_images = utils.resize(midas_depth.model(image).unsqueeze(1).repeat(1,3,1,1), 128)
+                depth_images = utils.resize(midas_depth.model(image).unsqueeze(1), 32)
+                depth_images = (depth_images / depth_images.view(depth_images.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(depth_images)).half()
+                depth_image_enc = depth_images # autoenc.encode(2*depth_images-1).latents * 0.18215
+            
+            if use_image_aug: 
+                image = img_augment(image)
+    
+            clip_target = clip_model.embed_image(image)
+            assert not torch.any(torch.isnan(clip_target))
+    
+            if epoch < int(mixup_pct * num_epochs):
+                voxel, perm, betas, select = utils.mixco(voxel)
+                past_voxel, _, _, _ = utils.mixco(voxel, perm=perm, betas=betas, select=select)
+                
+            for p in range(seq_len-1):
+                #print(past_behav.shape) #128, 15, 17
+                #print(past_behav[:,p,-1])
+                #print(past_15_voxels.shape) # 128, 1, 15724
+                mask = past_behav[:,p,-1] == torch.ones_like(past_behav[:,p,-1])
+                #print(mask) # 128
+                past_15_voxels[mask, p, :] = torch.zeros_like(past_15_voxels[0, p, :])
+                #print(past_15_voxels)
+                
+            past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+            past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+            past_15_times = past_15_times.reshape(-1)
+            time_embeddings = model.time_embedding(past_15_times)
+            past_info_full = torch.cat((past_15_voxels, time_embeddings), dim=-1)
+            
+            positional_current_voxel = torch.zeros((voxel.shape[0], time_embeddings.shape[-1])).to(voxel.device)
+            voxel = torch.cat((voxel, positional_current_voxel), dim=-1)
+            voxel_ridge = model.ridge(torch.cat((voxel, past_info_full), dim=-2))
+            voxel_ridge = voxel_ridge.view(int(voxel_ridge.shape[0]/seq_len), seq_len, hidden_dim)
+            #unsqueeze(1) # bz * 2, 1, 4096
+            
+            # past_voxel_ridge = model.ridge(past_voxel)
+            # voxel_ridge = torch.cat((voxel_ridge.unsqueeze(1), past_voxel_ridge.unsqueeze(1)), axis=1)
+            print(voxel_ridge.shape)
+    
+            clip_voxels, blurry_image_enc_, depth_image_enc_ = model.backbone(voxel_ridge)
+            
+            clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+            clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+    
+            if epoch < int(mixup_pct * num_epochs):                
+                loss_clip = utils.mixco_nce(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006, 
+                    perm=perm, betas=betas, select=select)
+            else:
+                epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=epoch_temp)
+
+            loss_clip_total += loss_clip.item()
+            loss_clip *= clip_scale
+            loss = loss_clip
+    
+            if blurry_recon:
+                downsampled_image = nn.functional.interpolate(image, size=(8, 8), mode='bilinear', align_corners=False)
+                re_upsampled_image = add_saturation(nn.functional.interpolate(downsampled_image, size=(128, 128), mode='nearest'))
+                re_upsampled_enc = autoenc.encode(2*re_upsampled_image-1).latents * 0.18215
+                
+                loss_blurry = (l1(blurry_image_enc_, blurry_image_enc) + l1(blurry_image_enc_, re_upsampled_enc))
+                loss_blurry += l1(torch.var(blurry_image_enc), torch.var(blurry_image_enc_))
+                loss_blurry_total += loss_blurry.item()
+                loss_blurry *= blur_scale
+                loss += loss_blurry
+
+            if depth_recon:
+                loss_depth = l1(depth_image_enc_, depth_image_enc)
+                # loss_depth += l1(torch.var(depth_image_enc_), torch.var(depth_image_enc))
+                loss_depth_total += loss_depth.item()
+                loss_depth *= depth_scale
+                loss += loss_depth
+    
+            # forward and backward top 1 accuracy        
+            labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+            fwd_percent_correct += utils.topk(torch.abs(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm)), labels, k=1).item()
+            bwd_percent_correct += utils.topk(torch.abs(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm)), labels, k=1).item()
+    
+            if blurry_recon:
+                with torch.no_grad():
+                    # only doing pixcorr eval on a subset of the samples per batch because its costly & slow to compute autoenc.decode()
+                    random_samps = np.random.choice(np.arange(len(voxel)), size=batch_size//5, replace=False)
+                    # random_samps = np.arange(batch_size//5)
+                    blurry_recon_images = (autoenc.decode(blurry_image_enc_[random_samps]/0.18215).sample/ 2 + 0.5).clamp(0,1)
+                    # pixcorr_origsize_nanmean is computationally less intense than utils.pixcorr and uses nanmean instead of mean
+                    pixcorr = utils.pixcorr_origsize_nanmean(image[random_samps], blurry_recon_images)
+                    # pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)
+                    # loss += (1 - pixcorr)
+                    blurry_pixcorr += pixcorr.item()
+                    # utils.check_loss(pixcorr)
+
+            utils.check_loss(loss)
+            accelerator.backward(loss)
+            optimizer.step()
+    
+            losses.append(loss.item())
+            lrs.append(optimizer.param_groups[0]['lr'])
+    
+            if lr_scheduler_type is not None:
+                lr_scheduler.step()
+
+    model.eval()
+    if local_rank==0:
+        with torch.no_grad(), torch.cuda.amp.autocast(dtype=data_type): 
+            for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):  
+                # all test samples should be loaded per batch such that test_i should never exceed 0
+                assert len(behav) == num_test
+                
+                ## Average same-image repeats ##
+                if test_image is None:
+                    voxel = voxels[behav[:,0,5].cpu().long()]
+                    image = behav[:,0,0].cpu().long()
+                    
+                    unique_image, sort_indices = torch.unique(image, return_inverse=True)
+                    for im in unique_image:
+                        locs = torch.where(im == image)[0]
+                        if test_image is None:
+                            test_image = images[im][None]
+                            test_voxel = torch.mean(voxel[locs],axis=0)[None]
+                        else:
+                            test_image = torch.vstack((test_image, images[im][None]))
+                            test_voxel = torch.vstack((test_voxel, torch.mean(voxel[locs],axis=0)[None]))
+    
+                # random sample of 300
+                random_indices = torch.arange(len(test_voxel))[:300]
+                voxel = test_voxel[random_indices].to(device)
+                image = test_image[random_indices].to(device)
+                assert len(image) == 300
+                
+                current_past_behav = past_behav[random_indices]
+
+                past_15_voxels = voxels[current_past_behav[:,:seq_len-1,5].cpu().long()].to(device) # batch_size, 15, 15279
+                past_15_times = torch.Tensor([i for i in range(seq_len-1)]).to(device) # 15
+
+                if blurry_recon:
+                    # blurry_image_enc = autoenc.encode(2*utils.resize(image,128)-1).latent_dist.mode() * 0.18215
+                    blurry_image_enc = autoenc.encode(2*utils.resize(add_saturation(image),128)-1).latents * 0.18215
+
+                if depth_recon:
+                    # depth_images = utils.resize(midas_depth.model(image).unsqueeze(1).repeat(1,3,1,1), 128)
+                    depth_images = utils.resize(midas_depth.model(image).unsqueeze(1), 32)
+                    depth_images = (depth_images / depth_images.view(depth_images.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(depth_images)).half()
+                    depth_image_enc = depth_images # autoenc.encode(2*depth_images-1).latents * 0.18215
+            
+                clip_target = clip_model.embed_image(image.float())
+                
+
+                past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+                past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+                past_15_times = past_15_times.reshape(-1)
+                time_embeddings = model.time_embedding(past_15_times)
+                past_info_full = torch.cat((past_15_voxels, time_embeddings), dim=-1)
+
+                positional_current_voxel = torch.zeros((voxel.shape[0], time_embeddings.shape[-1])).to(voxel.device)
+                voxel = torch.cat((voxel, positional_current_voxel), dim=-1)
+                voxel_ridge = model.ridge(torch.cat((voxel, past_info_full), dim=-2)).unsqueeze(1)
+                
+                #voxel_ridge = model.ridge(voxel).unsqueeze(1)
+
+                # voxel_ridge = torch.cat((voxel_ridge.unsqueeze(1),voxel_ridge.unsqueeze(1)),axis=1)
+                
+                clip_voxels, blurry_image_enc_, depth_image_enc_ = model.backbone(voxel_ridge)
+                
+                clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+                clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+        
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006)
+                test_loss_clip_total += loss_clip.item()
+                loss_clip = loss_clip * clip_scale
+                loss = loss_clip
+
+                if blurry_recon:
+                    downsampled_image = nn.functional.interpolate(image, size=(8, 8), mode='bilinear', align_corners=False)
+                    re_upsampled_image = add_saturation(nn.functional.interpolate(downsampled_image, size=(128, 128), mode='nearest'))
+                    re_upsampled_enc = autoenc.encode(2*re_upsampled_image-1).latents * 0.18215
+                    
+                    loss_blurry = (l1(blurry_image_enc_, blurry_image_enc) + l1(blurry_image_enc_, re_upsampled_enc))
+                    loss_blurry += l1(torch.var(blurry_image_enc), torch.var(blurry_image_enc_))
+                    test_loss_blurry_total += loss_blurry.item()
+                    loss_blurry *= blur_scale
+                    loss += loss_blurry
+    
+                    # halving the batch size because the decoder is computationally heavy
+                    blurry_recon_images = (autoenc.decode(blurry_image_enc_[:len(voxel)//2]/0.18215).sample / 2 + 0.5).clamp(0,1)
+                    blurry_recon_images = torch.vstack((blurry_recon_images, (autoenc.decode(blurry_image_enc_[len(voxel)//2:]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                    pixcorr = utils.pixcorr(image, blurry_recon_images)
+                    loss += (1 - pixcorr)
+                    test_blurry_pixcorr += pixcorr.item()
+
+                if depth_recon:
+                    loss_depth = l1(depth_image_enc_, depth_image_enc)
+                    # loss_depth += l1(torch.var(depth_image_enc_), torch.var(depth_image_enc))
+                    test_loss_depth_total += loss_depth.item()
+                    loss_depth *= depth_scale
+                    loss += loss_depth
+        
+                # forward and backward top 1 accuracy        
+                labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+                test_fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item()
+                test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()
+
+                utils.check_loss(loss)                
+                test_losses.append(loss.item())
+
+            # if utils.is_interactive(): clear_output(wait=True)
+            print("---")
+            
+            assert (test_i+1) == 1
+            logs = {"train/loss": np.mean(losses[-(train_i+1):]),
+                "test/loss": np.mean(test_losses[-(test_i+1):]),
+                "train/lr": lrs[-1],
+                "train/num_steps": len(losses),
+                "test/num_steps": len(test_losses),
+                "train/fwd_pct_correct": fwd_percent_correct / (train_i + 1),
+                "train/bwd_pct_correct": bwd_percent_correct / (train_i + 1),
+                "test/test_fwd_pct_correct": test_fwd_percent_correct / (test_i + 1),
+                "test/test_bwd_pct_correct": test_bwd_percent_correct / (test_i + 1),
+                "train/loss_clip_total": loss_clip_total / (train_i + 1),
+                "train/loss_blurry_total": loss_blurry_total / (train_i + 1),
+                "test/loss_clip_total": test_loss_clip_total / (test_i + 1),
+                "test/loss_blurry_total": test_loss_blurry_total / (test_i + 1),
+                "train/blurry_pixcorr": blurry_pixcorr / (train_i + 1),
+                "test/blurry_pixcorr": test_blurry_pixcorr / (test_i + 1),
+                "train/loss_depth_total": loss_depth_total / (train_i + 1),
+                "test/loss_depth_total": test_loss_depth_total / (test_i + 1),
+                }
+    
+            if blurry_recon:    
+                # transform blurry recon latents to images and plot it
+                fig, axes = plt.subplots(1, 8, figsize=(10, 4))
+                jj=-1
+                for j in [0,1,2,3]:
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image((autoenc.decode(blurry_image_enc[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                    axes[jj].axis('off')
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image((autoenc.decode(blurry_image_enc_[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                    axes[jj].axis('off')
+                
+                if wandb_log:
+                    logs[f"test/recons"] = wandb.Image(fig, caption=f"epoch{epoch:03d}")
+                    plt.close()
+                else:
+                    plt.show()
+
+            if depth_recon:
+                # transform blurry recon latents to images and plot it
+                fig, axes = plt.subplots(1, 8, figsize=(10, 4))
+                # axes[0].imshow(utils.torch_to_Image((autoenc.decode(depth_image_enc[[0]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                # axes[1].imshow(utils.torch_to_Image((autoenc.decode(depth_image_enc_[[0]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                jj=-1
+                for j in [0,1,2,3]:
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image(utils.resize(depth_image_enc[[j]].view(1,1,32,32).clamp(0,1), 224)))
+                    axes[jj].axis('off')
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image(utils.resize(depth_image_enc_[[j]].view(1,1,32,32).clamp(0,1), 224)))
+                    axes[jj].axis('off')
+                if wandb_log:
+                    logs[f"test/depth_recons"] = wandb.Image(fig, caption=f"epoch{epoch:03d}")
+                    plt.close()
+                else:
+                    plt.show()
+            
+            progress_bar.set_postfix(**logs)
+    
+            # Save model checkpoint and reconstruct
+            if epoch % ckpt_interval == 0:
+                if not utils.is_interactive():
+                    save_ckpt(f'last')
+                    
+            if wandb_log: wandb.log(logs)
+
+    # wait for other GPUs to catch up if needed
+    accelerator.wait_for_everyone()
+    torch.cuda.empty_cache()
+    gc.collect()
+
+print("\n===Finished!===\n")
+if ckpt_saving:
+    save_ckpt(f'last')
+if not utils.is_interactive():
+    sys.exit(0)
+
+
+# In[ ]:
+
+
+plt.plot(losses)
+plt.show()
+plt.plot(test_losses)
+plt.show()
+
+
+# # Retrieve nearest neighbor in the training set using test set data
+
+# In[ ]:
+
+
+annots = np.load("/fsx/proj-fmri/shared/mindeyev2_dataset/COCO_73k_annots_curated.npy")
+
+
+# In[ ]:
+
+
+ii=2
+all_indices = np.unique(train_73k_images) #np.hstack((test_vox_indices[ii],train_vox_indices))
+with torch.no_grad(), torch.cuda.amp.autocast():
+    for batch in tqdm(range(0,len(all_indices),512)):
+        if batch==0:
+            clip_target = clip_model.embed_image(images[all_indices[batch:batch+512]]).cpu()
+        else:
+            target = clip_model.embed_image(images[all_indices[batch:batch+512]]).cpu()
+            clip_target = torch.vstack((clip_target,target))
+    clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+    voxel = test_voxel[[ii]].to(device)
+    image = test_image[[ii]].to(device)
+
+    print("Original Image (test set)")
+    display(utils.torch_to_Image(image))
+    
+    clip_target = clip_model.embed_image(image).cpu()
+    # clip_target_norm = torch.vstack((clip_target_norm, nn.functional.normalize(clip_target.flatten(1), dim=-1)))
+    
+    voxel_ridge = model.ridge(voxel).unsqueeze(1)
+    clip_voxels, _, _ = model.backbone(voxel_ridge)    
+    clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+    clip_voxels_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+    print("clip_voxels_norm", clip_voxels_norm.shape)
+    print("clip_target_norm", clip_target_norm.shape)
+    
+    sortt = torch.argsort(utils.batchwise_cosine_similarity(clip_voxels_norm.cpu(), 
+                                                            clip_target_norm).flatten()).flip(0)
+    picks = all_indices[sortt[:5]]
+
+    print("\nNearest neighbors in training set")
+    for ip,p in enumerate(picks):
+        display(utils.torch_to_Image(images[[p]]))
+        # print(utils.select_annotations([annots[int(p)]]))
+        if ip==0: predicted_caption = utils.select_annotations([annots[int(p)]])[0]
+
+print("\n=====\npredicted_caption:\n", predicted_caption)
+
+
+# # Feed into Stable Diffusion XL for reconstructions
+
+# In[ ]:
+
+
+from diffusers import StableDiffusionXLPipeline
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    "/fsx/proj-fmri/shared/cache/models--stabilityai--stable-diffusion-xl-base-1.0/snapshots/f898a3e026e802f68796b95e9702464bac78d76f", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
+)
+pipe.to("cuda")
+pass
+
+
+# In[ ]:
+
+
+prompt = predicted_caption
+recon = pipe(prompt=prompt).images[0]
+
+
+# In[ ]:
+
+
+print("Seen image")
+display(utils.torch_to_Image(image))
+
+print("Reconstruction")
+utils.torch_to_Image(utils.resize(transforms.ToTensor()(recon),224))
+

src/Train_MLPMixer-Copy2.py

@@ -0,0 +1,1275 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+# # Code to convert this notebook to .py if you want to run it via command line or with Slurm
+#from subprocess import call
+#command = "jupyter nbconvert Train_MLPMixer-Copy1.ipynb --to python"
+#call(command,shell=True)
+
+
+# # Import packages & functions
+
+# In[2]:
+
+
+import os
+import sys
+import json
+import argparse
+import numpy as np
+import math
+from einops import rearrange
+import time
+import random
+import string
+import h5py
+from tqdm import tqdm
+
+import webdataset as wds
+import gc
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+
+from accelerate import Accelerator, DeepSpeedPlugin
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+
+
+# In[3]:
+
+
+### Multi-GPU config ###
+local_rank = os.getenv('RANK')
+if local_rank is None: 
+    local_rank = 0
+else:
+    local_rank = int(local_rank)
+print("LOCAL RANK ", local_rank)  
+
+num_devices = torch.cuda.device_count()
+if num_devices==0: num_devices = 1
+
+# ## UNCOMMENT BELOW SECTION AND COMMENT OUT DEEPSPEED SECTION TO AVOID USING DEEPSPEED ###
+# accelerator = Accelerator(split_batches=False, mixed_precision="fp16")
+# global_batch_size = batch_size = 32
+# data_type = torch.float16 # change depending on your mixed_precision
+
+### DEEPSPEED INITIALIZATION ###
+if num_devices <= 1 and utils.is_interactive():
+    global_batch_size = batch_size = 32
+    print(f"Setting batch_size to {batch_size}")
+    # can emulate a distributed environment for deepspeed to work in jupyter notebook
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = str(np.random.randint(10000)+9000)
+    os.environ["RANK"] = "0"
+    os.environ["LOCAL_RANK"] = "0"
+    os.environ["WORLD_SIZE"] = "1"
+    os.environ["GLOBAL_BATCH_SIZE"] = str(global_batch_size) # set this to your batch size!
+else:
+    global_batch_size = os.environ["GLOBAL_BATCH_SIZE"]    
+    batch_size = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices
+
+# alter the deepspeed config according to your global and local batch size
+if local_rank == 0:
+    with open('/fsx/proj-fmri/ckadirt/MindEyeV2/src/deepspeed_config_stage2.json', 'r') as file:
+        config = json.load(file)
+    config['train_batch_size'] = int(os.environ["GLOBAL_BATCH_SIZE"])
+    config['train_micro_batch_size_per_gpu'] = batch_size
+    config['bf16'] = {'enabled': False}
+    config['fp16'] = {'enabled': True}
+    with open('/fsx/proj-fmri/ckadirt/MindEyeV2/src/deepspeed_config_stage2.json', 'w') as file:
+        json.dump(config, file)
+else:
+    # give some time for the local_rank=0 gpu to prep new deepspeed config file
+    time.sleep(10)
+deepspeed_plugin = DeepSpeedPlugin("/fsx/proj-fmri/ckadirt/MindEyeV2/src/deepspeed_config_stage2.json")
+accelerator = Accelerator(split_batches=False, deepspeed_plugin=deepspeed_plugin)
+
+
+# In[4]:
+
+
+print("PID of this process =",os.getpid())
+device = accelerator.device
+print("device:",device)
+num_workers = num_devices
+print(accelerator.state)
+world_size = accelerator.state.num_processes
+distributed = not accelerator.state.distributed_type == 'NO'
+
+# set data_type to match your mixed precision (automatically set based on deepspeed config)
+if accelerator.mixed_precision == "bf16":
+    data_type = torch.bfloat16
+elif accelerator.mixed_precision == "fp16":
+    data_type = torch.float16
+else:
+    data_type = torch.float32
+
+print("distributed =",distributed, "num_devices =", num_devices, "local rank =", local_rank, "world size =", world_size, "data_type =", data_type)
+print = accelerator.print # only print if local_rank=0
+
+
+# # Configurations
+
+# In[5]:
+
+
+# if running this interactively, can specify jupyter_args here for argparser to use
+if utils.is_interactive():
+    # create random model_name
+    model_name = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
+    model_name = model_name + "_interactive"
+    print("model_name:", model_name)
+
+    # global_batch_size and batch_size should already be defined in the above cells
+    # other variables can be specified in the following string:
+    jupyter_args = f"--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset \
+                    --model_name={model_name} \
+                    --subj=1 --batch_size={batch_size} --no-blurry_recon --no-depth_recon --hidden_dim=4096 \
+                    --clip_scale=1. --blur_scale=100. --depth_scale=100. \
+                    --max_lr=3e-4 --mixup_pct=.66 --num_epochs=12 --ckpt_interval=999 --no-use_image_aug --no-ckpt_saving"
+
+    jupyter_args = jupyter_args.split()
+    print(jupyter_args)
+    
+    from IPython.display import clear_output # function to clear print outputs in cell
+    get_ipython().run_line_magic('load_ext', 'autoreload')
+    # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions
+    get_ipython().run_line_magic('autoreload', '2')
+
+
+# In[6]:
+
+
+parser = argparse.ArgumentParser(description="Model Training Configuration")
+parser.add_argument(
+    "--model_name", type=str, default="testing",
+    help="name of model, used for ckpt saving and wandb logging (if enabled)",
+)
+parser.add_argument(
+    "--data_path", type=str, default="/fsx/proj-fmri/shared/natural-scenes-dataset",
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,5,7],
+)
+parser.add_argument(
+    "--batch_size", type=int, default=32,
+    help="Batch size can be increased by 10x if only training v2c and not diffusion diffuser",
+)
+parser.add_argument(
+    "--wandb_log",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to log to wandb",
+)
+parser.add_argument(
+    "--resume_from_ckpt",action=argparse.BooleanOptionalAction,default=False,
+    help="if not using wandb and want to resume from a ckpt",
+)
+parser.add_argument(
+    "--wandb_project",type=str,default="stability",
+    help="wandb project name",
+)
+parser.add_argument(
+    "--mixup_pct",type=float,default=.33,
+    help="proportion of way through training when to switch from BiMixCo to SoftCLIP",
+)
+parser.add_argument(
+    "--blurry_recon",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to output blurry reconstructions",
+)
+parser.add_argument(
+    "--depth_recon",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to output depth reconstructions",
+)
+parser.add_argument(
+    "--blur_scale",type=float,default=100.,
+    help="multiply loss from blurry recons by this number",
+)
+parser.add_argument(
+    "--depth_scale",type=float,default=100.,
+    help="multiply loss from depth recons by this number",
+)
+parser.add_argument(
+    "--clip_scale",type=float,default=1.,
+    help="multiply contrastive loss by this number",
+)
+parser.add_argument(
+    "--use_image_aug",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to use image augmentation",
+)
+parser.add_argument(
+    "--num_epochs",type=int,default=120,
+    help="number of epochs of training",
+)
+parser.add_argument(
+    "--hidden_dim",type=int,default=4096,
+)
+parser.add_argument(
+    "--lr_scheduler_type",type=str,default='cycle',choices=['cycle','linear'],
+)
+parser.add_argument(
+    "--ckpt_saving",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--ckpt_interval",type=int,default=5,
+    help="save backup ckpt and reconstruct every x epochs",
+)
+parser.add_argument(
+    "--seed",type=int,default=42,
+)
+parser.add_argument(
+    "--max_lr",type=float,default=3e-4,
+)
+parser.add_argument(
+    "--seq_len",type=int,default=2,
+)
+
+if utils.is_interactive():
+    args = parser.parse_args(jupyter_args)
+else:
+    args = parser.parse_args()
+
+# create global variables without the args prefix
+for attribute_name in vars(args).keys():
+    globals()[attribute_name] = getattr(args, attribute_name)
+
+
+# In[7]:
+
+
+outdir = os.path.abspath(f'../train_logs/{model_name}')
+if not os.path.exists(outdir) and ckpt_saving:
+    os.makedirs(outdir,exist_ok=True)
+if use_image_aug:
+    import kornia
+    from kornia.augmentation.container import AugmentationSequential
+    img_augment = AugmentationSequential(
+        kornia.augmentation.RandomResizedCrop((224,224), (0.6,1), p=0.3),
+        kornia.augmentation.Resize((224, 224)),
+        kornia.augmentation.RandomHorizontalFlip(p=0.3),
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.3),
+        kornia.augmentation.RandomGrayscale(p=0.3),
+        same_on_batch=False,
+        data_keys=["input"],
+    )
+
+
+# # Prep data, models, and dataloaders
+
+# ## Dataloader
+
+# In[8]:
+
+
+if subj==1:
+    num_train = 24958
+    num_test = 2770
+test_batch_size = num_test
+
+def my_split_by_node(urls): return urls
+    
+train_url = f"{data_path}/wds/subj0{subj}/train/" + "{0..36}.tar"
+# train_url = f"{data_path}/wds/subj0{subj}/train/" + "{0..1}.tar"
+print(train_url)
+
+train_data = wds.WebDataset(train_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=True, pin_memory=True)
+
+test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+print(test_url)
+
+test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+test_dl = torch.utils.data.DataLoader(test_data, batch_size=test_batch_size, shuffle=False, drop_last=True, pin_memory=True)
+
+
+# ### check dataloaders are working
+
+# In[9]:
+
+
+test_vox_indices = []
+test_73k_images = []
+for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+    test_vox_indices = np.append(test_vox_indices, behav[:,0,5].cpu().numpy())
+    test_73k_images = np.append(test_73k_images, behav[:,0,0].cpu().numpy())
+test_vox_indices = test_vox_indices.astype(np.int16)
+print(test_i, (test_i+1) * test_batch_size, len(test_vox_indices))
+print("---\n")
+
+train_vox_indices = []
+train_73k_images = []
+for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+    train_vox_indices = np.append(train_vox_indices, behav[:,0,5].long().cpu().numpy())
+    train_73k_images = np.append(train_73k_images, behav[:,0,0].cpu().numpy())
+train_vox_indices = train_vox_indices.astype(np.int16)
+print(train_i, (train_i+1) * batch_size, len(train_vox_indices))
+
+
+# ## Load data and images
+
+# In[10]:
+
+
+# load betas
+f = h5py.File(f'{data_path}/betas_all_subj0{subj}.hdf5', 'r')
+# f = h5py.File(f'{data_path}/betas_subj0{subj}_thresholded_wholebrain.hdf5', 'r')
+
+voxels = f['betas'][:]
+print(f"subj0{subj} betas loaded into memory")
+voxels = torch.Tensor(voxels).to("cpu").to(data_type)
+print("voxels", voxels.shape)
+num_voxels = voxels.shape[-1]
+
+# load orig images
+f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+images = f['images'][:]
+images = torch.Tensor(images).to("cpu").to(data_type)
+print("images", images.shape)
+
+
+# ## Load models
+
+# ### CLIP image embeddings  model
+
+# In[11]:
+
+
+from models import Clipper
+clip_model = Clipper("ViT-L/14", device=torch.device(f"cuda:{local_rank}"), hidden_state=True, norm_embs=True)
+clip_seq_dim = 257
+clip_emb_dim = 768 #1024
+# hidden_dim = 4096
+#seq_len = 1 #2 #32 
+
+
+# ### SD VAE
+
+# In[12]:
+
+
+# if blurry_recon:
+#     from diffusers import AutoencoderKL
+#     autoenc = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, cache_dir="/fsx/proj-fmri/shared/cache")
+#     # autoenc.load_state_dict(torch.load('../train_logs/sdxl_vae_normed/best.pth')["model_state_dict"])
+#     autoenc.eval()
+#     autoenc.requires_grad_(False)
+#     autoenc.to(device)
+#     utils.count_params(autoenc)
+
+if blurry_recon:# or depth_recon:
+    from diffusers import VQModel
+    autoenc = VQModel.from_pretrained("/fsx/proj-fmri/shared/cache/models--microsoft--vq-diffusion-ithq/snapshots/3f796fb49ee559370dc638dea1d8116af131d993/vqvae", torch_dtype=data_type)
+    autoenc.eval()
+    autoenc.requires_grad_(False)
+    autoenc.to(device)
+    utils.count_params(autoenc)
+
+
+# #### downsampled images
+
+# In[13]:
+
+
+if blurry_recon:
+    if utils.is_interactive(): display(utils.torch_to_Image(images[[30]]))
+
+    input_batch = images[[30]].to(device)
+    print(input_batch.shape)
+
+    downsampled_image = nn.functional.interpolate(input_batch, size=(8, 8), mode='bilinear', align_corners=False)
+    re_upsampled_image = nn.functional.interpolate(downsampled_image, size=(128, 128), mode='nearest')
+    re_upsampled_enc = autoenc.encode(2*re_upsampled_image-1).latents * 0.18215
+    print(re_upsampled_enc.shape)
+    
+    if utils.is_interactive(): display(utils.torch_to_Image((autoenc.decode(re_upsampled_enc/0.18215).sample / 2 + 0.5).clamp(0,1)))
+
+
+# #### MiDaS depth
+
+# In[14]:
+
+
+if depth_recon:
+    from controlnet_aux.midas import MidasDetector
+    
+    midas_depth = MidasDetector.from_pretrained(
+      "valhalla/t2iadapter-aux-models", filename="dpt_large_384.pt", model_type="dpt_large", cache_dir="/fsx/proj-fmri/shared/cache").to(device)
+    midas_depth.model.eval()
+    midas_depth.model.requires_grad_(False)
+    midas_depth.model.to(device)
+    pass
+
+
+# In[15]:
+
+
+if depth_recon:
+    if utils.is_interactive(): display(utils.torch_to_Image(images[[30]]))
+
+    input_batch = images[[30,31]].float().to(device)
+    print(input_batch.shape)
+    
+    midas_emb = midas_depth.model(input_batch).unsqueeze(1)
+    print(midas_emb.shape)
+
+    prediction = utils.resize(midas_emb, 32) #/30).clamp(0,1).half() # 30 is roughly prediction.max()
+    print(prediction.shape)
+    
+    prediction = (prediction / prediction.view(prediction.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(prediction)).half()
+    midas_emb_size = prediction.flatten(1).shape[1]
+    print("midas_emb", prediction.shape, prediction.min(), prediction.max())
+    print("midas_emb_size", midas_emb_size)
+    
+    if utils.is_interactive(): display(utils.torch_to_Image(utils.resize(prediction, 224))) 
+
+    if blurry_recon:
+        prediction = utils.resize(midas_emb, 128).half().repeat(1,3,1,1)
+        prediction = (prediction / prediction.view(prediction.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(prediction)).half()
+        prediction_enc = autoenc.encode(2*prediction-1).latents * 0.18215
+        print("vae midas_emb", prediction_enc.shape, prediction_enc.min(), prediction_enc.max())
+    
+        if utils.is_interactive(): display(utils.torch_to_Image((autoenc.decode(prediction_enc/0.18215).sample / 2 + 0.5).clamp(0,1)))
+
+
+# ### MindEye modules
+
+# In[17]:
+
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+model
+
+
+# In[18]:
+
+
+time_embedding_dim = 512
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer
+    def __init__(self, input_size, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linear = torch.nn.Linear(input_size, out_features)
+    def forward(self, x):
+        return self.linear(x)
+        
+model.ridge = RidgeRegression(voxels.shape[1] + time_embedding_dim, out_features=hidden_dim)
+utils.count_params(model.ridge)
+utils.count_params(model)
+
+b = torch.randn((2,1,voxels.shape[1]))
+time_emb_test = torch.randn((2,1,time_embedding_dim))
+print(b.shape, model.ridge(torch.cat((b,time_emb_test),dim=-1)).shape)
+
+
+# In[24]:
+
+
+num_past_voxels = 15
+
+
+
+# In[25]:
+
+
+from functools import partial
+from diffusers.models.vae import Decoder
+class BrainNetwork(nn.Module):
+    def __init__(self, out_dim=768, in_dim=15724, seq_len=2, h=4096, n_blocks=4, drop=.15, clip_size=768):
+        super().__init__()
+        self.seq_len = seq_len
+        self.h = h
+        self.clip_size = clip_size
+        
+        # Initial linear layer to match the input dimensions to hidden dimensions
+        # self.lin0 = nn.Linear(in_dim, seq_len * h)
+        
+        # Mixer Blocks
+        self.mixer_blocks1 = nn.ModuleList([
+            self.mixer_block1(h, drop) for _ in range(n_blocks)
+        ])
+        self.mixer_blocks2 = nn.ModuleList([
+            self.mixer_block2(seq_len, drop) for _ in range(n_blocks)
+        ])
+        
+        # Output linear layer
+        self.clin1 = nn.Linear(h * seq_len, out_dim, bias=True)
+
+        # low-rank matrices
+        # self.rank = 500
+        # self.U = nn.Parameter(torch.randn(self.rank, out_dim))
+        # self.V = nn.Parameter(torch.randn(h * seq_len, self.rank))
+        # self.S = nn.Parameter(torch.randn(out_dim))
+
+        self.clip_proj = nn.Sequential(
+            nn.LayerNorm(clip_size),
+            nn.GELU(),
+            nn.Linear(clip_size, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, clip_size)
+        )
+
+        if blurry_recon:
+            # self.blin1 = nn.Sequential(
+            #     nn.Linear(out_dim, 4096, bias=True),
+            #     nn.LayerNorm(4096),
+            #     nn.GELU(),
+            #     nn.Linear(4096, 4096))
+            self.blin1 = nn.Linear(h*seq_len, 4096)
+            self.bgroupnorm = nn.GroupNorm(1, 256)
+            self.bupsampler = Decoder(
+                in_channels=256,
+                out_channels=128,
+                up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D"],
+                block_out_channels=[32, 64, 128],
+                layers_per_block=1,
+            )
+
+        if depth_recon:
+            # self.dlin1 = nn.Sequential(
+            #         nn.Linear(h, midas_emb_size),
+            #         nn.Sigmoid(),
+            #     )
+            self.dlin1 = nn.Linear(h*seq_len, 4096)
+            self.dgroupnorm = nn.GroupNorm(1, 256)
+            self.dupsampler = Decoder(
+                in_channels=256,
+                out_channels=1,#128,
+                up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D"],
+                block_out_channels=[32, 64, 128, 256],
+                layers_per_block=1,
+            )
+        
+    def mixer_block1(self, h, drop):
+        return nn.Sequential(
+            nn.LayerNorm(h),
+            self.mlp(h, h, drop),  # Token mixing
+        )
+
+    def mixer_block2(self, seq_len, drop):
+        return nn.Sequential(
+            nn.LayerNorm(seq_len),
+            self.mlp(seq_len, seq_len, drop)  # Channel mixing
+        )
+    
+    def mlp(self, in_dim, out_dim, drop):
+        return nn.Sequential(
+            nn.Linear(in_dim, out_dim),
+            nn.GELU(),
+            nn.Dropout(drop),
+            nn.Linear(out_dim, out_dim),
+        )
+        
+    def forward(self, x, idx = None):
+        print(idx)
+        # make empty tensors for blur and depth outputs
+        b,d = torch.Tensor([0.]), torch.Tensor([0.])
+        
+        # Initial linear layer
+        # x = self.lin0(x)
+        
+        # Reshape to seq_len by dim
+        # x = x.reshape(-1, self.seq_len, self.h)
+        
+        # Mixer blocks
+        #print("x shape ", x.shape)
+        residual1 = x
+        residual2 = x.permute(0,2,1)
+        #print("residual 2", residual2.shape)
+        for block1, block2 in zip(self.mixer_blocks1,self.mixer_blocks2):
+            x = block1(x) + residual1
+            #print("xblo", x.shape)
+            residual1 = x
+            x = x.permute(0,2,1)
+            
+            x = block2(x) + residual2
+            #print("xblo2", x.shape)
+            residual2 = x
+            x = x.permute(0,2,1)
+        
+        # Flatten
+        x = x.reshape(x.size(0), -1)
+        
+        c = self.clin1(x)
+
+        # low rank linear to out dim cuts # params by nearly half compared to full linear mapping
+        # c = (x @ (self.V/100) @ (self.U/100)) + self.S
+        
+        c = self.clip_proj(c.reshape(len(c), -1, self.clip_size))
+
+        if blurry_recon:
+            b = self.blin1(x)
+            b = b.reshape(len(b), 256, 4, 4)
+            b = self.bgroupnorm(b)
+            b = self.bupsampler(b)
+            
+        if depth_recon:
+            d = self.dlin1(x)#.reshape(len(x), 1, 32, 32)
+            d = d.reshape(len(d), 256, 4, 4)
+            d = self.dgroupnorm(d)
+            d = self.dupsampler(d)
+        
+        return c, b, d
+
+
+class TimeEmbedding(nn.Module):
+    def __init__(self, embedding_time_dim=512, num_past_voxels=15):
+        super().__init__()
+        self.embedding_time = nn.Embedding(num_past_voxels, embedding_time_dim)
+        self.num_past_voxels = num_past_voxels
+        self.embedding_time_dim = embedding_time_dim
+
+    def forward(self, time):
+        # time is (batch_size,)
+        time = time.long()
+        time = self.embedding_time(time)
+        return time # (batch_size, embedding_time_dim)
+    
+
+#model.memory_encoder = MemoryEncoder(in_dim=voxels.shape[1], out_dim=4096, num_past_voxels=15, embedding_time_dim=512)
+model.time_embedding = TimeEmbedding(embedding_time_dim=512, num_past_voxels=15)
+
+model.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, seq_len=seq_len, clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim) 
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+# test that the model works on some fake data
+b = torch.randn((1,seq_len,hidden_dim))
+print("b.shape",b.shape)
+with torch.no_grad():
+    clip_, blur_, depth_ = model.backbone(b)
+print(clip_.shape, blur_.shape, depth_.shape)
+
+
+# In[ ]:
+
+
+no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+opt_grouped_parameters = [
+    {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
+]
+
+optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr)
+
+if lr_scheduler_type == 'linear':
+    lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+        optimizer,
+        total_iters=int(np.floor(num_epochs*(num_train/num_devices/batch_size))),
+        last_epoch=-1
+    )
+elif lr_scheduler_type == 'cycle':
+    total_steps=int(np.floor(num_epochs*(num_train/num_devices/batch_size)))
+    print("total_steps", total_steps)
+    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer, 
+        max_lr=max_lr,
+        total_steps=total_steps,
+        final_div_factor=1000,
+        last_epoch=-1, pct_start=2/num_epochs
+    )
+    
+def save_ckpt(tag):    
+    ckpt_path = outdir+f'/{tag}.pth'
+    print(f'saving {ckpt_path}',flush=True)
+    unwrapped_model = accelerator.unwrap_model(model)
+    try:
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': unwrapped_model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'lr_scheduler': lr_scheduler.state_dict(),
+            'train_losses': losses,
+            'test_losses': test_losses,
+            'lrs': lrs,
+            }, ckpt_path)
+    except:
+        print("Couldn't save... moving on to prevent crashing.")
+    del unwrapped_model
+        
+print("\nDone with model preparations!")
+utils.count_params(model)
+
+
+# # Weights and Biases
+
+# In[ ]:
+
+
+if local_rank==0 and wandb_log: # only use main process for wandb logging
+    import wandb
+    wandb_project = 'mindeyev2'
+    wandb_run = model_name
+    wandb_notes = ''
+    
+    print(f"wandb {wandb_project} run {wandb_run}")
+    wandb.login(host='https://stability.wandb.io')#, relogin=True)
+    wandb_config = {
+      "model_name": model_name,
+      "global_batch_size": global_batch_size,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "clip_scale": clip_scale,
+      "blur_scale": blur_scale,
+      "use_image_aug": use_image_aug,
+      "max_lr": max_lr,
+      "mixup_pct": mixup_pct,
+      "num_train": num_train,
+      "num_test": num_test,
+      "ckpt_interval": ckpt_interval,
+      "ckpt_saving": ckpt_saving,
+      "seed": seed,
+      "distributed": distributed,
+      "num_devices": num_devices,
+      "world_size": world_size,
+      "train_url": train_url,
+      "test_url": test_url,
+    }
+    print("wandb_config:\n",wandb_config)
+    if False: # wandb_auto_resume
+        print("wandb_id:",model_name)
+        wandb.init(
+            id = model_name,
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+            resume="allow",
+        )
+    else:
+        wandb.init(
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+        )
+else:
+    wandb_log = False
+
+
+# # Main
+
+# In[ ]:
+
+
+epoch = 0
+losses, test_losses, lrs = [], [], []
+best_test_loss = 1e9
+soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))
+
+# Optionally resume from checkpoint #
+if resume_from_ckpt:
+    print("\n---resuming from last.pth ckpt---\n")
+    try:
+        checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+    except:
+        print('last.pth failed... trying last_backup.pth')
+        checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+    epoch = checkpoint['epoch']
+    print("Epoch",epoch)
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    model.load_state_dict(checkpoint['model_state_dict'])
+    del checkpoint
+elif wandb_log:
+    if wandb.run.resumed:
+        print("\n---resuming from last.pth ckpt---\n")
+        try:
+            checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+        except:
+            print('last.pth failed... trying last_backup.pth')
+            checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+        epoch = checkpoint['epoch']
+        print("Epoch",epoch)
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+        model.load_state_dict(checkpoint['model_state_dict'])
+        del checkpoint
+torch.cuda.empty_cache()
+
+
+# In[ ]:
+
+
+model, optimizer, train_dl, lr_scheduler = accelerator.prepare(
+model, optimizer, train_dl, lr_scheduler
+)
+# leaving out test_dl since we will only have local_rank 0 device do evals
+
+
+# In[ ]:
+
+
+def add_saturation(image, alpha=2):
+    gray_image = 0.2989 * image[:, 0, :, :] + 0.5870 * image[:, 1, :, :] + 0.1140 * image[:, 2, :, :]
+    gray_image = gray_image.unsqueeze(1).expand_as(image)
+    saturated_image = alpha * image + (1 - alpha) * gray_image
+    return torch.clamp(saturated_image, 0, 1)
+
+
+# In[ ]:
+
+
+print(f"{model_name} starting with epoch {epoch} / {num_epochs}")
+progress_bar = tqdm(range(epoch,num_epochs), ncols=1200, disable=(local_rank!=0))
+test_image, test_voxel = None, None
+mse = nn.MSELoss()
+l1 = nn.L1Loss()
+
+for epoch in progress_bar:
+    model.train()
+    
+    fwd_percent_correct = 0.
+    bwd_percent_correct = 0.
+    test_fwd_percent_correct = 0.
+    test_bwd_percent_correct = 0.
+
+    loss_clip_total = 0.
+    loss_blurry_total = 0.
+    loss_depth_total = 0.
+    test_loss_clip_total = 0.
+    test_loss_blurry_total = 0.
+    test_loss_depth_total = 0.
+
+    blurry_pixcorr = 0.
+    test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1
+    
+    for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+        with torch.cuda.amp.autocast(dtype=data_type):
+            optimizer.zero_grad()
+    
+            #voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+            #image = images[behav[:,0,0].cpu().long()].to(device).float()
+            
+            #past_15_voxels = voxels[past_behav[:,:,5].cpu().long()].to(device) # batch_size, 15, 15279
+            #past_15_times = torch.Tensor([i for i in range(seq_len - 1)]).to(device) # 15
+            
+            voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+            image = images[behav[:,0,0].cpu().long()].to(device).float()
+
+            past_15_voxels = voxels[past_behav[:,:seq_len-1,5].cpu().long()].to(device) # batch_size, 15, 15279
+            past_15_times = torch.Tensor([i for i in range(seq_len-1)]).to(device) # 15
+            #for past in range(1):
+            #    past_voxel = voxels[past_behav[:,past,5].cpu().long()].to(device)
+    
+            if blurry_recon:
+                # blurry_image_enc = autoenc.encode(2*utils.resize(image,128)-1).latent_dist.mode() * 0.18215
+                blurry_image_enc = autoenc.encode(2*utils.resize(add_saturation(image),128)-1).latents * 0.18215
+
+            if depth_recon:
+                # depth_images = utils.resize(midas_depth.model(image).unsqueeze(1).repeat(1,3,1,1), 128)
+                depth_images = utils.resize(midas_depth.model(image).unsqueeze(1), 32)
+                depth_images = (depth_images / depth_images.view(depth_images.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(depth_images)).half()
+                depth_image_enc = depth_images # autoenc.encode(2*depth_images-1).latents * 0.18215
+            
+            if use_image_aug: 
+                image = img_augment(image)
+    
+            clip_target = clip_model.embed_image(image)
+            assert not torch.any(torch.isnan(clip_target))
+    
+            if epoch < int(mixup_pct * num_epochs):
+                voxel, perm, betas, select = utils.mixco(voxel)
+                past_voxel, _, _, _ = utils.mixco(voxel, perm=perm, betas=betas, select=select)
+                
+            for p in range(seq_len-1):
+                #print(past_behav.shape) #128, 15, 17
+                #print(past_behav[:,p,-1])
+                #print(past_15_voxels.shape) # 128, 1, 15724
+                mask = past_behav[:,p,-1] == torch.ones_like(past_behav[:,p,-1])
+                #print(mask) # 128
+                past_15_voxels[mask, p, :] = torch.zeros_like(past_15_voxels[0, p, :])
+                #print(past_15_voxels)
+                
+            past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+            past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+            past_15_times = past_15_times.reshape(-1)
+            time_embeddings = model.time_embedding(past_15_times)
+            past_info_full = torch.cat((past_15_voxels, time_embeddings), dim=-1)
+            
+            positional_current_voxel = torch.zeros((voxel.shape[0], time_embeddings.shape[-1])).to(voxel.device)
+            voxel = torch.cat((voxel, positional_current_voxel), dim=-1)
+            voxel_ridge = model.ridge(torch.cat((voxel, past_info_full), dim=-2))
+            voxel_ridge = voxel_ridge.view( seq_len, int(voxel_ridge.shape[0]/seq_len), hidden_dim).permute(1,0,2)
+            #unsqueeze(1) # bz * 2, 1, 4096
+            
+            # past_voxel_ridge = model.ridge(past_voxel)
+            # voxel_ridge = torch.cat((voxel_ridge.unsqueeze(1), past_voxel_ridge.unsqueeze(1)), axis=1)
+            #print(voxel_ridge.shape)
+    
+            clip_voxels, blurry_image_enc_, depth_image_enc_ = model.backbone(voxel_ridge, idx = train_i)
+            
+            clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+            clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+    
+            if epoch < int(mixup_pct * num_epochs):                
+                loss_clip = utils.mixco_nce(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006, 
+                    perm=perm, betas=betas, select=select)
+            else:
+                epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=epoch_temp)
+
+            loss_clip_total += loss_clip.item()
+            loss_clip *= clip_scale
+            loss = loss_clip
+    
+            if blurry_recon:
+                downsampled_image = nn.functional.interpolate(image, size=(8, 8), mode='bilinear', align_corners=False)
+                re_upsampled_image = add_saturation(nn.functional.interpolate(downsampled_image, size=(128, 128), mode='nearest'))
+                re_upsampled_enc = autoenc.encode(2*re_upsampled_image-1).latents * 0.18215
+                
+                loss_blurry = (l1(blurry_image_enc_, blurry_image_enc) + l1(blurry_image_enc_, re_upsampled_enc))
+                loss_blurry += l1(torch.var(blurry_image_enc), torch.var(blurry_image_enc_))
+                loss_blurry_total += loss_blurry.item()
+                loss_blurry *= blur_scale
+                loss += loss_blurry
+
+            if depth_recon:
+                loss_depth = l1(depth_image_enc_, depth_image_enc)
+                # loss_depth += l1(torch.var(depth_image_enc_), torch.var(depth_image_enc))
+                loss_depth_total += loss_depth.item()
+                loss_depth *= depth_scale
+                loss += loss_depth
+    
+            # forward and backward top 1 accuracy        
+            labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+            fwd_percent_correct += utils.topk(torch.abs(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm)), labels, k=1).item()
+            bwd_percent_correct += utils.topk(torch.abs(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm)), labels, k=1).item()
+    
+            if blurry_recon:
+                with torch.no_grad():
+                    # only doing pixcorr eval on a subset of the samples per batch because its costly & slow to compute autoenc.decode()
+                    random_samps = np.random.choice(np.arange(len(voxel)), size=batch_size//5, replace=False)
+                    # random_samps = np.arange(batch_size//5)
+                    blurry_recon_images = (autoenc.decode(blurry_image_enc_[random_samps]/0.18215).sample/ 2 + 0.5).clamp(0,1)
+                    # pixcorr_origsize_nanmean is computationally less intense than utils.pixcorr and uses nanmean instead of mean
+                    pixcorr = utils.pixcorr_origsize_nanmean(image[random_samps], blurry_recon_images)
+                    # pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)
+                    # loss += (1 - pixcorr)
+                    blurry_pixcorr += pixcorr.item()
+                    # utils.check_loss(pixcorr)
+
+            utils.check_loss(loss)
+            accelerator.backward(loss)
+            optimizer.step()
+    
+            losses.append(loss.item())
+            lrs.append(optimizer.param_groups[0]['lr'])
+    
+            if lr_scheduler_type is not None:
+                lr_scheduler.step()
+
+    model.eval()
+    if local_rank==0:
+        with torch.no_grad(), torch.cuda.amp.autocast(dtype=data_type): 
+            for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):  
+                # all test samples should be loaded per batch such that test_i should never exceed 0
+                assert len(behav) == num_test
+                
+                ## Average same-image repeats ##
+                if test_image is None:
+                    voxel = voxels[behav[:,0,5].cpu().long()]
+                    image = behav[:,0,0].cpu().long()
+                    
+                    unique_image, sort_indices = torch.unique(image, return_inverse=True)
+                    for im in unique_image:
+                        locs = torch.where(im == image)[0]
+                        if test_image is None:
+                            test_image = images[im][None]
+                            test_voxel = torch.mean(voxel[locs],axis=0)[None]
+                        else:
+                            test_image = torch.vstack((test_image, images[im][None]))
+                            test_voxel = torch.vstack((test_voxel, torch.mean(voxel[locs],axis=0)[None]))
+    
+                # random sample of 300
+                random_indices = torch.arange(len(test_voxel))[:300]
+                voxel = test_voxel[random_indices].to(device)
+                image = test_image[random_indices].to(device)
+                assert len(image) == 300
+                
+                current_past_behav = past_behav[random_indices]
+
+                past_15_voxels = voxels[current_past_behav[:,:seq_len-1,5].cpu().long()].to(device) # batch_size, 15, 15279
+                past_15_times = torch.Tensor([i for i in range(seq_len-1)]).to(device) # 15
+
+                if blurry_recon:
+                    # blurry_image_enc = autoenc.encode(2*utils.resize(image,128)-1).latent_dist.mode() * 0.18215
+                    blurry_image_enc = autoenc.encode(2*utils.resize(add_saturation(image),128)-1).latents * 0.18215
+
+                if depth_recon:
+                    # depth_images = utils.resize(midas_depth.model(image).unsqueeze(1).repeat(1,3,1,1), 128)
+                    depth_images = utils.resize(midas_depth.model(image).unsqueeze(1), 32)
+                    depth_images = (depth_images / depth_images.view(depth_images.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(depth_images)).half()
+                    depth_image_enc = depth_images # autoenc.encode(2*depth_images-1).latents * 0.18215
+            
+                clip_target = clip_model.embed_image(image.float())
+                
+
+                past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+                past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+                past_15_times = past_15_times.reshape(-1)
+                time_embeddings = model.time_embedding(past_15_times)
+                past_info_full = torch.cat((past_15_voxels, time_embeddings), dim=-1)
+
+                positional_current_voxel = torch.zeros((voxel.shape[0], time_embeddings.shape[-1])).to(voxel.device)
+                voxel = torch.cat((voxel, positional_current_voxel), dim=-1)
+                voxel_ridge = model.ridge(torch.cat((voxel, past_info_full), dim=-2))
+                voxel_ridge = voxel_ridge.view(seq_len, int(voxel_ridge.shape[0]/seq_len), hidden_dim).permute(1,0,2)
+                
+                #voxel_ridge = model.ridge(voxel).unsqueeze(1)
+
+                # voxel_ridge = torch.cat((voxel_ridge.unsqueeze(1),voxel_ridge.unsqueeze(1)),axis=1)
+                
+                clip_voxels, blurry_image_enc_, depth_image_enc_ = model.backbone(voxel_ridge)
+                
+                clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+                clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+        
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006)
+                test_loss_clip_total += loss_clip.item()
+                loss_clip = loss_clip * clip_scale
+                loss = loss_clip
+
+                if blurry_recon:
+                    downsampled_image = nn.functional.interpolate(image, size=(8, 8), mode='bilinear', align_corners=False)
+                    re_upsampled_image = add_saturation(nn.functional.interpolate(downsampled_image, size=(128, 128), mode='nearest'))
+                    re_upsampled_enc = autoenc.encode(2*re_upsampled_image-1).latents * 0.18215
+                    
+                    loss_blurry = (l1(blurry_image_enc_, blurry_image_enc) + l1(blurry_image_enc_, re_upsampled_enc))
+                    loss_blurry += l1(torch.var(blurry_image_enc), torch.var(blurry_image_enc_))
+                    test_loss_blurry_total += loss_blurry.item()
+                    loss_blurry *= blur_scale
+                    loss += loss_blurry
+    
+                    # halving the batch size because the decoder is computationally heavy
+                    blurry_recon_images = (autoenc.decode(blurry_image_enc_[:len(voxel)//2]/0.18215).sample / 2 + 0.5).clamp(0,1)
+                    blurry_recon_images = torch.vstack((blurry_recon_images, (autoenc.decode(blurry_image_enc_[len(voxel)//2:]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                    pixcorr = utils.pixcorr(image, blurry_recon_images)
+                    loss += (1 - pixcorr)
+                    test_blurry_pixcorr += pixcorr.item()
+
+                if depth_recon:
+                    loss_depth = l1(depth_image_enc_, depth_image_enc)
+                    # loss_depth += l1(torch.var(depth_image_enc_), torch.var(depth_image_enc))
+                    test_loss_depth_total += loss_depth.item()
+                    loss_depth *= depth_scale
+                    loss += loss_depth
+        
+                # forward and backward top 1 accuracy        
+                labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+                test_fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item()
+                test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()
+
+                utils.check_loss(loss)                
+                test_losses.append(loss.item())
+
+            # if utils.is_interactive(): clear_output(wait=True)
+            print("---")
+            
+            assert (test_i+1) == 1
+            logs = {"train/loss": np.mean(losses[-(train_i+1):]),
+                "test/loss": np.mean(test_losses[-(test_i+1):]),
+                "train/lr": lrs[-1],
+                "train/num_steps": len(losses),
+                "test/num_steps": len(test_losses),
+                "train/fwd_pct_correct": fwd_percent_correct / (train_i + 1),
+                "train/bwd_pct_correct": bwd_percent_correct / (train_i + 1),
+                "test/test_fwd_pct_correct": test_fwd_percent_correct / (test_i + 1),
+                "test/test_bwd_pct_correct": test_bwd_percent_correct / (test_i + 1),
+                "train/loss_clip_total": loss_clip_total / (train_i + 1),
+                "train/loss_blurry_total": loss_blurry_total / (train_i + 1),
+                "test/loss_clip_total": test_loss_clip_total / (test_i + 1),
+                "test/loss_blurry_total": test_loss_blurry_total / (test_i + 1),
+                "train/blurry_pixcorr": blurry_pixcorr / (train_i + 1),
+                "test/blurry_pixcorr": test_blurry_pixcorr / (test_i + 1),
+                "train/loss_depth_total": loss_depth_total / (train_i + 1),
+                "test/loss_depth_total": test_loss_depth_total / (test_i + 1),
+                }
+    
+            if blurry_recon:    
+                # transform blurry recon latents to images and plot it
+                fig, axes = plt.subplots(1, 8, figsize=(10, 4))
+                jj=-1
+                for j in [0,1,2,3]:
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image((autoenc.decode(blurry_image_enc[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                    axes[jj].axis('off')
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image((autoenc.decode(blurry_image_enc_[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                    axes[jj].axis('off')
+                
+                if wandb_log:
+                    logs[f"test/recons"] = wandb.Image(fig, caption=f"epoch{epoch:03d}")
+                    plt.close()
+                else:
+                    plt.show()
+
+            if depth_recon:
+                # transform blurry recon latents to images and plot it
+                fig, axes = plt.subplots(1, 8, figsize=(10, 4))
+                # axes[0].imshow(utils.torch_to_Image((autoenc.decode(depth_image_enc[[0]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                # axes[1].imshow(utils.torch_to_Image((autoenc.decode(depth_image_enc_[[0]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                jj=-1
+                for j in [0,1,2,3]:
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image(utils.resize(depth_image_enc[[j]].view(1,1,32,32).clamp(0,1), 224)))
+                    axes[jj].axis('off')
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image(utils.resize(depth_image_enc_[[j]].view(1,1,32,32).clamp(0,1), 224)))
+                    axes[jj].axis('off')
+                if wandb_log:
+                    logs[f"test/depth_recons"] = wandb.Image(fig, caption=f"epoch{epoch:03d}")
+                    plt.close()
+                else:
+                    plt.show()
+            
+            progress_bar.set_postfix(**logs)
+    
+            # Save model checkpoint and reconstruct
+            if epoch % ckpt_interval == 0:
+                if not utils.is_interactive():
+                    save_ckpt(f'last')
+                    
+            if wandb_log: wandb.log(logs)
+
+    # wait for other GPUs to catch up if needed
+    accelerator.wait_for_everyone()
+    torch.cuda.empty_cache()
+    gc.collect()
+
+print("\n===Finished!===\n")
+if ckpt_saving:
+    save_ckpt(f'last')
+if not utils.is_interactive():
+    sys.exit(0)
+
+
+# In[ ]:
+
+
+plt.plot(losses)
+plt.show()
+plt.plot(test_losses)
+plt.show()
+
+
+# # Retrieve nearest neighbor in the training set using test set data
+
+# In[ ]:
+
+
+annots = np.load("/fsx/proj-fmri/shared/mindeyev2_dataset/COCO_73k_annots_curated.npy")
+
+
+# In[ ]:
+
+
+ii=2
+all_indices = np.unique(train_73k_images) #np.hstack((test_vox_indices[ii],train_vox_indices))
+with torch.no_grad(), torch.cuda.amp.autocast():
+    for batch in tqdm(range(0,len(all_indices),512)):
+        if batch==0:
+            clip_target = clip_model.embed_image(images[all_indices[batch:batch+512]]).cpu()
+        else:
+            target = clip_model.embed_image(images[all_indices[batch:batch+512]]).cpu()
+            clip_target = torch.vstack((clip_target,target))
+    clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+    voxel = test_voxel[[ii]].to(device)
+    image = test_image[[ii]].to(device)
+
+    print("Original Image (test set)")
+    display(utils.torch_to_Image(image))
+    
+    clip_target = clip_model.embed_image(image).cpu()
+    # clip_target_norm = torch.vstack((clip_target_norm, nn.functional.normalize(clip_target.flatten(1), dim=-1)))
+    
+    voxel_ridge = model.ridge(voxel).unsqueeze(1)
+    clip_voxels, _, _ = model.backbone(voxel_ridge)    
+    clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+    clip_voxels_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+    print("clip_voxels_norm", clip_voxels_norm.shape)
+    print("clip_target_norm", clip_target_norm.shape)
+    
+    sortt = torch.argsort(utils.batchwise_cosine_similarity(clip_voxels_norm.cpu(), 
+                                                            clip_target_norm).flatten()).flip(0)
+    picks = all_indices[sortt[:5]]
+
+    print("\nNearest neighbors in training set")
+    for ip,p in enumerate(picks):
+        display(utils.torch_to_Image(images[[p]]))
+        # print(utils.select_annotations([annots[int(p)]]))
+        if ip==0: predicted_caption = utils.select_annotations([annots[int(p)]])[0]
+
+print("\n=====\npredicted_caption:\n", predicted_caption)
+
+
+# # Feed into Stable Diffusion XL for reconstructions
+
+# In[ ]:
+
+
+from diffusers import StableDiffusionXLPipeline
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    "/fsx/proj-fmri/shared/cache/models--stabilityai--stable-diffusion-xl-base-1.0/snapshots/f898a3e026e802f68796b95e9702464bac78d76f", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
+)
+pipe.to("cuda")
+pass
+
+
+# In[ ]:
+
+
+prompt = predicted_caption
+recon = pipe(prompt=prompt).images[0]
+
+
+# In[ ]:
+
+
+print("Seen image")
+display(utils.torch_to_Image(image))
+
+print("Reconstruction")
+utils.torch_to_Image(utils.resize(transforms.ToTensor()(recon),224))
+

src/Train_MLPMixer-img.py

@@ -0,0 +1,1444 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+# # Code to convert this notebook to .py if you want to run it via command line or with Slurm
+#from subprocess import call
+#command = "jupyter nbconvert Train_MLPMixer-Copy1.ipynb --to python"
+#call(command,shell=True)
+
+
+# # Import packages & functions
+
+# In[2]:
+
+
+import os
+import sys
+import json
+import argparse
+import numpy as np
+import math
+from einops import rearrange
+import time
+import random
+import string
+import h5py
+from tqdm import tqdm
+
+import webdataset as wds
+import gc
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+
+from accelerate import Accelerator, DeepSpeedPlugin
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+
+global_batch_size = 16 #128
+
+import os
+os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+
+
+# In[3]:
+
+
+### Multi-GPU config ###
+local_rank = os.getenv('RANK')
+if local_rank is None: 
+    local_rank = 0
+else:
+    local_rank = int(local_rank)
+print("LOCAL RANK ", local_rank)  
+
+num_devices = torch.cuda.device_count()
+if num_devices==0: num_devices = 1
+
+#accelerator = Accelerator(split_batches=False)
+
+### UNCOMMENT BELOW STUFF TO USE DEEPSPEED (also comment out the immediately above "accelerator = " line) ###
+
+if num_devices <= 1 and utils.is_interactive():
+    # can emulate a distributed environment for deepspeed to work in jupyter notebook
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = str(np.random.randint(10000)+9000)
+    os.environ["RANK"] = "0"
+    os.environ["LOCAL_RANK"] = "0"
+    os.environ["WORLD_SIZE"] = "1"
+    os.environ["GLOBAL_BATCH_SIZE"] = str(global_batch_size) # set this to your batch size!
+    global_batch_size = os.environ["GLOBAL_BATCH_SIZE"]
+
+# alter the deepspeed config according to your global and local batch size
+if local_rank == 0:
+    with open('/fsx/proj-fmri/ckadirt/MindEyeV2/src/deepspeed_config_stage2_cpuoffload.json', 'r') as file:
+        config = json.load(file)
+    config['train_batch_size'] = int(os.environ["GLOBAL_BATCH_SIZE"])
+    config['train_micro_batch_size_per_gpu'] = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices
+    with open('deepspeed_config_stage2.json', 'w') as file:
+        json.dump(config, file)
+else:
+    # give some time for the local_rank=0 gpu to prep new deepspeed config file
+    time.sleep(10)
+deepspeed_plugin = DeepSpeedPlugin("/fsx/proj-fmri/ckadirt/MindEyeV2/src/deepspeed_config_stage2_cpuoffload.json")
+accelerator = Accelerator(split_batches=False, deepspeed_plugin=deepspeed_plugin)
+
+
+# In[4]:
+
+
+print("PID of this process =",os.getpid())
+device = accelerator.device
+print("device:",device)
+num_workers = num_devices
+print(accelerator.state)
+world_size = accelerator.state.num_processes
+distributed = not accelerator.state.distributed_type == 'NO'
+
+# set data_type to match your mixed precision (automatically set based on deepspeed config)
+if accelerator.mixed_precision == "bf16":
+    data_type = torch.bfloat16
+elif accelerator.mixed_precision == "fp16":
+    data_type = torch.float16
+else:
+    data_type = torch.float32
+
+print("distributed =",distributed, "num_devices =", num_devices, "local rank =", local_rank, "world size =", world_size, "data_type =", data_type)
+print = accelerator.print # only print if local_rank=0
+
+
+# In[5]:
+
+
+accelerator.state.distributed_type
+
+
+# # Configurations
+
+# In[6]:
+
+
+# if running this interactively, can specify jupyter_args here for argparser to use
+if utils.is_interactive():
+    # create random model_name
+    model_name = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
+    model_name = model_name + "_interactive"
+    print("model_name:", model_name)
+
+    # global_batch_size and batch_size should already be defined in the above cells
+    # other variables can be specified in the following string:
+    jupyter_args = f"--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset \
+                    --model_name={model_name} \
+                    --subj=1 --batch_size={global_batch_size} --no-blurry_recon --no-depth_recon --hidden_dim=1024 \
+                    --clip_scale=1. --blur_scale=100. --depth_scale=100. \
+                    --max_lr=3e-4 --mixup_pct=.66 --num_epochs=12 --ckpt_interval=999 --no-use_image_aug --no-ckpt_saving"
+
+    jupyter_args = jupyter_args.split()
+    print(jupyter_args)
+    
+    from IPython.display import clear_output # function to clear print outputs in cell
+    get_ipython().run_line_magic('load_ext', 'autoreload')
+    # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions
+    get_ipython().run_line_magic('autoreload', '2')
+
+
+# In[7]:
+
+
+parser = argparse.ArgumentParser(description="Model Training Configuration")
+parser.add_argument(
+    "--model_name", type=str, default="testing",
+    help="name of model, used for ckpt saving and wandb logging (if enabled)",
+)
+parser.add_argument(
+    "--data_path", type=str, default="/fsx/proj-fmri/shared/natural-scenes-dataset",
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,5,7],
+)
+parser.add_argument(
+    "--batch_size", type=int, default=32,
+    help="Batch size can be increased by 10x if only training v2c and not diffusion diffuser",
+)
+parser.add_argument(
+    "--wandb_log",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to log to wandb",
+)
+parser.add_argument(
+    "--resume_from_ckpt",action=argparse.BooleanOptionalAction,default=False,
+    help="if not using wandb and want to resume from a ckpt",
+)
+parser.add_argument(
+    "--wandb_project",type=str,default="stability",
+    help="wandb project name",
+)
+parser.add_argument(
+    "--mixup_pct",type=float,default=.33,
+    help="proportion of way through training when to switch from BiMixCo to SoftCLIP",
+)
+parser.add_argument(
+    "--blurry_recon",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to output blurry reconstructions",
+)
+parser.add_argument(
+    "--depth_recon",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to output depth reconstructions",
+)
+parser.add_argument(
+    "--blur_scale",type=float,default=100.,
+    help="multiply loss from blurry recons by this number",
+)
+parser.add_argument(
+    "--depth_scale",type=float,default=100.,
+    help="multiply loss from depth recons by this number",
+)
+parser.add_argument(
+    "--clip_scale",type=float,default=1.,
+    help="multiply contrastive loss by this number",
+)
+parser.add_argument(
+    "--use_image_aug",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to use image augmentation",
+)
+parser.add_argument(
+    "--num_epochs",type=int,default=120,
+    help="number of epochs of training",
+)
+parser.add_argument(
+    "--hidden_dim",type=int,default=4096,
+)
+parser.add_argument(
+    "--lr_scheduler_type",type=str,default='cycle',choices=['cycle','linear'],
+)
+parser.add_argument(
+    "--ckpt_saving",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--ckpt_interval",type=int,default=5,
+    help="save backup ckpt and reconstruct every x epochs",
+)
+parser.add_argument(
+    "--seed",type=int,default=42,
+)
+parser.add_argument(
+    "--max_lr",type=float,default=3e-4,
+)
+parser.add_argument(
+    "--seq_len",type=int,default=2,
+)
+
+if utils.is_interactive():
+    args = parser.parse_args(jupyter_args)
+else:
+    args = parser.parse_args()
+
+# create global variables without the args prefix
+for attribute_name in vars(args).keys():
+    globals()[attribute_name] = getattr(args, attribute_name)
+
+
+# In[8]:
+
+
+outdir = os.path.abspath(f'../train_logs/{model_name}')
+if not os.path.exists(outdir) and ckpt_saving:
+    os.makedirs(outdir,exist_ok=True)
+if use_image_aug:
+    import kornia
+    from kornia.augmentation.container import AugmentationSequential
+    img_augment = AugmentationSequential(
+        kornia.augmentation.RandomResizedCrop((224,224), (0.6,1), p=0.3),
+        kornia.augmentation.Resize((224, 224)),
+        kornia.augmentation.RandomHorizontalFlip(p=0.3),
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.3),
+        kornia.augmentation.RandomGrayscale(p=0.3),
+        same_on_batch=False,
+        data_keys=["input"],
+    )
+
+
+# # Prep data, models, and dataloaders
+
+# ## Dataloader
+
+# In[9]:
+
+
+if subj==1:
+    num_train = 24958
+    num_test = 2770
+test_batch_size = num_test
+
+def my_split_by_node(urls): return urls
+    
+train_url = f"{data_path}/wds/subj0{subj}/train/" + "{0..36}.tar"
+# train_url = f"{data_path}/wds/subj0{subj}/train/" + "{0..1}.tar"
+print(train_url)
+
+train_data = wds.WebDataset(train_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=True, pin_memory=True)
+
+test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+print(test_url)
+
+test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+test_dl = torch.utils.data.DataLoader(test_data, batch_size=test_batch_size, shuffle=False, drop_last=True, pin_memory=True)
+
+
+# ### check dataloaders are working
+
+# In[10]:
+
+
+test_vox_indices = []
+test_73k_images = []
+for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+    test_vox_indices = np.append(test_vox_indices, behav[:,0,5].cpu().numpy())
+    test_73k_images = np.append(test_73k_images, behav[:,0,0].cpu().numpy())
+test_vox_indices = test_vox_indices.astype(np.int16)
+print(test_i, (test_i+1) * test_batch_size, len(test_vox_indices))
+print("---\n")
+
+train_vox_indices = []
+train_73k_images = []
+for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+    train_vox_indices = np.append(train_vox_indices, behav[:,0,5].long().cpu().numpy())
+    train_73k_images = np.append(train_73k_images, behav[:,0,0].cpu().numpy())
+train_vox_indices = train_vox_indices.astype(np.int16)
+print(train_i, (train_i+1) * batch_size, len(train_vox_indices))
+
+
+# ## Load data and images
+
+# In[11]:
+
+
+# load betas
+f = h5py.File(f'{data_path}/betas_all_subj0{subj}.hdf5', 'r')
+# f = h5py.File(f'{data_path}/betas_subj0{subj}_thresholded_wholebrain.hdf5', 'r')
+
+voxels = f['betas'][:]
+print(f"subj0{subj} betas loaded into memory")
+voxels = torch.Tensor(voxels).to("cpu").to(data_type)
+print("voxels", voxels.shape)
+num_voxels = voxels.shape[-1]
+
+# load orig images
+f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+images = f['images'][:]
+images = torch.Tensor(images).to("cpu").to(data_type)
+print("images", images.shape)
+
+
+# ## Load models
+
+# ### CLIP image embeddings  model
+
+# In[12]:
+
+
+from models import Clipper
+clip_model = Clipper("ViT-L/14", device=torch.device(f"cuda:{local_rank}"), hidden_state=True, norm_embs=True)
+clip_seq_dim = 257
+clip_emb_dim = 768 #1024
+# hidden_dim = 4096
+#seq_len = 1 #2 #32 
+
+
+# In[13]:
+
+
+clip_model2 = Clipper("ViT-L/14", device=torch.device(f"cuda:{local_rank}"), hidden_state=False, norm_embs=True)
+
+
+# In[14]:
+
+
+#out2t = clip_model2.embed_image(torch.randn(32,3,224,224))
+
+
+# In[15]:
+
+
+#out2t.shape
+
+
+# ### SD VAE
+
+# In[16]:
+
+
+# if blurry_recon:
+#     from diffusers import AutoencoderKL
+#     autoenc = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, cache_dir="/fsx/proj-fmri/shared/cache")
+#     # autoenc.load_state_dict(torch.load('../train_logs/sdxl_vae_normed/best.pth')["model_state_dict"])
+#     autoenc.eval()
+#     autoenc.requires_grad_(False)
+#     autoenc.to(device)
+#     utils.count_params(autoenc)
+
+if blurry_recon:# or depth_recon:
+    from diffusers import VQModel
+    autoenc = VQModel.from_pretrained("/fsx/proj-fmri/shared/cache/models--microsoft--vq-diffusion-ithq/snapshots/3f796fb49ee559370dc638dea1d8116af131d993/vqvae", torch_dtype=data_type)
+    autoenc.eval()
+    autoenc.requires_grad_(False)
+    autoenc.to(device)
+    utils.count_params(autoenc)
+
+
+# #### downsampled images
+
+# In[17]:
+
+
+if blurry_recon:
+    if utils.is_interactive(): display(utils.torch_to_Image(images[[30]]))
+
+    input_batch = images[[30]].to(device)
+    print(input_batch.shape)
+
+    downsampled_image = nn.functional.interpolate(input_batch, size=(8, 8), mode='bilinear', align_corners=False)
+    re_upsampled_image = nn.functional.interpolate(downsampled_image, size=(128, 128), mode='nearest')
+    re_upsampled_enc = autoenc.encode(2*re_upsampled_image-1).latents * 0.18215
+    print(re_upsampled_enc.shape)
+    
+    if utils.is_interactive(): display(utils.torch_to_Image((autoenc.decode(re_upsampled_enc/0.18215).sample / 2 + 0.5).clamp(0,1)))
+
+
+# #### MiDaS depth
+
+# In[18]:
+
+
+if depth_recon:
+    from controlnet_aux.midas import MidasDetector
+    
+    midas_depth = MidasDetector.from_pretrained(
+      "valhalla/t2iadapter-aux-models", filename="dpt_large_384.pt", model_type="dpt_large", cache_dir="/fsx/proj-fmri/shared/cache").to(device)
+    midas_depth.model.eval()
+    midas_depth.model.requires_grad_(False)
+    midas_depth.model.to(device)
+    pass
+
+
+# In[19]:
+
+
+if depth_recon:
+    if utils.is_interactive(): display(utils.torch_to_Image(images[[30]]))
+
+    input_batch = images[[30,31]].float().to(device)
+    print(input_batch.shape)
+    
+    midas_emb = midas_depth.model(input_batch).unsqueeze(1)
+    print(midas_emb.shape)
+
+    prediction = utils.resize(midas_emb, 32) #/30).clamp(0,1).half() # 30 is roughly prediction.max()
+    print(prediction.shape)
+    
+    prediction = (prediction / prediction.view(prediction.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(prediction)).half()
+    midas_emb_size = prediction.flatten(1).shape[1]
+    print("midas_emb", prediction.shape, prediction.min(), prediction.max())
+    print("midas_emb_size", midas_emb_size)
+    
+    if utils.is_interactive(): display(utils.torch_to_Image(utils.resize(prediction, 224))) 
+
+    if blurry_recon:
+        prediction = utils.resize(midas_emb, 128).half().repeat(1,3,1,1)
+        prediction = (prediction / prediction.view(prediction.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(prediction)).half()
+        prediction_enc = autoenc.encode(2*prediction-1).latents * 0.18215
+        print("vae midas_emb", prediction_enc.shape, prediction_enc.min(), prediction_enc.max())
+    
+        if utils.is_interactive(): display(utils.torch_to_Image((autoenc.decode(prediction_enc/0.18215).sample / 2 + 0.5).clamp(0,1)))
+
+
+# ### MindEye modules
+
+# In[20]:
+
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+model
+
+
+# In[21]:
+
+
+time_embedding_dim = 512
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer
+    def __init__(self, input_size, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linear = torch.nn.Linear(input_size, out_features)
+    def forward(self, x):
+        return self.linear(x)
+        
+model.ridge = RidgeRegression(voxels.shape[1] + time_embedding_dim, out_features=hidden_dim)
+utils.count_params(model.ridge)
+utils.count_params(model)
+
+b = torch.randn((2,1,voxels.shape[1]))
+time_emb_test = torch.randn((2,1,time_embedding_dim))
+print(b.shape, model.ridge(torch.cat((b,time_emb_test),dim=-1)).shape)
+
+
+# In[22]:
+
+
+num_past_voxels = 15
+#seq_len = 1 + 1
+
+
+# In[23]:
+
+
+from functools import partial
+from diffusers.models.vae import Decoder
+class BrainNetwork(nn.Module):
+    def __init__(self, out_dim=768, in_dim=15724, seq_len=2, h=4096, n_blocks=4, drop=.15, clip_size=768):
+        super().__init__()
+        self.seq_len = seq_len
+        self.h = h
+        self.clip_size = clip_size
+        
+        # Initial linear layer to match the input dimensions to hidden dimensions
+        # self.lin0 = nn.Linear(in_dim, seq_len * h)
+        
+        # Mixer Blocks
+        self.mixer_blocks1 = nn.ModuleList([
+            self.mixer_block1(h, drop) for _ in range(n_blocks)
+        ])
+        self.mixer_blocks2 = nn.ModuleList([
+            self.mixer_block2(seq_len, drop) for _ in range(n_blocks)
+        ])
+        
+        # Output linear layer
+        self.clin1 = nn.Linear(h * seq_len, out_dim, bias=True)
+
+        # low-rank matrices
+        # self.rank = 500
+        # self.U = nn.Parameter(torch.randn(self.rank, out_dim))
+        # self.V = nn.Parameter(torch.randn(h * seq_len, self.rank))
+        # self.S = nn.Parameter(torch.randn(out_dim))
+
+        self.clip_proj = nn.Sequential(
+            nn.LayerNorm(clip_size),
+            nn.GELU(),
+            nn.Linear(clip_size, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, clip_size)
+        )
+
+        if blurry_recon:
+            # self.blin1 = nn.Sequential(
+            #     nn.Linear(out_dim, 4096, bias=True),
+            #     nn.LayerNorm(4096),
+            #     nn.GELU(),
+            #     nn.Linear(4096, 4096))
+            self.blin1 = nn.Linear(h*seq_len, 4096)
+            self.bgroupnorm = nn.GroupNorm(1, 256)
+            self.bupsampler = Decoder(
+                in_channels=256,
+                out_channels=128,
+                up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D"],
+                block_out_channels=[32, 64, 128],
+                layers_per_block=1,
+            )
+
+        if depth_recon:
+            # self.dlin1 = nn.Sequential(
+            #         nn.Linear(h, midas_emb_size),
+            #         nn.Sigmoid(),
+            #     )
+            self.dlin1 = nn.Linear(h*seq_len, 4096)
+            self.dgroupnorm = nn.GroupNorm(1, 256)
+            self.dupsampler = Decoder(
+                in_channels=256,
+                out_channels=1,#128,
+                up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D"],
+                block_out_channels=[32, 64, 128, 256],
+                layers_per_block=1,
+            )
+        
+    def mixer_block1(self, h, drop):
+        return nn.Sequential(
+            nn.LayerNorm(h),
+            self.mlp(h, h, drop),  # Token mixing
+        )
+
+    def mixer_block2(self, seq_len, drop):
+        return nn.Sequential(
+            nn.LayerNorm(seq_len),
+            self.mlp(seq_len, seq_len, drop)  # Channel mixing
+        )
+    
+    def mlp(self, in_dim, out_dim, drop):
+        return nn.Sequential(
+            nn.Linear(in_dim, out_dim),
+            nn.GELU(),
+            nn.Dropout(drop),
+            nn.Linear(out_dim, out_dim),
+        )
+        
+    def forward(self, x, idx = None):
+        print(idx)
+        # make empty tensors for blur and depth outputs
+        b,d = torch.Tensor([0.]), torch.Tensor([0.])
+        
+        # Initial linear layer
+        # x = self.lin0(x)
+        
+        # Reshape to seq_len by dim
+        # x = x.reshape(-1, self.seq_len, self.h)
+        
+        # Mixer blocks
+        #print("x shape ", x.shape)
+        residual1 = x
+        residual2 = x.permute(0,2,1)
+        #print("residual 2", residual2.shape)
+        for block1, block2 in zip(self.mixer_blocks1,self.mixer_blocks2):
+            x = block1(x) + residual1
+            #print("xblo", x.shape)
+            residual1 = x
+            x = x.permute(0,2,1)
+            
+            x = block2(x) + residual2
+            #print("xblo2", x.shape)
+            residual2 = x
+            x = x.permute(0,2,1)
+        
+        # Flatten
+        x = x.reshape(x.size(0), -1)
+        
+        c = self.clin1(x)
+
+        # low rank linear to out dim cuts # params by nearly half compared to full linear mapping
+        # c = (x @ (self.V/100) @ (self.U/100)) + self.S
+        
+        c = self.clip_proj(c.reshape(len(c), -1, self.clip_size))
+
+        if blurry_recon:
+            b = self.blin1(x)
+            b = b.reshape(len(b), 256, 4, 4)
+            b = self.bgroupnorm(b)
+            b = self.bupsampler(b)
+            
+        if depth_recon:
+            d = self.dlin1(x)#.reshape(len(x), 1, 32, 32)
+            d = d.reshape(len(d), 256, 4, 4)
+            d = self.dgroupnorm(d)
+            d = self.dupsampler(d)
+        
+        return c, b, d
+
+
+class TimeEmbedding(nn.Module):
+    def __init__(self, embedding_time_dim=512, num_past_voxels=15):
+        super().__init__()
+        self.embedding_time = nn.Embedding(num_past_voxels, embedding_time_dim)
+        self.num_past_voxels = num_past_voxels
+        self.embedding_time_dim = embedding_time_dim
+
+    def forward(self, time):
+        # time is (batch_size,)
+        time = time.long()
+        time = self.embedding_time(time)
+        return time # (batch_size, embedding_time_dim)
+    
+
+#model.memory_encoder = MemoryEncoder(in_dim=voxels.shape[1], out_dim=4096, num_past_voxels=15, embedding_time_dim=512)
+model.time_embedding = TimeEmbedding(embedding_time_dim=512, num_past_voxels=15)
+
+model.backbone = BrainNetwork(h=hidden_dim + clip_emb_dim, in_dim=hidden_dim + clip_emb_dim, seq_len=seq_len, clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim) 
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+# test that the model works on some fake data
+b = torch.randn((1,seq_len,hidden_dim + clip_emb_dim))
+print("b.shape",b.shape)
+with torch.no_grad():
+    clip_, blur_, depth_ = model.backbone(b)
+print(clip_.shape, blur_.shape, depth_.shape)
+
+
+# In[24]:
+
+
+"""
+voxel_ridge = torch.randn(512,4096)
+voxel_ridge = voxel_ridge.view(int(voxel_ridge.shape[0]/seq_len), seq_len, hidden_dim)
+print("b.shape",voxel_ridge.shape)
+with torch.no_grad():
+    clip_, blur_, depth_ = model.backbone(voxel_ridge)
+print(clip_.shape, blur_.shape, depth_.shape)"""
+
+
+# In[25]:
+
+
+no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+opt_grouped_parameters = [
+    {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
+]
+
+optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr)
+
+if lr_scheduler_type == 'linear':
+    lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+        optimizer,
+        total_iters=int(np.floor(num_epochs*(num_train/num_devices/batch_size))),
+        last_epoch=-1
+    )
+elif lr_scheduler_type == 'cycle':
+    total_steps=int(np.floor(num_epochs*(num_train/num_devices/batch_size)))
+    print("total_steps", total_steps)
+    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer, 
+        max_lr=max_lr,
+        total_steps=total_steps,
+        final_div_factor=1000,
+        last_epoch=-1, pct_start=2/num_epochs
+    )
+    
+def save_ckpt(tag):    
+    ckpt_path = outdir+f'/{tag}.pth'
+    print(f'saving {ckpt_path}',flush=True)
+    unwrapped_model = accelerator.unwrap_model(model)
+    try:
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': unwrapped_model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'lr_scheduler': lr_scheduler.state_dict(),
+            'train_losses': losses,
+            'test_losses': test_losses,
+            'lrs': lrs,
+            }, ckpt_path)
+    except:
+        print("Couldn't save... moving on to prevent crashing.")
+    del unwrapped_model
+        
+print("\nDone with model preparations!")
+utils.count_params(model)
+
+
+# In[26]:
+
+
+#nn++
+
+
+# In[27]:
+
+
+"""pp = None
+for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+    #with torch.cuda.amp.autocast(dtype=data_type):
+        #optimizer.zero_grad()
+
+        voxel = voxels[behav[:,0,5].cpu().long()]#.to(device)
+        image = images[behav[:,0,0].cpu().long()].float()#.to(device).float()
+
+        past_15_voxels = voxels[past_behav[:,:seq_len-1,5].cpu().long()]#.to(device) # batch_size, 15, 15279
+        past_15_times = torch.Tensor([i for i in range(seq_len)])#.to(device) # 15
+        print(past_behav[:,:seq_len-1,0].cpu().long())
+        past_15_images = images[past_behav[:,:seq_len-1,0].cpu().long()]
+        
+        break
+    
+        print(past_15_times)
+        #for past in range(1):
+        #    past_voxel = voxels[past_behav[:,past,5].cpu().long()].to(device)
+        
+        #if blurry_recon:
+        #    blurry_image_enc = autoenc.encode(2*utils.resize(image,128)-1).latent_dist.mode() * 0.18215
+        blurry_image_enc = autoenc.encode(2*utils.resize(add_saturation(image),128)-1).latents * 0.18215
+
+        if depth_recon:
+            # depth_images = utils.resize(midas_depth.model(image).unsqueeze(1).repeat(1,3,1,1), 128)
+            depth_images = utils.resize(midas_depth.model(image).unsqueeze(1), 32)
+            depth_images = (depth_images / depth_images.view(depth_images.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(depth_images)).half()
+            depth_image_enc = depth_images # autoenc.encode(2*depth_images-1).latents * 0.18215
+
+        if use_image_aug: 
+            image = img_augment(image)
+
+        clip_target = clip_model.embed_image(image)
+        assert not torch.any(torch.isnan(clip_target))
+
+        if epoch < int(mixup_pct * num_epochs):
+            voxel, perm, betas, select = utils.mixco(voxel)
+            past_voxel, _, _, _ = utils.mixco(voxel, perm=perm, betas=betas, select=select)
+        
+        for p in range(seq_len-1):
+            print(past_behav.shape) #128, 15, 17
+            print(past_behav[:,p,-1])
+            print(past_15_voxels.shape) # 128, 1, 15724
+            mask = past_behav[:,p,-1] == torch.ones_like(past_behav[:,p,-1])
+            print(mask) # 128
+            past_15_voxels[mask, p, :] = torch.zeros_like(past_15_voxels[0, p, :])
+            print(past_15_voxels)
+            pp = past_15_voxels
+            
+        break"""
+
+
+# In[28]:
+
+
+#pp[20, 0, :]
+
+
+# # Weights and Biases
+
+# In[29]:
+
+
+if local_rank==0 and wandb_log: # only use main process for wandb logging
+    import wandb
+    wandb_project = 'mindeyev2'
+    wandb_run = model_name
+    wandb_notes = ''
+    
+    print(f"wandb {wandb_project} run {wandb_run}")
+    wandb.login(host='https://stability.wandb.io')#, relogin=True)
+    wandb_config = {
+      "model_name": model_name,
+      "global_batch_size": global_batch_size,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "clip_scale": clip_scale,
+      "blur_scale": blur_scale,
+      "use_image_aug": use_image_aug,
+      "max_lr": max_lr,
+      "mixup_pct": mixup_pct,
+      "num_train": num_train,
+      "num_test": num_test,
+      "ckpt_interval": ckpt_interval,
+      "ckpt_saving": ckpt_saving,
+      "seed": seed,
+      "distributed": distributed,
+      "num_devices": num_devices,
+      "world_size": world_size,
+      "train_url": train_url,
+      "test_url": test_url,
+    }
+    print("wandb_config:\n",wandb_config)
+    if False: # wandb_auto_resume
+        print("wandb_id:",model_name)
+        wandb.init(
+            id = model_name,
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+            resume="allow",
+        )
+    else:
+        wandb.init(
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+        )
+else:
+    wandb_log = False
+
+
+# # Main
+
+# In[30]:
+
+
+epoch = 0
+losses, test_losses, lrs = [], [], []
+best_test_loss = 1e9
+soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))
+
+# Optionally resume from checkpoint #
+if resume_from_ckpt:
+    print("\n---resuming from last.pth ckpt---\n")
+    try:
+        checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+    except:
+        print('last.pth failed... trying last_backup.pth')
+        checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+    epoch = checkpoint['epoch']
+    print("Epoch",epoch)
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    model.load_state_dict(checkpoint['model_state_dict'])
+    del checkpoint
+elif wandb_log:
+    if wandb.run.resumed:
+        print("\n---resuming from last.pth ckpt---\n")
+        try:
+            checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+        except:
+            print('last.pth failed... trying last_backup.pth')
+            checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+        epoch = checkpoint['epoch']
+        print("Epoch",epoch)
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+        model.load_state_dict(checkpoint['model_state_dict'])
+        del checkpoint
+torch.cuda.empty_cache()
+
+
+# In[31]:
+
+
+model, optimizer, train_dl, lr_scheduler = accelerator.prepare(
+model, optimizer, train_dl, lr_scheduler
+)
+# leaving out test_dl since we will only have local_rank 0 device do evals
+
+
+# In[32]:
+
+
+def add_saturation(image, alpha=2):
+    gray_image = 0.2989 * image[:, 0, :, :] + 0.5870 * image[:, 1, :, :] + 0.1140 * image[:, 2, :, :]
+    gray_image = gray_image.unsqueeze(1).expand_as(image)
+    saturated_image = alpha * image + (1 - alpha) * gray_image
+    return torch.clamp(saturated_image, 0, 1)
+
+
+# In[33]:
+
+
+#b = torch.randn(1,2)
+#b.to(device)
+
+
+# In[34]:
+
+
+#device
+
+
+# In[35]:
+
+
+#past_15_times = torch.Tensor([i for i in range(seq_len-1)]).long() # 15
+#past_15_times.to(device)
+
+
+# In[36]:
+
+
+#nn++
+
+
+# In[ ]:
+
+
+#images.shape
+
+
+# In[94]:
+
+
+print(f"{model_name} starting with epoch {epoch} / {num_epochs}")
+progress_bar = tqdm(range(epoch,num_epochs), ncols=1200, disable=(local_rank!=0))
+test_image, test_voxel = None, None
+mse = nn.MSELoss()
+l1 = nn.L1Loss()
+
+for epoch in progress_bar:
+    model.train()
+    
+    fwd_percent_correct = 0.
+    bwd_percent_correct = 0.
+    test_fwd_percent_correct = 0.
+    test_bwd_percent_correct = 0.
+
+    loss_clip_total = 0.
+    loss_blurry_total = 0.
+    loss_depth_total = 0.
+    test_loss_clip_total = 0.
+    test_loss_blurry_total = 0.
+    test_loss_depth_total = 0.
+
+    blurry_pixcorr = 0.
+    test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1
+    
+    for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+        with torch.cuda.amp.autocast():
+            optimizer.zero_grad()
+    
+            #voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+            #image = images[behav[:,0,0].cpu().long()].to(device).float()
+            
+            #past_15_voxels = voxels[past_behav[:,:,5].cpu().long()].to(device) # batch_size, 15, 15279
+            #past_15_times = torch.Tensor([i for i in range(seq_len - 1)]).to(device) # 15
+            
+            voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+            image = images[behav[:,0,0].cpu().long()].to(device).float()
+
+            past_15_voxels = voxels[past_behav[:,:seq_len-1,5].cpu().long()].to(device) # batch_size, 15, 15279
+            #print(past_behav[:,:seq_len-1,0].cpu().long(), behav[:,0,0].cpu().long(), past_behav[:,:seq_len-1,0].cpu().long()[0])
+            past_15_images = images[past_behav[:,:seq_len-1,0].cpu().long()].to(device).float()
+            past_array = [i for i in range(seq_len-1)]
+            past_15_times = torch.Tensor(past_array) # 15
+            #print(past_15_times)
+            #print(past_15_voxels.shape, past_behav[:,:seq_len-1,5].cpu().long())
+            past_15_times = past_15_times.to(device)
+            #for past in range(1):
+            #    past_voxel = voxels[past_behav[:,past,5].cpu().long()].to(device)
+    
+            if blurry_recon:
+                # blurry_image_enc = autoenc.encode(2*utils.resize(image,128)-1).latent_dist.mode() * 0.18215
+                blurry_image_enc = autoenc.encode(2*utils.resize(add_saturation(image),128)-1).latents * 0.18215
+
+            if depth_recon:
+                # depth_images = utils.resize(midas_depth.model(image).unsqueeze(1).repeat(1,3,1,1), 128)
+                depth_images = utils.resize(midas_depth.model(image).unsqueeze(1), 32)
+                depth_images = (depth_images / depth_images.view(depth_images.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(depth_images)).half()
+                depth_image_enc = depth_images # autoenc.encode(2*depth_images-1).latents * 0.18215
+            
+            if use_image_aug: 
+                image = img_augment(image)
+    
+            clip_target = clip_model.embed_image(image)
+            assert not torch.any(torch.isnan(clip_target))
+    
+            if epoch < int(mixup_pct * num_epochs):
+                voxel, perm, betas, select = utils.mixco(voxel)
+                past_voxel, _, _, _ = utils.mixco(voxel, perm=perm, betas=betas, select=select)
+                
+            #print(past_15_images.shape)
+                
+            for p in range(seq_len-1):
+                #print(past_behav.shape) #128, 15, 17
+                #print(past_behav[:,p,-1])
+                #print(past_15_voxels.shape) # 128, 1, 15724
+                mask = past_behav[:,p,-1] == torch.ones_like(past_behav[:,p,-1])
+                #print(mask) # 128
+                past_15_voxels[mask, p, :] = torch.zeros_like(past_15_voxels[0, p, :])
+                past_15_images[mask, p, :] = torch.zeros_like(past_15_images[0, p, :])
+                #print(past_15_voxels)
+                
+            past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+            past_15_images = past_15_images.reshape(-1, past_15_images.shape[-3], past_15_images.shape[-2], past_15_images.shape[-1])
+            #print(past_15_images.shape)
+            past_15_embeddings = clip_model2.embed_image(past_15_images)
+            #print(past_15_embeddings.shape, 'uteho')
+            past_15_embeddings = torch.cat([torch.zeros(batch_size, past_15_embeddings.shape[-1]).to(past_15_embeddings.device), past_15_embeddings], dim = 0)
+            #print('tuhet', past_15_embeddings.shape)
+            #print('yepe', past_15_embeddings[0,:])
+            #print('yepe', past_15_embeddings[17,:])
+            past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+            past_15_times = past_15_times.reshape(-1)
+            time_embeddings = model.time_embedding(past_15_times)
+            
+            past_info_full = torch.cat((past_15_voxels, time_embeddings), dim=-1)
+            
+            positional_current_voxel = torch.zeros((voxel.shape[0], time_embeddings.shape[-1])).to(voxel.device)
+            voxel = torch.cat((voxel, positional_current_voxel), dim=-1)
+            voxel_ridge = model.ridge(torch.cat((voxel, past_info_full), dim=-2))
+            voxel_ridge = voxel_ridge.view(seq_len,int(voxel_ridge.shape[0]/seq_len),  hidden_dim).permute(1,0,2)
+            #past_15_embeddings = torch.split(past_15_embeddings, seq_len)
+            #print(past_15_embeddings, 'ttt')
+            past_15_embeddings = past_15_embeddings.reshape(seq_len, int(past_15_embeddings.shape[0]/seq_len), clip_emb_dim).permute(1,0,2)
+            #unsqueeze(1) # bz * 2, 1, 4096
+            #print(voxel_ridge.shape, past_15_embeddings.shape)
+            #print('yepe', past_15_embeddings[10,0,:])
+            #print('yepe', past_15_embeddings[10,1,:])
+            voxel_ridge = torch.cat((voxel_ridge, past_15_embeddings), dim=-1)
+            #print(voxel_ridge[:,0,-10:-1])
+            #print(voxel_ridge[:,0,10:20])
+            #raise("uehot")
+            # past_voxel_ridge = model.ridge(past_voxel)
+            # voxel_ridge = torch.cat((voxel_ridge.unsqueeze(1), past_voxel_ridge.unsqueeze(1)), axis=1)
+            #print(voxel_ridge.shape)
+    
+            clip_voxels, blurry_image_enc_, depth_image_enc_ = model.backbone(voxel_ridge)
+            
+            clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+            clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+    
+            if epoch < int(mixup_pct * num_epochs):                
+                loss_clip = utils.mixco_nce(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006, 
+                    perm=perm, betas=betas, select=select)
+            else:
+                epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=epoch_temp)
+
+            loss_clip_total += loss_clip.item()
+            loss_clip *= clip_scale
+            loss = loss_clip
+    
+            if blurry_recon:
+                downsampled_image = nn.functional.interpolate(image, size=(8, 8), mode='bilinear', align_corners=False)
+                re_upsampled_image = add_saturation(nn.functional.interpolate(downsampled_image, size=(128, 128), mode='nearest'))
+                re_upsampled_enc = autoenc.encode(2*re_upsampled_image-1).latents * 0.18215
+                
+                loss_blurry = (l1(blurry_image_enc_, blurry_image_enc) + l1(blurry_image_enc_, re_upsampled_enc))
+                loss_blurry += l1(torch.var(blurry_image_enc), torch.var(blurry_image_enc_))
+                loss_blurry_total += loss_blurry.item()
+                loss_blurry *= blur_scale
+                loss += loss_blurry
+
+            if depth_recon:
+                loss_depth = l1(depth_image_enc_, depth_image_enc)
+                # loss_depth += l1(torch.var(depth_image_enc_), torch.var(depth_image_enc))
+                loss_depth_total += loss_depth.item()
+                loss_depth *= depth_scale
+                loss += loss_depth
+    
+            # forward and backward top 1 accuracy        
+            labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+            fwd_percent_correct += utils.topk(torch.abs(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm)), labels, k=1).item()
+            bwd_percent_correct += utils.topk(torch.abs(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm)), labels, k=1).item()
+    
+            if blurry_recon:
+                with torch.no_grad():
+                    # only doing pixcorr eval on a subset of the samples per batch because its costly & slow to compute autoenc.decode()
+                    random_samps = np.random.choice(np.arange(len(voxel)), size=batch_size//5, replace=False)
+                    # random_samps = np.arange(batch_size//5)
+                    blurry_recon_images = (autoenc.decode(blurry_image_enc_[random_samps]/0.18215).sample/ 2 + 0.5).clamp(0,1)
+                    # pixcorr_origsize_nanmean is computationally less intense than utils.pixcorr and uses nanmean instead of mean
+                    pixcorr = utils.pixcorr_origsize_nanmean(image[random_samps], blurry_recon_images)
+                    # pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)
+                    # loss += (1 - pixcorr)
+                    blurry_pixcorr += pixcorr.item()
+                    # utils.check_loss(pixcorr)
+
+            utils.check_loss(loss)
+            accelerator.backward(loss)
+            optimizer.step()
+    
+            losses.append(loss.item())
+            lrs.append(optimizer.param_groups[0]['lr'])
+    
+            if lr_scheduler_type is not None:
+                lr_scheduler.step()
+
+    model.eval()
+    if local_rank==0:
+        with torch.no_grad(), torch.cuda.amp.autocast(dtype=data_type): 
+            for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):  
+                # all test samples should be loaded per batch such that test_i should never exceed 0
+                assert len(behav) == num_test
+                
+                ## Average same-image repeats ##
+                if test_image is None:
+                    voxel = voxels[behav[:,0,5].cpu().long()]
+                    image = behav[:,0,0].cpu().long()
+                    
+                    unique_image, sort_indices = torch.unique(image, return_inverse=True)
+                    for im in unique_image:
+                        locs = torch.where(im == image)[0]
+                        if test_image is None:
+                            test_image = images[im][None]
+                            test_voxel = torch.mean(voxel[locs],axis=0)[None]
+                        else:
+                            test_image = torch.vstack((test_image, images[im][None]))
+                            test_voxel = torch.vstack((test_voxel, torch.mean(voxel[locs],axis=0)[None]))
+    
+                # random sample of 300
+                random_indices = torch.arange(len(test_voxel))[:300]
+                voxel = test_voxel[random_indices].to(device)
+                image = test_image[random_indices].to(device)
+                assert len(image) == 300
+                
+                current_past_behav = past_behav[random_indices]
+
+                past_15_voxels = voxels[current_past_behav[:,:seq_len-1,5].cpu().long()].to(device) # batch_size, 15, 15279
+                past_15_images = images[current_past_behav[:,:seq_len-1,0].cpu().long()].to(device).float()
+                past_15_times = torch.Tensor([i for i in range(seq_len-1)]).to(device) # 15
+
+                if blurry_recon:
+                    # blurry_image_enc = autoenc.encode(2*utils.resize(image,128)-1).latent_dist.mode() * 0.18215
+                    blurry_image_enc = autoenc.encode(2*utils.resize(add_saturation(image),128)-1).latents * 0.18215
+
+                if depth_recon:
+                    # depth_images = utils.resize(midas_depth.model(image).unsqueeze(1).repeat(1,3,1,1), 128)
+                    depth_images = utils.resize(midas_depth.model(image).unsqueeze(1), 32)
+                    depth_images = (depth_images / depth_images.view(depth_images.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(depth_images)).half()
+                    depth_image_enc = depth_images # autoenc.encode(2*depth_images-1).latents * 0.18215
+            
+                clip_target = clip_model.embed_image(image.float())
+                
+
+                past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+                past_15_images = past_15_images.reshape(-1, past_15_images.shape[-3], past_15_images.shape[-2], past_15_images.shape[-1])
+                #print(past_15_images.shape)
+                past_15_embeddings = clip_model2.embed_image(past_15_images)
+                #print(past_15_embeddings.shape)
+                past_15_embeddings = torch.cat([torch.zeros(image.shape[0], past_15_embeddings.shape[-1]).to(past_15_embeddings.device), past_15_embeddings], dim = 0)
+                #print(past_15_embeddings.shape)
+                past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+                past_15_times = past_15_times.reshape(-1)
+                time_embeddings = model.time_embedding(past_15_times)
+                past_info_full = torch.cat((past_15_voxels, time_embeddings), dim=-1)
+
+                positional_current_voxel = torch.zeros((voxel.shape[0], time_embeddings.shape[-1])).to(voxel.device)
+                voxel = torch.cat((voxel, positional_current_voxel), dim=-1)
+                voxel_ridge = model.ridge(torch.cat((voxel, past_info_full), dim=-2))
+                voxel_ridge = voxel_ridge.view(seq_len, int(voxel_ridge.shape[0]/seq_len), hidden_dim).permute(1,0,2)
+                past_15_embeddings = past_15_embeddings.view(seq_len, int(past_15_embeddings.shape[0]/seq_len), clip_emb_dim).permute(1,0,2)
+                #print(past_15_embeddings.shape, voxel_ridge.shape)
+                voxel_ridge = torch.cat((voxel_ridge, past_15_embeddings), dim=-1)
+                
+                #voxel_ridge = model.ridge(voxel).unsqueeze(1)
+
+                # voxel_ridge = torch.cat((voxel_ridge.unsqueeze(1),voxel_ridge.unsqueeze(1)),axis=1)
+                
+                clip_voxels, blurry_image_enc_, depth_image_enc_ = model.backbone(voxel_ridge)
+                
+                clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+                clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+        
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006)
+                test_loss_clip_total += loss_clip.item()
+                loss_clip = loss_clip * clip_scale
+                loss = loss_clip
+
+                if blurry_recon:
+                    downsampled_image = nn.functional.interpolate(image, size=(8, 8), mode='bilinear', align_corners=False)
+                    re_upsampled_image = add_saturation(nn.functional.interpolate(downsampled_image, size=(128, 128), mode='nearest'))
+                    re_upsampled_enc = autoenc.encode(2*re_upsampled_image-1).latents * 0.18215
+                    
+                    loss_blurry = (l1(blurry_image_enc_, blurry_image_enc) + l1(blurry_image_enc_, re_upsampled_enc))
+                    loss_blurry += l1(torch.var(blurry_image_enc), torch.var(blurry_image_enc_))
+                    test_loss_blurry_total += loss_blurry.item()
+                    loss_blurry *= blur_scale
+                    loss += loss_blurry
+    
+                    # halving the batch size because the decoder is computationally heavy
+                    blurry_recon_images = (autoenc.decode(blurry_image_enc_[:len(voxel)//2]/0.18215).sample / 2 + 0.5).clamp(0,1)
+                    blurry_recon_images = torch.vstack((blurry_recon_images, (autoenc.decode(blurry_image_enc_[len(voxel)//2:]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                    pixcorr = utils.pixcorr(image, blurry_recon_images)
+                    loss += (1 - pixcorr)
+                    test_blurry_pixcorr += pixcorr.item()
+
+                if depth_recon:
+                    loss_depth = l1(depth_image_enc_, depth_image_enc)
+                    # loss_depth += l1(torch.var(depth_image_enc_), torch.var(depth_image_enc))
+                    test_loss_depth_total += loss_depth.item()
+                    loss_depth *= depth_scale
+                    loss += loss_depth
+        
+                # forward and backward top 1 accuracy        
+                labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+                test_fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item()
+                test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()
+
+                utils.check_loss(loss)                
+                test_losses.append(loss.item())
+
+            # if utils.is_interactive(): clear_output(wait=True)
+            print("---")
+            
+            assert (test_i+1) == 1
+            logs = {"train/loss": np.mean(losses[-(train_i+1):]),
+                "test/loss": np.mean(test_losses[-(test_i+1):]),
+                "train/lr": lrs[-1],
+                "train/num_steps": len(losses),
+                "test/num_steps": len(test_losses),
+                "train/fwd_pct_correct": fwd_percent_correct / (train_i + 1),
+                "train/bwd_pct_correct": bwd_percent_correct / (train_i + 1),
+                "test/test_fwd_pct_correct": test_fwd_percent_correct / (test_i + 1),
+                "test/test_bwd_pct_correct": test_bwd_percent_correct / (test_i + 1),
+                "train/loss_clip_total": loss_clip_total / (train_i + 1),
+                "train/loss_blurry_total": loss_blurry_total / (train_i + 1),
+                "test/loss_clip_total": test_loss_clip_total / (test_i + 1),
+                "test/loss_blurry_total": test_loss_blurry_total / (test_i + 1),
+                "train/blurry_pixcorr": blurry_pixcorr / (train_i + 1),
+                "test/blurry_pixcorr": test_blurry_pixcorr / (test_i + 1),
+                "train/loss_depth_total": loss_depth_total / (train_i + 1),
+                "test/loss_depth_total": test_loss_depth_total / (test_i + 1),
+                }
+    
+            if blurry_recon:    
+                # transform blurry recon latents to images and plot it
+                fig, axes = plt.subplots(1, 8, figsize=(10, 4))
+                jj=-1
+                for j in [0,1,2,3]:
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image((autoenc.decode(blurry_image_enc[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                    axes[jj].axis('off')
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image((autoenc.decode(blurry_image_enc_[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                    axes[jj].axis('off')
+                
+                if wandb_log:
+                    logs[f"test/recons"] = wandb.Image(fig, caption=f"epoch{epoch:03d}")
+                    plt.close()
+                else:
+                    plt.show()
+
+            if depth_recon:
+                # transform blurry recon latents to images and plot it
+                fig, axes = plt.subplots(1, 8, figsize=(10, 4))
+                # axes[0].imshow(utils.torch_to_Image((autoenc.decode(depth_image_enc[[0]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                # axes[1].imshow(utils.torch_to_Image((autoenc.decode(depth_image_enc_[[0]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                jj=-1
+                for j in [0,1,2,3]:
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image(utils.resize(depth_image_enc[[j]].view(1,1,32,32).clamp(0,1), 224)))
+                    axes[jj].axis('off')
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image(utils.resize(depth_image_enc_[[j]].view(1,1,32,32).clamp(0,1), 224)))
+                    axes[jj].axis('off')
+                if wandb_log:
+                    logs[f"test/depth_recons"] = wandb.Image(fig, caption=f"epoch{epoch:03d}")
+                    plt.close()
+                else:
+                    plt.show()
+            
+            progress_bar.set_postfix(**logs)
+    
+            # Save model checkpoint and reconstruct
+            if epoch % ckpt_interval == 0:
+                if not utils.is_interactive():
+                    save_ckpt(f'last')
+                    
+            if wandb_log: wandb.log(logs)
+
+    # wait for other GPUs to catch up if needed
+    accelerator.wait_for_everyone()
+    torch.cuda.empty_cache()
+    gc.collect()
+
+print("\n===Finished!===\n")
+if ckpt_saving:
+    save_ckpt(f'last')
+if not utils.is_interactive():
+    sys.exit(0)
+
+
+# In[ ]:
+
+
+plt.plot(losses)
+plt.show()
+plt.plot(test_losses)
+plt.show()
+
+
+# # Retrieve nearest neighbor in the training set using test set data
+
+# In[ ]:
+
+
+annots = np.load("/fsx/proj-fmri/shared/mindeyev2_dataset/COCO_73k_annots_curated.npy")
+
+
+# In[ ]:
+
+
+ii=2
+all_indices = np.unique(train_73k_images) #np.hstack((test_vox_indices[ii],train_vox_indices))
+with torch.no_grad(), torch.cuda.amp.autocast():
+    for batch in tqdm(range(0,len(all_indices),512)):
+        if batch==0:
+            clip_target = clip_model.embed_image(images[all_indices[batch:batch+512]]).cpu()
+        else:
+            target = clip_model.embed_image(images[all_indices[batch:batch+512]]).cpu()
+            clip_target = torch.vstack((clip_target,target))
+    clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+    voxel = test_voxel[[ii]].to(device)
+    image = test_image[[ii]].to(device)
+
+    print("Original Image (test set)")
+    display(utils.torch_to_Image(image))
+    
+    clip_target = clip_model.embed_image(image).cpu()
+    # clip_target_norm = torch.vstack((clip_target_norm, nn.functional.normalize(clip_target.flatten(1), dim=-1)))
+    
+    voxel_ridge = model.ridge(voxel).unsqueeze(1)
+    clip_voxels, _, _ = model.backbone(voxel_ridge)    
+    clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+    clip_voxels_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+    print("clip_voxels_norm", clip_voxels_norm.shape)
+    print("clip_target_norm", clip_target_norm.shape)
+    
+    sortt = torch.argsort(utils.batchwise_cosine_similarity(clip_voxels_norm.cpu(), 
+                                                            clip_target_norm).flatten()).flip(0)
+    picks = all_indices[sortt[:5]]
+
+    print("\nNearest neighbors in training set")
+    for ip,p in enumerate(picks):
+        display(utils.torch_to_Image(images[[p]]))
+        # print(utils.select_annotations([annots[int(p)]]))
+        if ip==0: predicted_caption = utils.select_annotations([annots[int(p)]])[0]
+
+print("\n=====\npredicted_caption:\n", predicted_caption)
+
+
+# # Feed into Stable Diffusion XL for reconstructions
+
+# In[ ]:
+
+
+from diffusers import StableDiffusionXLPipeline
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    "/fsx/proj-fmri/shared/cache/models--stabilityai--stable-diffusion-xl-base-1.0/snapshots/f898a3e026e802f68796b95e9702464bac78d76f", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
+)
+pipe.to("cuda")
+pass
+
+
+# In[ ]:
+
+
+prompt = predicted_caption
+recon = pipe(prompt=prompt).images[0]
+
+
+# In[ ]:
+
+
+print("Seen image")
+display(utils.torch_to_Image(image))
+
+print("Reconstruction")
+utils.torch_to_Image(utils.resize(transforms.ToTensor()(recon),224))
+

src/Train_MLPMixer.py

@@ -0,0 +1,1275 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+# # Code to convert this notebook to .py if you want to run it via command line or with Slurm
+#from subprocess import call
+#command = "jupyter nbconvert Train_MLPMixer-Copy1.ipynb --to python"
+#call(command,shell=True)
+
+
+# # Import packages & functions
+
+# In[2]:
+
+
+import os
+import sys
+import json
+import argparse
+import numpy as np
+import math
+from einops import rearrange
+import time
+import random
+import string
+import h5py
+from tqdm import tqdm
+
+import webdataset as wds
+import gc
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+
+from accelerate import Accelerator, DeepSpeedPlugin
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+
+
+# In[3]:
+
+
+### Multi-GPU config ###
+local_rank = os.getenv('RANK')
+if local_rank is None: 
+    local_rank = 0
+else:
+    local_rank = int(local_rank)
+print("LOCAL RANK ", local_rank)  
+
+num_devices = torch.cuda.device_count()
+if num_devices==0: num_devices = 1
+
+# ## UNCOMMENT BELOW SECTION AND COMMENT OUT DEEPSPEED SECTION TO AVOID USING DEEPSPEED ###
+# accelerator = Accelerator(split_batches=False, mixed_precision="fp16")
+# global_batch_size = batch_size = 32
+# data_type = torch.float16 # change depending on your mixed_precision
+
+### DEEPSPEED INITIALIZATION ###
+if num_devices <= 1 and utils.is_interactive():
+    global_batch_size = batch_size = 32
+    print(f"Setting batch_size to {batch_size}")
+    # can emulate a distributed environment for deepspeed to work in jupyter notebook
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = str(np.random.randint(10000)+9000)
+    os.environ["RANK"] = "0"
+    os.environ["LOCAL_RANK"] = "0"
+    os.environ["WORLD_SIZE"] = "1"
+    os.environ["GLOBAL_BATCH_SIZE"] = str(global_batch_size) # set this to your batch size!
+else:
+    global_batch_size = os.environ["GLOBAL_BATCH_SIZE"]    
+    batch_size = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices
+
+# alter the deepspeed config according to your global and local batch size
+if local_rank == 0:
+    with open('/fsx/proj-fmri/ckadirt/MindEyeV2/src/deepspeed_config_stage2.json', 'r') as file:
+        config = json.load(file)
+    config['train_batch_size'] = int(os.environ["GLOBAL_BATCH_SIZE"])
+    config['train_micro_batch_size_per_gpu'] = batch_size
+    config['bf16'] = {'enabled': False}
+    config['fp16'] = {'enabled': True}
+    with open('/fsx/proj-fmri/ckadirt/MindEyeV2/src/deepspeed_config_stage2.json', 'w') as file:
+        json.dump(config, file)
+else:
+    # give some time for the local_rank=0 gpu to prep new deepspeed config file
+    time.sleep(10)
+deepspeed_plugin = DeepSpeedPlugin("/fsx/proj-fmri/ckadirt/MindEyeV2/src/deepspeed_config_stage2_cpuoffload.json")
+accelerator = Accelerator(split_batches=False, deepspeed_plugin=deepspeed_plugin)
+
+
+# In[4]:
+
+
+print("PID of this process =",os.getpid())
+device = accelerator.device
+print("device:",device)
+num_workers = num_devices
+print(accelerator.state)
+world_size = accelerator.state.num_processes
+distributed = not accelerator.state.distributed_type == 'NO'
+
+# set data_type to match your mixed precision (automatically set based on deepspeed config)
+if accelerator.mixed_precision == "bf16":
+    data_type = torch.bfloat16
+elif accelerator.mixed_precision == "fp16":
+    data_type = torch.float16
+else:
+    data_type = torch.float32
+
+print("distributed =",distributed, "num_devices =", num_devices, "local rank =", local_rank, "world size =", world_size, "data_type =", data_type)
+print = accelerator.print # only print if local_rank=0
+
+
+# # Configurations
+
+# In[5]:
+
+
+# if running this interactively, can specify jupyter_args here for argparser to use
+if utils.is_interactive():
+    # create random model_name
+    model_name = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
+    model_name = model_name + "_interactive"
+    print("model_name:", model_name)
+
+    # global_batch_size and batch_size should already be defined in the above cells
+    # other variables can be specified in the following string:
+    jupyter_args = f"--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset \
+                    --model_name={model_name} \
+                    --subj=1 --batch_size={batch_size} --no-blurry_recon --no-depth_recon --hidden_dim=4096 \
+                    --clip_scale=1. --blur_scale=100. --depth_scale=100. \
+                    --max_lr=3e-4 --mixup_pct=.66 --num_epochs=12 --ckpt_interval=999 --no-use_image_aug --no-ckpt_saving"
+
+    jupyter_args = jupyter_args.split()
+    print(jupyter_args)
+    
+    from IPython.display import clear_output # function to clear print outputs in cell
+    get_ipython().run_line_magic('load_ext', 'autoreload')
+    # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions
+    get_ipython().run_line_magic('autoreload', '2')
+
+
+# In[6]:
+
+
+parser = argparse.ArgumentParser(description="Model Training Configuration")
+parser.add_argument(
+    "--model_name", type=str, default="testing",
+    help="name of model, used for ckpt saving and wandb logging (if enabled)",
+)
+parser.add_argument(
+    "--data_path", type=str, default="/fsx/proj-fmri/shared/natural-scenes-dataset",
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,5,7],
+)
+parser.add_argument(
+    "--batch_size", type=int, default=32,
+    help="Batch size can be increased by 10x if only training v2c and not diffusion diffuser",
+)
+parser.add_argument(
+    "--wandb_log",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to log to wandb",
+)
+parser.add_argument(
+    "--resume_from_ckpt",action=argparse.BooleanOptionalAction,default=False,
+    help="if not using wandb and want to resume from a ckpt",
+)
+parser.add_argument(
+    "--wandb_project",type=str,default="stability",
+    help="wandb project name",
+)
+parser.add_argument(
+    "--mixup_pct",type=float,default=.33,
+    help="proportion of way through training when to switch from BiMixCo to SoftCLIP",
+)
+parser.add_argument(
+    "--blurry_recon",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to output blurry reconstructions",
+)
+parser.add_argument(
+    "--depth_recon",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to output depth reconstructions",
+)
+parser.add_argument(
+    "--blur_scale",type=float,default=100.,
+    help="multiply loss from blurry recons by this number",
+)
+parser.add_argument(
+    "--depth_scale",type=float,default=100.,
+    help="multiply loss from depth recons by this number",
+)
+parser.add_argument(
+    "--clip_scale",type=float,default=1.,
+    help="multiply contrastive loss by this number",
+)
+parser.add_argument(
+    "--use_image_aug",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to use image augmentation",
+)
+parser.add_argument(
+    "--num_epochs",type=int,default=120,
+    help="number of epochs of training",
+)
+parser.add_argument(
+    "--hidden_dim",type=int,default=4096,
+)
+parser.add_argument(
+    "--lr_scheduler_type",type=str,default='cycle',choices=['cycle','linear'],
+)
+parser.add_argument(
+    "--ckpt_saving",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--ckpt_interval",type=int,default=5,
+    help="save backup ckpt and reconstruct every x epochs",
+)
+parser.add_argument(
+    "--seed",type=int,default=42,
+)
+parser.add_argument(
+    "--max_lr",type=float,default=3e-4,
+)
+parser.add_argument(
+    "--seq_len",type=int,default=2,
+)
+
+if utils.is_interactive():
+    args = parser.parse_args(jupyter_args)
+else:
+    args = parser.parse_args()
+
+# create global variables without the args prefix
+for attribute_name in vars(args).keys():
+    globals()[attribute_name] = getattr(args, attribute_name)
+
+
+# In[7]:
+
+
+outdir = os.path.abspath(f'../train_logs/{model_name}')
+if not os.path.exists(outdir) and ckpt_saving:
+    os.makedirs(outdir,exist_ok=True)
+if use_image_aug:
+    import kornia
+    from kornia.augmentation.container import AugmentationSequential
+    img_augment = AugmentationSequential(
+        kornia.augmentation.RandomResizedCrop((224,224), (0.6,1), p=0.3),
+        kornia.augmentation.Resize((224, 224)),
+        kornia.augmentation.RandomHorizontalFlip(p=0.3),
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.3),
+        kornia.augmentation.RandomGrayscale(p=0.3),
+        same_on_batch=False,
+        data_keys=["input"],
+    )
+
+
+# # Prep data, models, and dataloaders
+
+# ## Dataloader
+
+# In[8]:
+
+
+if subj==1:
+    num_train = 24958
+    num_test = 2770
+test_batch_size = num_test
+
+def my_split_by_node(urls): return urls
+    
+train_url = f"{data_path}/wds/subj0{subj}/train/" + "{0..36}.tar"
+# train_url = f"{data_path}/wds/subj0{subj}/train/" + "{0..1}.tar"
+print(train_url)
+
+train_data = wds.WebDataset(train_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=True, pin_memory=True)
+
+test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+print(test_url)
+
+test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+test_dl = torch.utils.data.DataLoader(test_data, batch_size=test_batch_size, shuffle=False, drop_last=True, pin_memory=True)
+
+
+# ### check dataloaders are working
+
+# In[9]:
+
+
+test_vox_indices = []
+test_73k_images = []
+for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+    test_vox_indices = np.append(test_vox_indices, behav[:,0,5].cpu().numpy())
+    test_73k_images = np.append(test_73k_images, behav[:,0,0].cpu().numpy())
+test_vox_indices = test_vox_indices.astype(np.int16)
+print(test_i, (test_i+1) * test_batch_size, len(test_vox_indices))
+print("---\n")
+
+train_vox_indices = []
+train_73k_images = []
+for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+    train_vox_indices = np.append(train_vox_indices, behav[:,0,5].long().cpu().numpy())
+    train_73k_images = np.append(train_73k_images, behav[:,0,0].cpu().numpy())
+train_vox_indices = train_vox_indices.astype(np.int16)
+print(train_i, (train_i+1) * batch_size, len(train_vox_indices))
+
+
+# ## Load data and images
+
+# In[10]:
+
+
+# load betas
+f = h5py.File(f'{data_path}/betas_all_subj0{subj}.hdf5', 'r')
+# f = h5py.File(f'{data_path}/betas_subj0{subj}_thresholded_wholebrain.hdf5', 'r')
+
+voxels = f['betas'][:]
+print(f"subj0{subj} betas loaded into memory")
+voxels = torch.Tensor(voxels).to("cpu").to(data_type)
+print("voxels", voxels.shape)
+num_voxels = voxels.shape[-1]
+
+# load orig images
+f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+images = f['images'][:]
+images = torch.Tensor(images).to("cpu").to(data_type)
+print("images", images.shape)
+
+
+# ## Load models
+
+# ### CLIP image embeddings  model
+
+# In[11]:
+
+
+from models import Clipper
+clip_model = Clipper("ViT-L/14", device=torch.device(f"cuda:{local_rank}"), hidden_state=True, norm_embs=True)
+clip_seq_dim = 257
+clip_emb_dim = 768 #1024
+# hidden_dim = 4096
+#seq_len = 1 #2 #32 
+
+
+# ### SD VAE
+
+# In[12]:
+
+
+# if blurry_recon:
+#     from diffusers import AutoencoderKL
+#     autoenc = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, cache_dir="/fsx/proj-fmri/shared/cache")
+#     # autoenc.load_state_dict(torch.load('../train_logs/sdxl_vae_normed/best.pth')["model_state_dict"])
+#     autoenc.eval()
+#     autoenc.requires_grad_(False)
+#     autoenc.to(device)
+#     utils.count_params(autoenc)
+
+if blurry_recon:# or depth_recon:
+    from diffusers import VQModel
+    autoenc = VQModel.from_pretrained("/fsx/proj-fmri/shared/cache/models--microsoft--vq-diffusion-ithq/snapshots/3f796fb49ee559370dc638dea1d8116af131d993/vqvae", torch_dtype=data_type)
+    autoenc.eval()
+    autoenc.requires_grad_(False)
+    autoenc.to(device)
+    utils.count_params(autoenc)
+
+
+# #### downsampled images
+
+# In[13]:
+
+
+if blurry_recon:
+    if utils.is_interactive(): display(utils.torch_to_Image(images[[30]]))
+
+    input_batch = images[[30]].to(device)
+    print(input_batch.shape)
+
+    downsampled_image = nn.functional.interpolate(input_batch, size=(8, 8), mode='bilinear', align_corners=False)
+    re_upsampled_image = nn.functional.interpolate(downsampled_image, size=(128, 128), mode='nearest')
+    re_upsampled_enc = autoenc.encode(2*re_upsampled_image-1).latents * 0.18215
+    print(re_upsampled_enc.shape)
+    
+    if utils.is_interactive(): display(utils.torch_to_Image((autoenc.decode(re_upsampled_enc/0.18215).sample / 2 + 0.5).clamp(0,1)))
+
+
+# #### MiDaS depth
+
+# In[14]:
+
+
+if depth_recon:
+    from controlnet_aux.midas import MidasDetector
+    
+    midas_depth = MidasDetector.from_pretrained(
+      "valhalla/t2iadapter-aux-models", filename="dpt_large_384.pt", model_type="dpt_large", cache_dir="/fsx/proj-fmri/shared/cache").to(device)
+    midas_depth.model.eval()
+    midas_depth.model.requires_grad_(False)
+    midas_depth.model.to(device)
+    pass
+
+
+# In[15]:
+
+
+if depth_recon:
+    if utils.is_interactive(): display(utils.torch_to_Image(images[[30]]))
+
+    input_batch = images[[30,31]].float().to(device)
+    print(input_batch.shape)
+    
+    midas_emb = midas_depth.model(input_batch).unsqueeze(1)
+    print(midas_emb.shape)
+
+    prediction = utils.resize(midas_emb, 32) #/30).clamp(0,1).half() # 30 is roughly prediction.max()
+    print(prediction.shape)
+    
+    prediction = (prediction / prediction.view(prediction.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(prediction)).half()
+    midas_emb_size = prediction.flatten(1).shape[1]
+    print("midas_emb", prediction.shape, prediction.min(), prediction.max())
+    print("midas_emb_size", midas_emb_size)
+    
+    if utils.is_interactive(): display(utils.torch_to_Image(utils.resize(prediction, 224))) 
+
+    if blurry_recon:
+        prediction = utils.resize(midas_emb, 128).half().repeat(1,3,1,1)
+        prediction = (prediction / prediction.view(prediction.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(prediction)).half()
+        prediction_enc = autoenc.encode(2*prediction-1).latents * 0.18215
+        print("vae midas_emb", prediction_enc.shape, prediction_enc.min(), prediction_enc.max())
+    
+        if utils.is_interactive(): display(utils.torch_to_Image((autoenc.decode(prediction_enc/0.18215).sample / 2 + 0.5).clamp(0,1)))
+
+
+# ### MindEye modules
+
+# In[17]:
+
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+model
+
+
+# In[18]:
+
+
+time_embedding_dim = 512
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer
+    def __init__(self, input_size, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linear = torch.nn.Linear(input_size, out_features)
+    def forward(self, x):
+        return self.linear(x)
+        
+model.ridge = RidgeRegression(voxels.shape[1] + time_embedding_dim, out_features=hidden_dim)
+utils.count_params(model.ridge)
+utils.count_params(model)
+
+b = torch.randn((2,1,voxels.shape[1]))
+time_emb_test = torch.randn((2,1,time_embedding_dim))
+print(b.shape, model.ridge(torch.cat((b,time_emb_test),dim=-1)).shape)
+
+
+# In[24]:
+
+
+num_past_voxels = 15
+
+
+
+# In[25]:
+
+
+from functools import partial
+from diffusers.models.vae import Decoder
+class BrainNetwork(nn.Module):
+    def __init__(self, out_dim=768, in_dim=15724, seq_len=2, h=4096, n_blocks=4, drop=.15, clip_size=768):
+        super().__init__()
+        self.seq_len = seq_len
+        self.h = h
+        self.clip_size = clip_size
+        
+        # Initial linear layer to match the input dimensions to hidden dimensions
+        # self.lin0 = nn.Linear(in_dim, seq_len * h)
+        
+        # Mixer Blocks
+        self.mixer_blocks1 = nn.ModuleList([
+            self.mixer_block1(h, drop) for _ in range(n_blocks)
+        ])
+        self.mixer_blocks2 = nn.ModuleList([
+            self.mixer_block2(seq_len, drop) for _ in range(n_blocks)
+        ])
+        
+        # Output linear layer
+        self.clin1 = nn.Linear(h * seq_len, out_dim, bias=True)
+
+        # low-rank matrices
+        # self.rank = 500
+        # self.U = nn.Parameter(torch.randn(self.rank, out_dim))
+        # self.V = nn.Parameter(torch.randn(h * seq_len, self.rank))
+        # self.S = nn.Parameter(torch.randn(out_dim))
+
+        self.clip_proj = nn.Sequential(
+            nn.LayerNorm(clip_size),
+            nn.GELU(),
+            nn.Linear(clip_size, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, clip_size)
+        )
+
+        if blurry_recon:
+            # self.blin1 = nn.Sequential(
+            #     nn.Linear(out_dim, 4096, bias=True),
+            #     nn.LayerNorm(4096),
+            #     nn.GELU(),
+            #     nn.Linear(4096, 4096))
+            self.blin1 = nn.Linear(h*seq_len, 4096)
+            self.bgroupnorm = nn.GroupNorm(1, 256)
+            self.bupsampler = Decoder(
+                in_channels=256,
+                out_channels=128,
+                up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D"],
+                block_out_channels=[32, 64, 128],
+                layers_per_block=1,
+            )
+
+        if depth_recon:
+            # self.dlin1 = nn.Sequential(
+            #         nn.Linear(h, midas_emb_size),
+            #         nn.Sigmoid(),
+            #     )
+            self.dlin1 = nn.Linear(h*seq_len, 4096)
+            self.dgroupnorm = nn.GroupNorm(1, 256)
+            self.dupsampler = Decoder(
+                in_channels=256,
+                out_channels=1,#128,
+                up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D"],
+                block_out_channels=[32, 64, 128, 256],
+                layers_per_block=1,
+            )
+        
+    def mixer_block1(self, h, drop):
+        return nn.Sequential(
+            nn.LayerNorm(h),
+            self.mlp(h, h, drop),  # Token mixing
+        )
+
+    def mixer_block2(self, seq_len, drop):
+        return nn.Sequential(
+            nn.LayerNorm(seq_len),
+            self.mlp(seq_len, seq_len, drop)  # Channel mixing
+        )
+    
+    def mlp(self, in_dim, out_dim, drop):
+        return nn.Sequential(
+            nn.Linear(in_dim, out_dim),
+            nn.GELU(),
+            nn.Dropout(drop),
+            nn.Linear(out_dim, out_dim),
+        )
+        
+    def forward(self, x, idx = None):
+        print(idx)
+        # make empty tensors for blur and depth outputs
+        b,d = torch.Tensor([0.]), torch.Tensor([0.])
+        
+        # Initial linear layer
+        # x = self.lin0(x)
+        
+        # Reshape to seq_len by dim
+        # x = x.reshape(-1, self.seq_len, self.h)
+        
+        # Mixer blocks
+        #print("x shape ", x.shape)
+        residual1 = x
+        residual2 = x.permute(0,2,1)
+        #print("residual 2", residual2.shape)
+        for block1, block2 in zip(self.mixer_blocks1,self.mixer_blocks2):
+            x = block1(x) + residual1
+            #print("xblo", x.shape)
+            residual1 = x
+            x = x.permute(0,2,1)
+            
+            x = block2(x) + residual2
+            #print("xblo2", x.shape)
+            residual2 = x
+            x = x.permute(0,2,1)
+        
+        # Flatten
+        x = x.reshape(x.size(0), -1)
+        
+        c = self.clin1(x)
+
+        # low rank linear to out dim cuts # params by nearly half compared to full linear mapping
+        # c = (x @ (self.V/100) @ (self.U/100)) + self.S
+        
+        c = self.clip_proj(c.reshape(len(c), -1, self.clip_size))
+
+        if blurry_recon:
+            b = self.blin1(x)
+            b = b.reshape(len(b), 256, 4, 4)
+            b = self.bgroupnorm(b)
+            b = self.bupsampler(b)
+            
+        if depth_recon:
+            d = self.dlin1(x)#.reshape(len(x), 1, 32, 32)
+            d = d.reshape(len(d), 256, 4, 4)
+            d = self.dgroupnorm(d)
+            d = self.dupsampler(d)
+        
+        return c, b, d
+
+
+class TimeEmbedding(nn.Module):
+    def __init__(self, embedding_time_dim=512, num_past_voxels=15):
+        super().__init__()
+        self.embedding_time = nn.Embedding(num_past_voxels, embedding_time_dim)
+        self.num_past_voxels = num_past_voxels
+        self.embedding_time_dim = embedding_time_dim
+
+    def forward(self, time):
+        # time is (batch_size,)
+        time = time.long()
+        time = self.embedding_time(time)
+        return time # (batch_size, embedding_time_dim)
+    
+
+#model.memory_encoder = MemoryEncoder(in_dim=voxels.shape[1], out_dim=4096, num_past_voxels=15, embedding_time_dim=512)
+model.time_embedding = TimeEmbedding(embedding_time_dim=512, num_past_voxels=15)
+
+model.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, seq_len=seq_len, clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim) 
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+# test that the model works on some fake data
+b = torch.randn((1,seq_len,hidden_dim))
+print("b.shape",b.shape)
+with torch.no_grad():
+    clip_, blur_, depth_ = model.backbone(b)
+print(clip_.shape, blur_.shape, depth_.shape)
+
+
+# In[ ]:
+
+
+no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+opt_grouped_parameters = [
+    {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
+]
+
+optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr)
+
+if lr_scheduler_type == 'linear':
+    lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+        optimizer,
+        total_iters=int(np.floor(num_epochs*(num_train/num_devices/batch_size))),
+        last_epoch=-1
+    )
+elif lr_scheduler_type == 'cycle':
+    total_steps=int(np.floor(num_epochs*(num_train/num_devices/batch_size)))
+    print("total_steps", total_steps)
+    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer, 
+        max_lr=max_lr,
+        total_steps=total_steps,
+        final_div_factor=1000,
+        last_epoch=-1, pct_start=2/num_epochs
+    )
+    
+def save_ckpt(tag):    
+    ckpt_path = outdir+f'/{tag}.pth'
+    print(f'saving {ckpt_path}',flush=True)
+    unwrapped_model = accelerator.unwrap_model(model)
+    try:
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': unwrapped_model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'lr_scheduler': lr_scheduler.state_dict(),
+            'train_losses': losses,
+            'test_losses': test_losses,
+            'lrs': lrs,
+            }, ckpt_path)
+    except:
+        print("Couldn't save... moving on to prevent crashing.")
+    del unwrapped_model
+        
+print("\nDone with model preparations!")
+utils.count_params(model)
+
+
+# # Weights and Biases
+
+# In[ ]:
+
+
+if local_rank==0 and wandb_log: # only use main process for wandb logging
+    import wandb
+    wandb_project = 'mindeyev2'
+    wandb_run = model_name
+    wandb_notes = ''
+    
+    print(f"wandb {wandb_project} run {wandb_run}")
+    wandb.login(host='https://stability.wandb.io')#, relogin=True)
+    wandb_config = {
+      "model_name": model_name,
+      "global_batch_size": global_batch_size,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "clip_scale": clip_scale,
+      "blur_scale": blur_scale,
+      "use_image_aug": use_image_aug,
+      "max_lr": max_lr,
+      "mixup_pct": mixup_pct,
+      "num_train": num_train,
+      "num_test": num_test,
+      "ckpt_interval": ckpt_interval,
+      "ckpt_saving": ckpt_saving,
+      "seed": seed,
+      "distributed": distributed,
+      "num_devices": num_devices,
+      "world_size": world_size,
+      "train_url": train_url,
+      "test_url": test_url,
+    }
+    print("wandb_config:\n",wandb_config)
+    if False: # wandb_auto_resume
+        print("wandb_id:",model_name)
+        wandb.init(
+            id = model_name,
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+            resume="allow",
+        )
+    else:
+        wandb.init(
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+        )
+else:
+    wandb_log = False
+
+
+# # Main
+
+# In[ ]:
+
+
+epoch = 0
+losses, test_losses, lrs = [], [], []
+best_test_loss = 1e9
+soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))
+
+# Optionally resume from checkpoint #
+if resume_from_ckpt:
+    print("\n---resuming from last.pth ckpt---\n")
+    try:
+        checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+    except:
+        print('last.pth failed... trying last_backup.pth')
+        checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+    epoch = checkpoint['epoch']
+    print("Epoch",epoch)
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    model.load_state_dict(checkpoint['model_state_dict'])
+    del checkpoint
+elif wandb_log:
+    if wandb.run.resumed:
+        print("\n---resuming from last.pth ckpt---\n")
+        try:
+            checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+        except:
+            print('last.pth failed... trying last_backup.pth')
+            checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+        epoch = checkpoint['epoch']
+        print("Epoch",epoch)
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+        model.load_state_dict(checkpoint['model_state_dict'])
+        del checkpoint
+torch.cuda.empty_cache()
+
+
+# In[ ]:
+
+
+model, optimizer, train_dl, lr_scheduler = accelerator.prepare(
+model, optimizer, train_dl, lr_scheduler
+)
+# leaving out test_dl since we will only have local_rank 0 device do evals
+
+
+# In[ ]:
+
+
+def add_saturation(image, alpha=2):
+    gray_image = 0.2989 * image[:, 0, :, :] + 0.5870 * image[:, 1, :, :] + 0.1140 * image[:, 2, :, :]
+    gray_image = gray_image.unsqueeze(1).expand_as(image)
+    saturated_image = alpha * image + (1 - alpha) * gray_image
+    return torch.clamp(saturated_image, 0, 1)
+
+
+# In[ ]:
+
+
+print(f"{model_name} starting with epoch {epoch} / {num_epochs}")
+progress_bar = tqdm(range(epoch,num_epochs), ncols=1200, disable=(local_rank!=0))
+test_image, test_voxel = None, None
+mse = nn.MSELoss()
+l1 = nn.L1Loss()
+
+for epoch in progress_bar:
+    model.train()
+    
+    fwd_percent_correct = 0.
+    bwd_percent_correct = 0.
+    test_fwd_percent_correct = 0.
+    test_bwd_percent_correct = 0.
+
+    loss_clip_total = 0.
+    loss_blurry_total = 0.
+    loss_depth_total = 0.
+    test_loss_clip_total = 0.
+    test_loss_blurry_total = 0.
+    test_loss_depth_total = 0.
+
+    blurry_pixcorr = 0.
+    test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1
+    
+    for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+        with torch.cuda.amp.autocast(dtype=data_type):
+            optimizer.zero_grad()
+    
+            #voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+            #image = images[behav[:,0,0].cpu().long()].to(device).float()
+            
+            #past_15_voxels = voxels[past_behav[:,:,5].cpu().long()].to(device) # batch_size, 15, 15279
+            #past_15_times = torch.Tensor([i for i in range(seq_len - 1)]).to(device) # 15
+            
+            voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+            image = images[behav[:,0,0].cpu().long()].to(device).float()
+
+            past_15_voxels = voxels[past_behav[:,:seq_len-1,5].cpu().long()].to(device) # batch_size, 15, 15279
+            past_15_times = torch.Tensor([i for i in range(seq_len-1)]).to(device) # 15
+            #for past in range(1):
+            #    past_voxel = voxels[past_behav[:,past,5].cpu().long()].to(device)
+    
+            if blurry_recon:
+                # blurry_image_enc = autoenc.encode(2*utils.resize(image,128)-1).latent_dist.mode() * 0.18215
+                blurry_image_enc = autoenc.encode(2*utils.resize(add_saturation(image),128)-1).latents * 0.18215
+
+            if depth_recon:
+                # depth_images = utils.resize(midas_depth.model(image).unsqueeze(1).repeat(1,3,1,1), 128)
+                depth_images = utils.resize(midas_depth.model(image).unsqueeze(1), 32)
+                depth_images = (depth_images / depth_images.view(depth_images.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(depth_images)).half()
+                depth_image_enc = depth_images # autoenc.encode(2*depth_images-1).latents * 0.18215
+            
+            if use_image_aug: 
+                image = img_augment(image)
+    
+            clip_target = clip_model.embed_image(image)
+            assert not torch.any(torch.isnan(clip_target))
+    
+            if epoch < int(mixup_pct * num_epochs):
+                voxel, perm, betas, select = utils.mixco(voxel)
+                past_voxel, _, _, _ = utils.mixco(voxel, perm=perm, betas=betas, select=select)
+                
+            for p in range(seq_len-1):
+                #print(past_behav.shape) #128, 15, 17
+                #print(past_behav[:,p,-1])
+                #print(past_15_voxels.shape) # 128, 1, 15724
+                mask = past_behav[:,p,-1] == torch.ones_like(past_behav[:,p,-1])
+                #print(mask) # 128
+                past_15_voxels[mask, p, :] = torch.zeros_like(past_15_voxels[0, p, :])
+                #print(past_15_voxels)
+                
+            past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+            past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+            past_15_times = past_15_times.reshape(-1)
+            time_embeddings = model.time_embedding(past_15_times)
+            past_info_full = torch.cat((past_15_voxels, time_embeddings), dim=-1)
+            
+            positional_current_voxel = torch.zeros((voxel.shape[0], time_embeddings.shape[-1])).to(voxel.device)
+            voxel = torch.cat((voxel, positional_current_voxel), dim=-1)
+            voxel_ridge = model.ridge(torch.cat((voxel, past_info_full), dim=-2))
+            voxel_ridge = voxel_ridge.view(int(voxel_ridge.shape[0]/seq_len), seq_len, hidden_dim)
+            #unsqueeze(1) # bz * 2, 1, 4096
+            
+            # past_voxel_ridge = model.ridge(past_voxel)
+            # voxel_ridge = torch.cat((voxel_ridge.unsqueeze(1), past_voxel_ridge.unsqueeze(1)), axis=1)
+            #print(voxel_ridge.shape)
+    
+            clip_voxels, blurry_image_enc_, depth_image_enc_ = model.backbone(voxel_ridge, idx = train_i)
+            
+            clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+            clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+    
+            if epoch < int(mixup_pct * num_epochs):                
+                loss_clip = utils.mixco_nce(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006, 
+                    perm=perm, betas=betas, select=select)
+            else:
+                epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=epoch_temp)
+
+            loss_clip_total += loss_clip.item()
+            loss_clip *= clip_scale
+            loss = loss_clip
+    
+            if blurry_recon:
+                downsampled_image = nn.functional.interpolate(image, size=(8, 8), mode='bilinear', align_corners=False)
+                re_upsampled_image = add_saturation(nn.functional.interpolate(downsampled_image, size=(128, 128), mode='nearest'))
+                re_upsampled_enc = autoenc.encode(2*re_upsampled_image-1).latents * 0.18215
+                
+                loss_blurry = (l1(blurry_image_enc_, blurry_image_enc) + l1(blurry_image_enc_, re_upsampled_enc))
+                loss_blurry += l1(torch.var(blurry_image_enc), torch.var(blurry_image_enc_))
+                loss_blurry_total += loss_blurry.item()
+                loss_blurry *= blur_scale
+                loss += loss_blurry
+
+            if depth_recon:
+                loss_depth = l1(depth_image_enc_, depth_image_enc)
+                # loss_depth += l1(torch.var(depth_image_enc_), torch.var(depth_image_enc))
+                loss_depth_total += loss_depth.item()
+                loss_depth *= depth_scale
+                loss += loss_depth
+    
+            # forward and backward top 1 accuracy        
+            labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+            fwd_percent_correct += utils.topk(torch.abs(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm)), labels, k=1).item()
+            bwd_percent_correct += utils.topk(torch.abs(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm)), labels, k=1).item()
+    
+            if blurry_recon:
+                with torch.no_grad():
+                    # only doing pixcorr eval on a subset of the samples per batch because its costly & slow to compute autoenc.decode()
+                    random_samps = np.random.choice(np.arange(len(voxel)), size=batch_size//5, replace=False)
+                    # random_samps = np.arange(batch_size//5)
+                    blurry_recon_images = (autoenc.decode(blurry_image_enc_[random_samps]/0.18215).sample/ 2 + 0.5).clamp(0,1)
+                    # pixcorr_origsize_nanmean is computationally less intense than utils.pixcorr and uses nanmean instead of mean
+                    pixcorr = utils.pixcorr_origsize_nanmean(image[random_samps], blurry_recon_images)
+                    # pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)
+                    # loss += (1 - pixcorr)
+                    blurry_pixcorr += pixcorr.item()
+                    # utils.check_loss(pixcorr)
+
+            utils.check_loss(loss)
+            accelerator.backward(loss)
+            optimizer.step()
+    
+            losses.append(loss.item())
+            lrs.append(optimizer.param_groups[0]['lr'])
+    
+            if lr_scheduler_type is not None:
+                lr_scheduler.step()
+
+    model.eval()
+    if local_rank==0:
+        with torch.no_grad(), torch.cuda.amp.autocast(dtype=data_type): 
+            for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):  
+                # all test samples should be loaded per batch such that test_i should never exceed 0
+                assert len(behav) == num_test
+                
+                ## Average same-image repeats ##
+                if test_image is None:
+                    voxel = voxels[behav[:,0,5].cpu().long()]
+                    image = behav[:,0,0].cpu().long()
+                    
+                    unique_image, sort_indices = torch.unique(image, return_inverse=True)
+                    for im in unique_image:
+                        locs = torch.where(im == image)[0]
+                        if test_image is None:
+                            test_image = images[im][None]
+                            test_voxel = torch.mean(voxel[locs],axis=0)[None]
+                        else:
+                            test_image = torch.vstack((test_image, images[im][None]))
+                            test_voxel = torch.vstack((test_voxel, torch.mean(voxel[locs],axis=0)[None]))
+    
+                # random sample of 300
+                random_indices = torch.arange(len(test_voxel))[:300]
+                voxel = test_voxel[random_indices].to(device)
+                image = test_image[random_indices].to(device)
+                assert len(image) == 300
+                
+                current_past_behav = past_behav[random_indices]
+
+                past_15_voxels = voxels[current_past_behav[:,:seq_len-1,5].cpu().long()].to(device) # batch_size, 15, 15279
+                past_15_times = torch.Tensor([i for i in range(seq_len-1)]).to(device) # 15
+
+                if blurry_recon:
+                    # blurry_image_enc = autoenc.encode(2*utils.resize(image,128)-1).latent_dist.mode() * 0.18215
+                    blurry_image_enc = autoenc.encode(2*utils.resize(add_saturation(image),128)-1).latents * 0.18215
+
+                if depth_recon:
+                    # depth_images = utils.resize(midas_depth.model(image).unsqueeze(1).repeat(1,3,1,1), 128)
+                    depth_images = utils.resize(midas_depth.model(image).unsqueeze(1), 32)
+                    depth_images = (depth_images / depth_images.view(depth_images.shape[0], -1).max(dim=1)[0].view(-1, 1, 1, 1).expand_as(depth_images)).half()
+                    depth_image_enc = depth_images # autoenc.encode(2*depth_images-1).latents * 0.18215
+            
+                clip_target = clip_model.embed_image(image.float())
+                
+
+                past_15_voxels = past_15_voxels.reshape(-1, past_15_voxels.shape[-1])
+                past_15_times = past_15_times.repeat(voxel.shape[0], 1)
+                past_15_times = past_15_times.reshape(-1)
+                time_embeddings = model.time_embedding(past_15_times)
+                past_info_full = torch.cat((past_15_voxels, time_embeddings), dim=-1)
+
+                positional_current_voxel = torch.zeros((voxel.shape[0], time_embeddings.shape[-1])).to(voxel.device)
+                voxel = torch.cat((voxel, positional_current_voxel), dim=-1)
+                voxel_ridge = model.ridge(torch.cat((voxel, past_info_full), dim=-2))
+                voxel_ridge = voxel_ridge.view(int(voxel_ridge.shape[0]/seq_len), seq_len, hidden_dim)
+                
+                #voxel_ridge = model.ridge(voxel).unsqueeze(1)
+
+                # voxel_ridge = torch.cat((voxel_ridge.unsqueeze(1),voxel_ridge.unsqueeze(1)),axis=1)
+                
+                clip_voxels, blurry_image_enc_, depth_image_enc_ = model.backbone(voxel_ridge)
+                
+                clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+                clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+        
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006)
+                test_loss_clip_total += loss_clip.item()
+                loss_clip = loss_clip * clip_scale
+                loss = loss_clip
+
+                if blurry_recon:
+                    downsampled_image = nn.functional.interpolate(image, size=(8, 8), mode='bilinear', align_corners=False)
+                    re_upsampled_image = add_saturation(nn.functional.interpolate(downsampled_image, size=(128, 128), mode='nearest'))
+                    re_upsampled_enc = autoenc.encode(2*re_upsampled_image-1).latents * 0.18215
+                    
+                    loss_blurry = (l1(blurry_image_enc_, blurry_image_enc) + l1(blurry_image_enc_, re_upsampled_enc))
+                    loss_blurry += l1(torch.var(blurry_image_enc), torch.var(blurry_image_enc_))
+                    test_loss_blurry_total += loss_blurry.item()
+                    loss_blurry *= blur_scale
+                    loss += loss_blurry
+    
+                    # halving the batch size because the decoder is computationally heavy
+                    blurry_recon_images = (autoenc.decode(blurry_image_enc_[:len(voxel)//2]/0.18215).sample / 2 + 0.5).clamp(0,1)
+                    blurry_recon_images = torch.vstack((blurry_recon_images, (autoenc.decode(blurry_image_enc_[len(voxel)//2:]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                    pixcorr = utils.pixcorr(image, blurry_recon_images)
+                    loss += (1 - pixcorr)
+                    test_blurry_pixcorr += pixcorr.item()
+
+                if depth_recon:
+                    loss_depth = l1(depth_image_enc_, depth_image_enc)
+                    # loss_depth += l1(torch.var(depth_image_enc_), torch.var(depth_image_enc))
+                    test_loss_depth_total += loss_depth.item()
+                    loss_depth *= depth_scale
+                    loss += loss_depth
+        
+                # forward and backward top 1 accuracy        
+                labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+                test_fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item()
+                test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()
+
+                utils.check_loss(loss)                
+                test_losses.append(loss.item())
+
+            # if utils.is_interactive(): clear_output(wait=True)
+            print("---")
+            
+            assert (test_i+1) == 1
+            logs = {"train/loss": np.mean(losses[-(train_i+1):]),
+                "test/loss": np.mean(test_losses[-(test_i+1):]),
+                "train/lr": lrs[-1],
+                "train/num_steps": len(losses),
+                "test/num_steps": len(test_losses),
+                "train/fwd_pct_correct": fwd_percent_correct / (train_i + 1),
+                "train/bwd_pct_correct": bwd_percent_correct / (train_i + 1),
+                "test/test_fwd_pct_correct": test_fwd_percent_correct / (test_i + 1),
+                "test/test_bwd_pct_correct": test_bwd_percent_correct / (test_i + 1),
+                "train/loss_clip_total": loss_clip_total / (train_i + 1),
+                "train/loss_blurry_total": loss_blurry_total / (train_i + 1),
+                "test/loss_clip_total": test_loss_clip_total / (test_i + 1),
+                "test/loss_blurry_total": test_loss_blurry_total / (test_i + 1),
+                "train/blurry_pixcorr": blurry_pixcorr / (train_i + 1),
+                "test/blurry_pixcorr": test_blurry_pixcorr / (test_i + 1),
+                "train/loss_depth_total": loss_depth_total / (train_i + 1),
+                "test/loss_depth_total": test_loss_depth_total / (test_i + 1),
+                }
+    
+            if blurry_recon:    
+                # transform blurry recon latents to images and plot it
+                fig, axes = plt.subplots(1, 8, figsize=(10, 4))
+                jj=-1
+                for j in [0,1,2,3]:
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image((autoenc.decode(blurry_image_enc[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                    axes[jj].axis('off')
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image((autoenc.decode(blurry_image_enc_[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                    axes[jj].axis('off')
+                
+                if wandb_log:
+                    logs[f"test/recons"] = wandb.Image(fig, caption=f"epoch{epoch:03d}")
+                    plt.close()
+                else:
+                    plt.show()
+
+            if depth_recon:
+                # transform blurry recon latents to images and plot it
+                fig, axes = plt.subplots(1, 8, figsize=(10, 4))
+                # axes[0].imshow(utils.torch_to_Image((autoenc.decode(depth_image_enc[[0]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                # axes[1].imshow(utils.torch_to_Image((autoenc.decode(depth_image_enc_[[0]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                jj=-1
+                for j in [0,1,2,3]:
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image(utils.resize(depth_image_enc[[j]].view(1,1,32,32).clamp(0,1), 224)))
+                    axes[jj].axis('off')
+                    jj+=1
+                    axes[jj].imshow(utils.torch_to_Image(utils.resize(depth_image_enc_[[j]].view(1,1,32,32).clamp(0,1), 224)))
+                    axes[jj].axis('off')
+                if wandb_log:
+                    logs[f"test/depth_recons"] = wandb.Image(fig, caption=f"epoch{epoch:03d}")
+                    plt.close()
+                else:
+                    plt.show()
+            
+            progress_bar.set_postfix(**logs)
+    
+            # Save model checkpoint and reconstruct
+            if epoch % ckpt_interval == 0:
+                if not utils.is_interactive():
+                    save_ckpt(f'last')
+                    
+            if wandb_log: wandb.log(logs)
+
+    # wait for other GPUs to catch up if needed
+    accelerator.wait_for_everyone()
+    torch.cuda.empty_cache()
+    gc.collect()
+
+print("\n===Finished!===\n")
+if ckpt_saving:
+    save_ckpt(f'last')
+if not utils.is_interactive():
+    sys.exit(0)
+
+
+# In[ ]:
+
+
+plt.plot(losses)
+plt.show()
+plt.plot(test_losses)
+plt.show()
+
+
+# # Retrieve nearest neighbor in the training set using test set data
+
+# In[ ]:
+
+
+annots = np.load("/fsx/proj-fmri/shared/mindeyev2_dataset/COCO_73k_annots_curated.npy")
+
+
+# In[ ]:
+
+
+ii=2
+all_indices = np.unique(train_73k_images) #np.hstack((test_vox_indices[ii],train_vox_indices))
+with torch.no_grad(), torch.cuda.amp.autocast():
+    for batch in tqdm(range(0,len(all_indices),512)):
+        if batch==0:
+            clip_target = clip_model.embed_image(images[all_indices[batch:batch+512]]).cpu()
+        else:
+            target = clip_model.embed_image(images[all_indices[batch:batch+512]]).cpu()
+            clip_target = torch.vstack((clip_target,target))
+    clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+    voxel = test_voxel[[ii]].to(device)
+    image = test_image[[ii]].to(device)
+
+    print("Original Image (test set)")
+    display(utils.torch_to_Image(image))
+    
+    clip_target = clip_model.embed_image(image).cpu()
+    # clip_target_norm = torch.vstack((clip_target_norm, nn.functional.normalize(clip_target.flatten(1), dim=-1)))
+    
+    voxel_ridge = model.ridge(voxel).unsqueeze(1)
+    clip_voxels, _, _ = model.backbone(voxel_ridge)    
+    clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+    clip_voxels_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+    print("clip_voxels_norm", clip_voxels_norm.shape)
+    print("clip_target_norm", clip_target_norm.shape)
+    
+    sortt = torch.argsort(utils.batchwise_cosine_similarity(clip_voxels_norm.cpu(), 
+                                                            clip_target_norm).flatten()).flip(0)
+    picks = all_indices[sortt[:5]]
+
+    print("\nNearest neighbors in training set")
+    for ip,p in enumerate(picks):
+        display(utils.torch_to_Image(images[[p]]))
+        # print(utils.select_annotations([annots[int(p)]]))
+        if ip==0: predicted_caption = utils.select_annotations([annots[int(p)]])[0]
+
+print("\n=====\npredicted_caption:\n", predicted_caption)
+
+
+# # Feed into Stable Diffusion XL for reconstructions
+
+# In[ ]:
+
+
+from diffusers import StableDiffusionXLPipeline
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    "/fsx/proj-fmri/shared/cache/models--stabilityai--stable-diffusion-xl-base-1.0/snapshots/f898a3e026e802f68796b95e9702464bac78d76f", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
+)
+pipe.to("cuda")
+pass
+
+
+# In[ ]:
+
+
+prompt = predicted_caption
+recon = pipe(prompt=prompt).images[0]
+
+
+# In[ ]:
+
+
+print("Seen image")
+display(utils.torch_to_Image(image))
+
+print("Reconstruction")
+utils.torch_to_Image(utils.resize(transforms.ToTensor()(recon),224))
+

src/accel.slurm

@@ -0,0 +1,38 @@
+#!/bin/bash
+#SBATCH --account=topfmri
+#SBATCH --partition=g40x
+#SBATCH --job-name=ms
+#SBATCH --nodes=1              
+#SBATCH --ntasks-per-node=4     # should = number of gpus
+#SBATCH --gres=gpu:4
+#SBATCH --time=32:00:00          # total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=topfmri
+
+module load cuda/11.7 # should match torch.cuda.version
+
+export NUM_GPUS=4 # Set to equal gres=gpu:#
+export GLOBAL_BATCH_SIZE=512
+
+# Make sure another job doesnt use same port, here using random number
+export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000)) 
+
+export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export COUNT_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)
+
+export WANDB_DIR="/fsx/proj-fmri/paulscotti/MindEyeV2/wandb/"
+export WANDB_CACHE_DIR="/fsx/home-paulscotti/.cache"
+export WANDB_MODE="online"
+
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo WORLD_SIZE=${COUNT_NODE}
+
+###########
+
+cd /fsx/proj-fmri/paulscotti/MindEyeV2
+accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT Train.py --data_path=/fsx/proj-fmri/shared/mindeyev2_dataset --model_name=test --subj=1 --batch_size=${GLOBAL_BATCH_SIZE} --n_samples_save=0 --max_lr=3e-4 --mixup_pct=.33 --num_epochs=240 --ckpt_interval=999 --no-use_image_aug
+
+# --wandb_log

src/accel2.slurm

@@ -0,0 +1,40 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --partition=g40x
+#SBATCH --job-name=ms
+#SBATCH --nodes=1              
+#SBATCH --ntasks-per-node=4     # should = number of gpus
+#SBATCH --gres=gpu:8
+#SBATCH --time=32:00:00          # total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=fmri
+
+
+
+export NUM_GPUS=8 # Set to equal gres=gpu:#
+export GLOBAL_BATCH_SIZE=512
+
+# Make sure another job doesnt use same port, here using random number
+export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000)) 
+
+export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export COUNT_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)
+
+export WANDB_DIR="/fsx/proj-fmri/ckadirt/MindEyeV2/src/wandb" 
+export WANDB_CACHE_DIR="/admin/home-ckadirt/.cache"
+export WANDB_MODE="online"
+
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo WORLD_SIZE=${COUNT_NODE}
+
+source /admin/home-ckadirt/.bashrc
+
+###########
+
+cd /fsx/proj-fmri/ckadirt/MindEyeV2/src/
+accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT Train-with-memory-cat.py --data_path=/fsx/proj-fmri/shared/mindeyev2_dataset --model_name=test_mem_cat_r --subj=1 --batch_size=${GLOBAL_BATCH_SIZE} --n_samples_save=0 --max_lr=3e-5 --mixup_pct=.33 --num_epochs=240 --ckpt_interval=999 --no-use_image_aug
+
+# --wandb_log

src/accel3.slurm

@@ -0,0 +1,40 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --partition=g40x
+#SBATCH --job-name=ms
+#SBATCH --nodes=1              
+#SBATCH --ntasks-per-node=4     # should = number of gpus
+#SBATCH --gres=gpu:4
+#SBATCH --time=32:00:00          # total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=fmri
+
+
+
+export NUM_GPUS=4 # Set to equal gres=gpu:#
+export GLOBAL_BATCH_SIZE=128
+
+# Make sure another job doesnt use same port, here using random number
+export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000)) 
+
+export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export COUNT_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)
+
+export WANDB_DIR="/fsx/proj-fmri/ckadirt/MindEyeV2/src/wandb" 
+export WANDB_CACHE_DIR="/admin/home-ckadirt/.cache"
+export WANDB_MODE="online"
+
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo WORLD_SIZE=${COUNT_NODE}
+
+source /admin/home-ckadirt/.bashrc
+
+###########
+
+cd /fsx/proj-fmri/ckadirt/MindEyeV2/src/
+accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT Train-with-memory.py --data_path=/fsx/proj-fmri/shared/mindeyev2_dataset --model_name=test_mem --subj=1 --batch_size=${GLOBAL_BATCH_SIZE} --n_samples_save=0 --max_lr=3e-4 --mixup_pct=.33 --num_epochs=40 --ckpt_interval=999 --no-use_image_aug
+
+# --wandb_log

src/accel4.slurm

@@ -0,0 +1,40 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --partition=g40x
+#SBATCH --job-name=ms
+#SBATCH --nodes=1              
+#SBATCH --ntasks-per-node=4     # should = number of gpus
+#SBATCH --gres=gpu:8
+#SBATCH --time=32:00:00          # total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=fmri
+
+
+
+export NUM_GPUS=8 # Set to equal gres=gpu:#
+export GLOBAL_BATCH_SIZE=512
+
+# Make sure another job doesnt use same port, here using random number
+export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000)) 
+
+export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export COUNT_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)
+
+export WANDB_DIR="/fsx/proj-fmri/ckadirt/MindEyeV2/src/wandb" 
+export WANDB_CACHE_DIR="/admin/home-ckadirt/.cache"
+export WANDB_MODE="online"
+
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo WORLD_SIZE=${COUNT_NODE}
+
+source /admin/home-ckadirt/.bashrc
+
+###########
+
+cd /fsx/proj-fmri/ckadirt/MindEyeV2/src/
+accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT Train.py --data_path=/fsx/proj-fmri/shared/mindeyev2_dataset --model_name=baseline_r --subj=1 --batch_size=${GLOBAL_BATCH_SIZE} --n_samples_save=0 --max_lr=3e-5 --mixup_pct=.33 --num_epochs=240 --ckpt_interval=999 --no-use_image_aug
+
+# --wandb_log

src/accel5.slurm

@@ -0,0 +1,40 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --partition=g40x
+#SBATCH --job-name=memoryrr
+#SBATCH --nodes=1              
+#SBATCH --ntasks-per-node=4     # should = number of gpus
+#SBATCH --gres=gpu:8
+#SBATCH --time=32:00:00          # total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=fmri
+
+
+
+export NUM_GPUS=8 # Set to equal gres=gpu:#
+export GLOBAL_BATCH_SIZE=512
+
+# Make sure another job doesnt use same port, here using random number
+export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000)) 
+
+export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export COUNT_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)
+
+export WANDB_DIR="/fsx/proj-fmri/ckadirt/MindEyeV2/src/wandb" 
+export WANDB_CACHE_DIR="/admin/home-ckadirt/.cache"
+export WANDB_MODE="online"
+
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo WORLD_SIZE=${COUNT_NODE}
+
+source /admin/home-ckadirt/.bashrc
+
+###########
+
+cd /fsx/proj-fmri/ckadirt/MindEyeV2/src/
+accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT Train-with-memory-rr.py --data_path=/fsx/proj-fmri/shared/mindeyev2_dataset --model_name=testing-rr-uni_r --subj=1 --batch_size=${GLOBAL_BATCH_SIZE} --n_samples_save=0 --max_lr=1e-5 --mixup_pct=.66 --num_epochs=120 --ckpt_interval=999 --no-use_image_aug
+
+# --wandb_log

src/accel6.slurm

@@ -0,0 +1,40 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --partition=g40x
+#SBATCH --job-name=memoryrr
+#SBATCH --nodes=1              
+#SBATCH --ntasks-per-node=4     # should = number of gpus
+#SBATCH --gres=gpu:8
+#SBATCH --time=32:00:00          # total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=fmri
+
+
+
+export NUM_GPUS=8 # Set to equal gres=gpu:#
+export GLOBAL_BATCH_SIZE=512
+
+# Make sure another job doesnt use same port, here using random number
+export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000)) 
+
+export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export COUNT_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)
+
+export WANDB_DIR="/fsx/proj-fmri/ckadirt/MindEyeV2/src/wandb" 
+export WANDB_CACHE_DIR="/admin/home-ckadirt/.cache"
+export WANDB_MODE="online"
+
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo WORLD_SIZE=${COUNT_NODE}
+
+source /admin/home-ckadirt/.bashrc
+
+###########
+
+cd /fsx/proj-fmri/ckadirt/MindEyeV2/src/
+accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT Train-with-memory-rr-dropout.py --data_path=/fsx/proj-fmri/shared/mindeyev2_dataset --model_name=testing-rr-uni_r --subj=1 --batch_size=${GLOBAL_BATCH_SIZE} --n_samples_save=0 --max_lr=3e-5 --mixup_pct=.66 --num_epochs=120 --ckpt_interval=999 --no-use_image_aug
+
+# --wandb_log

src/accel7.slurm

@@ -0,0 +1,41 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --partition=g40x
+#SBATCH --job-name=memoryrr
+#SBATCH --nodes=1              
+#SBATCH --ntasks-per-node=8     # should = number of gpus
+#SBATCH --gres=gpu:8
+#SBATCH --time=32:00:00          # total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=fmri
+
+
+
+export NUM_GPUS=8  # Set to equal gres=gpu:#!
+export BATCH_SIZE=32
+export GLOBAL_BATCH_SIZE=$((BATCH_SIZE * NUM_GPUS))
+
+# Make sure another job doesnt use same port, here using random number
+export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000)) 
+
+export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export COUNT_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)
+
+export WANDB_DIR="/fsx/proj-fmri/ckadirt/MindEyeV2/src/wandb" 
+export WANDB_CACHE_DIR="/admin/home-ckadirt/.cache"
+export WANDB_MODE="online"
+
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo WORLD_SIZE=${COUNT_NODE}
+
+source /admin/home-ckadirt/.bashrc
+
+###########
+
+cd /fsx/proj-fmri/ckadirt/MindEyeV2/src/
+accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT Train_MLPMixer-Copy2.py --data_path=/fsx/proj-fmri/shared/mindeyev2_dataset --model_name=testing-rr-1024-past-5 --subj=1 --batch_size=${GLOBAL_BATCH_SIZE} --max_lr=3e-4 --mixup_pct=.66 --num_epochs=120 --ckpt_interval=10 --no-use_image_aug --hidden_dim=1024 --seq_len=5
+
+# --wandb_log

src/accel8.slurm

@@ -0,0 +1,41 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --partition=g40x
+#SBATCH --job-name=memoryrr
+#SBATCH --nodes=1              
+#SBATCH --ntasks-per-node=8     # should = number of gpus
+#SBATCH --gres=gpu:8
+#SBATCH --time=32:00:00          # total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=fmri
+
+
+
+export NUM_GPUS=8  # Set to equal gres=gpu:#!
+export BATCH_SIZE=32
+export GLOBAL_BATCH_SIZE=$((BATCH_SIZE * NUM_GPUS))
+
+# Make sure another job doesnt use same port, here using random number
+export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000)) 
+
+export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export COUNT_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)
+
+export WANDB_DIR="/fsx/proj-fmri/ckadirt/MindEyeV2/src/wandb" 
+export WANDB_CACHE_DIR="/admin/home-ckadirt/.cache"
+export WANDB_MODE="online"
+
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo WORLD_SIZE=${COUNT_NODE}
+
+source /admin/home-ckadirt/.bashrc
+
+###########
+
+cd /fsx/proj-fmri/ckadirt/MindEyeV2/src/
+accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT Train_MLPMixer-img.py --data_path=/fsx/proj-fmri/shared/mindeyev2_dataset --model_name=testing-rr-1024-img-past-2 --subj=1 --batch_size=${GLOBAL_BATCH_SIZE} --max_lr=3e-4 --mixup_pct=.66 --num_epochs=120 --ckpt_interval=10 --no-use_image_aug --hidden_dim=1024 --seq_len=2
+
+# --wandb_log

src/accel9.slurm

@@ -0,0 +1,44 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --partition=g40x
+#SBATCH --job-name=blip2captions          
+#SBATCH --nodes=1              
+#SBATCH --ntasks-per-node=1    # should = number of gpus
+#SBATCH --gres=gpu:1
+#SBATCH --time=24:00:00          # total run time limit (HH:MM:SS)
+#SBATCH --comment=medarc
+#SBATCH --requeue
+#SBATCH -e slurms/%j.err
+#SBATCH -o slurms/%j.out
+
+
+
+
+export NUM_GPUS=1  # Set to equal gres=gpu:#!
+export BATCH_SIZE=128
+export GLOBAL_BATCH_SIZE=$((BATCH_SIZE * NUM_GPUS))
+
+# Make sure another job doesnt use same port, here using random number
+export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000)) 
+
+export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export COUNT_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)
+
+export WANDB_DIR="/fsx/proj-fmri/ckadirt/MindEyeV2/src/wandb" 
+export WANDB_CACHE_DIR="/admin/home-ckadirt/.cache"
+export WANDB_MODE="online"
+
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo WORLD_SIZE=${COUNT_NODE}
+
+source /admin/home-ckadirt/.bashrc
+
+###########
+
+cd /fsx/proj-fmri/ckadirt/MindEyeV2/src/
+# accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT 
+python train2.py --data_path=/fsx/proj-fmri/shared/mindeyev2_dataset --model_name=caption_clip_0.5_bz --subj=1 --batch_size=${GLOBAL_BATCH_SIZE} --max_lr=1e-4 --mixup_pct=.66 --num_epochs=50 --use_image_aug --ckpt_interval=15 --clip_mse_ratio=0.5
+
+# --wandb_log

src/blip2_captions.py

@@ -0,0 +1,71 @@
+import os
+import sys
+import json
+import argparse
+import numpy as np
+import math
+from einops import rearrange
+import time
+import random
+import h5py
+from tqdm import tqdm
+
+import webdataset as wds
+import gc
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from torchvision.transforms import ToPILImage   #CHANGED (added)
+
+from accelerate import Accelerator, DeepSpeedPlugin
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+
+global_batch_size = 128 #128
+
+### Multi-GPU config ###
+local_rank = os.getenv('RANK')
+if local_rank is None: 
+    local_rank = 0
+else:
+    local_rank = int(local_rank)
+print("LOCAL RANK ", local_rank)  
+
+num_devices = torch.cuda.device_count()
+if num_devices==0: num_devices = 1
+
+accelerator = Accelerator(split_batches=False)
+
+### UNCOMMENT BELOW STUFF TO USE DEEPSPEED (also comment out the immediately above "accelerator = " line) ###
+
+# if num_devices <= 1 and utils.is_interactive():
+#     # can emulate a distributed environment for deepspeed to work in jupyter notebook
+#     os.environ["MASTER_ADDR"] = "localhost"
+#     os.environ["MASTER_PORT"] = str(np.random.randint(10000)+9000)
+#     os.environ["RANK"] = "0"
+#     os.environ["LOCAL_RANK"] = "0"
+#     os.environ["WORLD_SIZE"] = "1"
+#     os.environ["GLOBAL_BATCH_SIZE"] = str(global_batch_size) # set this to your batch size!
+#     global_batch_size = os.environ["GLOBAL_BATCH_SIZE"]
+
+# # alter the deepspeed config according to your global and local batch size
+# if local_rank == 0:
+#     with open('deepspeed_config_stage2.json', 'r') as file:
+#         config = json.load(file)
+#     config['train_batch_size'] = int(os.environ["GLOBAL_BATCH_SIZE"])
+#     config['train_micro_batch_size_per_gpu'] = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices
+#     with open('deepspeed_config_stage2.json', 'w') as file:
+#         json.dump(config, file)
+# else:
+#     # give some time for the local_rank=0 gpu to prep new deepspeed config file
+#     time.sleep(10)
+# deepspeed_plugin = DeepSpeedPlugin("deepspeed_config_stage2.json")
+# accelerator = Accelerator(split_batches=False, deepspeed_plugin=deepspeed_plugin)
+
+

src/checking_models.ipynb

@@ -0,0 +1,1526 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "ef9e1556-7840-4004-b181-a2c97ac2ab17",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b6f12dd4-f3aa-4981-b604-b72e67229011",
+   "metadata": {},
+   "source": [
+    "# DinoV2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "id": "2a604617-b602-4503-b288-e9828684505e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Using cache found in /fsx/proj-fmri/shared/cache/dinov2/hub/facebookresearch_dinov2_main\n"
+     ]
+    }
+   ],
+   "source": [
+    "# need to change TORCH_HOME env variable to specify pretrained model should go in shared folder, not home directory\n",
+    "os.environ['TORCH_HOME'] = '/fsx/proj-fmri/shared/cache/dinov2'\n",
+    "dinov2_model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14')\n",
+    "# remove initial image patching\n",
+    "dinov2_model.patch_embed = nn.Identity()\n",
+    "dinov2_model.patch_embed = nn.Identity()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "id": "32da913d-d931-4967-a5e8-bd40c21d1ad9",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([2, 33, 1024])\n"
+     ]
+    }
+   ],
+   "source": [
+    "dinov2_model.to(\"cuda\")\n",
+    "input = torch.randn((2,33,1024)).to(\"cuda\")\n",
+    "\n",
+    "for block in dinov2_model.blocks: input = block(input)\n",
+    "input = dinov2_model.norm(input)\n",
+    "\n",
+    "print(input.shape)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "febe89c0-06d0-4309-b378-a8d58b99bf4c",
+   "metadata": {},
+   "source": [
+    "# eva"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "id": "690204d0-13d7-452b-97af-14d144800e81",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from urllib.request import urlopen\n",
+    "from PIL import Image\n",
+    "import timm\n",
+    "\n",
+    "img = Image.open(urlopen(\n",
+    "    'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'\n",
+    "))\n",
+    "\n",
+    "model = timm.create_model(\n",
+    "    \"eva02_enormous_patch14_clip_224.laion2b\",\n",
+    "    pretrained=True,\n",
+    "    num_classes=0,  # remove classifier nn.Linear\n",
+    ")\n",
+    "model = model.eval()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 39,
+   "id": "035e3e9d-86c9-4ddf-b760-7b78dded7d2e",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "ValueError",
+     "evalue": "You have to specify pixel_values",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[39], line 5\u001b[0m\n\u001b[1;32m      2\u001b[0m data_config \u001b[38;5;241m=\u001b[39m timm\u001b[38;5;241m.\u001b[39mdata\u001b[38;5;241m.\u001b[39mresolve_model_data_config(model)\n\u001b[1;32m      3\u001b[0m transforms \u001b[38;5;241m=\u001b[39m timm\u001b[38;5;241m.\u001b[39mdata\u001b[38;5;241m.\u001b[39mcreate_transform(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mdata_config, is_training\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m)\n\u001b[0;32m----> 5\u001b[0m output \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtransforms\u001b[49m\u001b[43m(\u001b[49m\u001b[43mimg\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43munsqueeze\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# output is (batch_size, num_features) shaped tensor\u001b[39;00m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28mprint\u001b[39m(output\u001b[38;5;241m.\u001b[39mshape)\n\u001b[1;32m      7\u001b[0m \u001b[38;5;66;03m# or equivalently (without needing to set num_classes=0)\u001b[39;00m\n",
+      "File \u001b[0;32m~/miniconda3/envs/mindeye/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
+      "File \u001b[0;32m~/miniconda3/envs/mindeye/lib/python3.10/site-packages/transformers/models/clipseg/modeling_clipseg.py:1433\u001b[0m, in \u001b[0;36mCLIPSegForImageSegmentation.forward\u001b[0;34m(self, input_ids, pixel_values, conditional_pixel_values, conditional_embeddings, attention_mask, position_ids, labels, output_attentions, output_hidden_states, return_dict)\u001b[0m\n\u001b[1;32m   1431\u001b[0m \u001b[38;5;66;03m# step 1: forward the query images through the frozen CLIP vision encoder\u001b[39;00m\n\u001b[1;32m   1432\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mno_grad():\n\u001b[0;32m-> 1433\u001b[0m     vision_outputs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mclip\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvision_model\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1434\u001b[0m \u001b[43m        \u001b[49m\u001b[43mpixel_values\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpixel_values\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1435\u001b[0m \u001b[43m        \u001b[49m\u001b[43moutput_attentions\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moutput_attentions\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1436\u001b[0m \u001b[43m        \u001b[49m\u001b[43moutput_hidden_states\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# we need the intermediate hidden states\u001b[39;49;00m\n\u001b[1;32m   1437\u001b[0m \u001b[43m        \u001b[49m\u001b[43mreturn_dict\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_dict\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1438\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1439\u001b[0m     pooled_output \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mclip\u001b[38;5;241m.\u001b[39mvisual_projection(vision_outputs[\u001b[38;5;241m1\u001b[39m])\n\u001b[1;32m   1441\u001b[0m     hidden_states \u001b[38;5;241m=\u001b[39m vision_outputs\u001b[38;5;241m.\u001b[39mhidden_states \u001b[38;5;28;01mif\u001b[39;00m return_dict \u001b[38;5;28;01melse\u001b[39;00m vision_outputs[\u001b[38;5;241m2\u001b[39m]\n",
+      "File \u001b[0;32m~/miniconda3/envs/mindeye/lib/python3.10/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
+      "File \u001b[0;32m~/miniconda3/envs/mindeye/lib/python3.10/site-packages/transformers/models/clipseg/modeling_clipseg.py:872\u001b[0m, in \u001b[0;36mCLIPSegVisionTransformer.forward\u001b[0;34m(self, pixel_values, output_attentions, output_hidden_states, return_dict)\u001b[0m\n\u001b[1;32m    869\u001b[0m return_dict \u001b[38;5;241m=\u001b[39m return_dict \u001b[38;5;28;01mif\u001b[39;00m return_dict \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconfig\u001b[38;5;241m.\u001b[39muse_return_dict\n\u001b[1;32m    871\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m pixel_values \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 872\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mYou have to specify pixel_values\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    874\u001b[0m hidden_states \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39membeddings(pixel_values)\n\u001b[1;32m    875\u001b[0m hidden_states \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpre_layrnorm(hidden_states)\n",
+      "\u001b[0;31mValueError\u001b[0m: You have to specify pixel_values"
+     ]
+    }
+   ],
+   "source": [
+    "# get model specific transforms (normalization, resize)\n",
+    "data_config = timm.data.resolve_model_data_config(model)\n",
+    "transforms = timm.data.create_transform(**data_config, is_training=False)\n",
+    "\n",
+    "output = model(transforms(img).unsqueeze(0))  # output is (batch_size, num_features) shaped tensor\n",
+    "print(output.shape)\n",
+    "# or equivalently (without needing to set num_classes=0)\n",
+    "\n",
+    "output = model.forward_features(transforms(img).unsqueeze(0))\n",
+    "# output is unpooled, a (1, 257, 768) shaped tensor\n",
+    "print(output.shape)\n",
+    "\n",
+    "output = model.forward_head(output, pre_logits=True)\n",
+    "# output is a (1, num_features) shaped tensor\n",
+    "print(output.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "54275c4c-e506-4959-92f1-29e584f5ce51",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.forward_features(transforms(img).unsqueeze(0)).shape"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6546c673-f3ab-4d43-a051-cab20e782bab",
+   "metadata": {},
+   "source": [
+    "# Eva02-clip"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "id": "dfbc95de-9af9-4583-98fc-b8061114ef64",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import timm \n",
+    "# couldnt figure out how to load pretrained model from shared folder rather than home directory using timm...\n",
+    "eva02_model = timm.create_model(\"eva02_enormous_patch14_clip_224.laion2b\", pretrained=True)\n",
+    "# eva02_model.head_drop = nn.Identity()\n",
+    "# eva02_model.head = nn.Identity()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "id": "97e3ea29-ae6b-4bd2-b3d7-17839098a6e4",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([2, 1024])"
+      ]
+     },
+     "execution_count": 30,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "eva02_model(torch.randn((2,3,224,224))).shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "id": "069b76f0-029f-42b1-85f5-a492ee1cc5d1",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([2, 256, 1024])\n"
+     ]
+    }
+   ],
+   "source": [
+    "image = torch.randn((2,3,224,224))\n",
+    "\n",
+    "input = eva02_model.patch_embed(image)\n",
+    "input = eva02_model.pos_drop(input)\n",
+    "for block in eva02_model.blocks: input = block(input)\n",
+    "input = eva02_model.norm(input)\n",
+    "input = eva02_model.fc_norm(input)\n",
+    "input = eva02_model.head_drop(input)\n",
+    "input = eva02_model.head(input)\n",
+    "\n",
+    "print(input.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "id": "90e4e8e7-3dd1-43b0-a305-066a6ec13c2e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Help on Eva in module timm.models.eva object:\n",
+      "\n",
+      "class Eva(torch.nn.modules.module.Module)\n",
+      " |  Eva(img_size: Union[int, Tuple[int, int]] = 224, patch_size: Union[int, Tuple[int, int]] = 16, in_chans: int = 3, num_classes: int = 1000, global_pool: str = 'avg', embed_dim: int = 768, depth: int = 12, num_heads: int = 12, qkv_bias: bool = True, qkv_fused: bool = True, mlp_ratio: float = 4.0, swiglu_mlp: bool = False, scale_mlp: bool = False, scale_attn_inner: bool = False, drop_rate: float = 0.0, pos_drop_rate: float = 0.0, patch_drop_rate: float = 0.0, proj_drop_rate: float = 0.0, attn_drop_rate: float = 0.0, drop_path_rate: float = 0.0, norm_layer: Callable = <class 'timm.layers.norm.LayerNorm'>, init_values: Optional[float] = None, class_token: bool = True, use_abs_pos_emb: bool = True, use_rot_pos_emb: bool = False, use_post_norm: bool = False, ref_feat_shape: Union[int, Tuple[int, int], NoneType] = None, head_init_scale: float = 0.001)\n",
+      " |  \n",
+      " |  Eva Vision Transformer w/ Abs & Rotary Pos Embed\n",
+      " |  \n",
+      " |  This class implements the EVA and EVA02 models that were based on the BEiT ViT variant\n",
+      " |    * EVA - abs pos embed, global avg pool\n",
+      " |    * EVA02 - abs + rope pos embed, global avg pool, SwiGLU, scale Norm in MLP (ala normformer)\n",
+      " |  \n",
+      " |  Method resolution order:\n",
+      " |      Eva\n",
+      " |      torch.nn.modules.module.Module\n",
+      " |      builtins.object\n",
+      " |  \n",
+      " |  Methods defined here:\n",
+      " |  \n",
+      " |  __init__(self, img_size: Union[int, Tuple[int, int]] = 224, patch_size: Union[int, Tuple[int, int]] = 16, in_chans: int = 3, num_classes: int = 1000, global_pool: str = 'avg', embed_dim: int = 768, depth: int = 12, num_heads: int = 12, qkv_bias: bool = True, qkv_fused: bool = True, mlp_ratio: float = 4.0, swiglu_mlp: bool = False, scale_mlp: bool = False, scale_attn_inner: bool = False, drop_rate: float = 0.0, pos_drop_rate: float = 0.0, patch_drop_rate: float = 0.0, proj_drop_rate: float = 0.0, attn_drop_rate: float = 0.0, drop_path_rate: float = 0.0, norm_layer: Callable = <class 'timm.layers.norm.LayerNorm'>, init_values: Optional[float] = None, class_token: bool = True, use_abs_pos_emb: bool = True, use_rot_pos_emb: bool = False, use_post_norm: bool = False, ref_feat_shape: Union[int, Tuple[int, int], NoneType] = None, head_init_scale: float = 0.001)\n",
+      " |      Args:\n",
+      " |          img_size:\n",
+      " |          patch_size:\n",
+      " |          in_chans:\n",
+      " |          num_classes:\n",
+      " |          global_pool:\n",
+      " |          embed_dim:\n",
+      " |          depth:\n",
+      " |          num_heads:\n",
+      " |          qkv_bias:\n",
+      " |          qkv_fused:\n",
+      " |          mlp_ratio:\n",
+      " |          swiglu_mlp:\n",
+      " |          scale_mlp:\n",
+      " |          scale_attn_inner:\n",
+      " |          drop_rate:\n",
+      " |          pos_drop_rate:\n",
+      " |          proj_drop_rate:\n",
+      " |          attn_drop_rate:\n",
+      " |          drop_path_rate:\n",
+      " |          norm_layer:\n",
+      " |          init_values:\n",
+      " |          class_token:\n",
+      " |          use_abs_pos_emb:\n",
+      " |          use_rot_pos_emb:\n",
+      " |          use_post_norm:\n",
+      " |          ref_feat_shape:\n",
+      " |          head_init_scale:\n",
+      " |  \n",
+      " |  fix_init_weight(self)\n",
+      " |  \n",
+      " |  forward(self, x)\n",
+      " |      Defines the computation performed at every call.\n",
+      " |      \n",
+      " |      Should be overridden by all subclasses.\n",
+      " |      \n",
+      " |      .. note::\n",
+      " |          Although the recipe for forward pass needs to be defined within\n",
+      " |          this function, one should call the :class:`Module` instance afterwards\n",
+      " |          instead of this since the former takes care of running the\n",
+      " |          registered hooks while the latter silently ignores them.\n",
+      " |  \n",
+      " |  forward_features(self, x)\n",
+      " |  \n",
+      " |  forward_head(self, x, pre_logits: bool = False)\n",
+      " |  \n",
+      " |  get_classifier(self)\n",
+      " |  \n",
+      " |  group_matcher(self, coarse=False)\n",
+      " |  \n",
+      " |  no_weight_decay(self)\n",
+      " |  \n",
+      " |  reset_classifier(self, num_classes, global_pool=None)\n",
+      " |  \n",
+      " |  set_grad_checkpointing(self, enable=True)\n",
+      " |  \n",
+      " |  ----------------------------------------------------------------------\n",
+      " |  Data and other attributes defined here:\n",
+      " |  \n",
+      " |  __annotations__ = {}\n",
+      " |  \n",
+      " |  ----------------------------------------------------------------------\n",
+      " |  Methods inherited from torch.nn.modules.module.Module:\n",
+      " |  \n",
+      " |  __call__ = _call_impl(self, *args, **kwargs)\n",
+      " |  \n",
+      " |  __delattr__(self, name)\n",
+      " |      Implement delattr(self, name).\n",
+      " |  \n",
+      " |  __dir__(self)\n",
+      " |      Default dir() implementation.\n",
+      " |  \n",
+      " |  __getattr__(self, name: str) -> Union[torch.Tensor, ForwardRef('Module')]\n",
+      " |  \n",
+      " |  __repr__(self)\n",
+      " |      Return repr(self).\n",
+      " |  \n",
+      " |  __setattr__(self, name: str, value: Union[torch.Tensor, ForwardRef('Module')]) -> None\n",
+      " |      Implement setattr(self, name, value).\n",
+      " |  \n",
+      " |  __setstate__(self, state)\n",
+      " |  \n",
+      " |  add_module(self, name: str, module: Optional[ForwardRef('Module')]) -> None\n",
+      " |      Adds a child module to the current module.\n",
+      " |      \n",
+      " |      The module can be accessed as an attribute using the given name.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          name (str): name of the child module. The child module can be\n",
+      " |              accessed from this module using the given name\n",
+      " |          module (Module): child module to be added to the module.\n",
+      " |  \n",
+      " |  apply(self: ~T, fn: Callable[[ForwardRef('Module')], NoneType]) -> ~T\n",
+      " |      Applies ``fn`` recursively to every submodule (as returned by ``.children()``)\n",
+      " |      as well as self. Typical use includes initializing the parameters of a model\n",
+      " |      (see also :ref:`nn-init-doc`).\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          fn (:class:`Module` -> None): function to be applied to each submodule\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          Module: self\n",
+      " |      \n",
+      " |      Example::\n",
+      " |      \n",
+      " |          >>> @torch.no_grad()\n",
+      " |          >>> def init_weights(m):\n",
+      " |          >>>     print(m)\n",
+      " |          >>>     if type(m) == nn.Linear:\n",
+      " |          >>>         m.weight.fill_(1.0)\n",
+      " |          >>>         print(m.weight)\n",
+      " |          >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))\n",
+      " |          >>> net.apply(init_weights)\n",
+      " |          Linear(in_features=2, out_features=2, bias=True)\n",
+      " |          Parameter containing:\n",
+      " |          tensor([[1., 1.],\n",
+      " |                  [1., 1.]], requires_grad=True)\n",
+      " |          Linear(in_features=2, out_features=2, bias=True)\n",
+      " |          Parameter containing:\n",
+      " |          tensor([[1., 1.],\n",
+      " |                  [1., 1.]], requires_grad=True)\n",
+      " |          Sequential(\n",
+      " |            (0): Linear(in_features=2, out_features=2, bias=True)\n",
+      " |            (1): Linear(in_features=2, out_features=2, bias=True)\n",
+      " |          )\n",
+      " |  \n",
+      " |  bfloat16(self: ~T) -> ~T\n",
+      " |      Casts all floating point parameters and buffers to ``bfloat16`` datatype.\n",
+      " |      \n",
+      " |      .. note::\n",
+      " |          This method modifies the module in-place.\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          Module: self\n",
+      " |  \n",
+      " |  buffers(self, recurse: bool = True) -> Iterator[torch.Tensor]\n",
+      " |      Returns an iterator over module buffers.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          recurse (bool): if True, then yields buffers of this module\n",
+      " |              and all submodules. Otherwise, yields only buffers that\n",
+      " |              are direct members of this module.\n",
+      " |      \n",
+      " |      Yields:\n",
+      " |          torch.Tensor: module buffer\n",
+      " |      \n",
+      " |      Example::\n",
+      " |      \n",
+      " |          >>> # xdoctest: +SKIP(\"undefined vars\")\n",
+      " |          >>> for buf in model.buffers():\n",
+      " |          >>>     print(type(buf), buf.size())\n",
+      " |          <class 'torch.Tensor'> (20L,)\n",
+      " |          <class 'torch.Tensor'> (20L, 1L, 5L, 5L)\n",
+      " |  \n",
+      " |  children(self) -> Iterator[ForwardRef('Module')]\n",
+      " |      Returns an iterator over immediate children modules.\n",
+      " |      \n",
+      " |      Yields:\n",
+      " |          Module: a child module\n",
+      " |  \n",
+      " |  cpu(self: ~T) -> ~T\n",
+      " |      Moves all model parameters and buffers to the CPU.\n",
+      " |      \n",
+      " |      .. note::\n",
+      " |          This method modifies the module in-place.\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          Module: self\n",
+      " |  \n",
+      " |  cuda(self: ~T, device: Union[int, torch.device, NoneType] = None) -> ~T\n",
+      " |      Moves all model parameters and buffers to the GPU.\n",
+      " |      \n",
+      " |      This also makes associated parameters and buffers different objects. So\n",
+      " |      it should be called before constructing optimizer if the module will\n",
+      " |      live on GPU while being optimized.\n",
+      " |      \n",
+      " |      .. note::\n",
+      " |          This method modifies the module in-place.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          device (int, optional): if specified, all parameters will be\n",
+      " |              copied to that device\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          Module: self\n",
+      " |  \n",
+      " |  double(self: ~T) -> ~T\n",
+      " |      Casts all floating point parameters and buffers to ``double`` datatype.\n",
+      " |      \n",
+      " |      .. note::\n",
+      " |          This method modifies the module in-place.\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          Module: self\n",
+      " |  \n",
+      " |  eval(self: ~T) -> ~T\n",
+      " |      Sets the module in evaluation mode.\n",
+      " |      \n",
+      " |      This has any effect only on certain modules. See documentations of\n",
+      " |      particular modules for details of their behaviors in training/evaluation\n",
+      " |      mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,\n",
+      " |      etc.\n",
+      " |      \n",
+      " |      This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.\n",
+      " |      \n",
+      " |      See :ref:`locally-disable-grad-doc` for a comparison between\n",
+      " |      `.eval()` and several similar mechanisms that may be confused with it.\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          Module: self\n",
+      " |  \n",
+      " |  extra_repr(self) -> str\n",
+      " |      Set the extra representation of the module\n",
+      " |      \n",
+      " |      To print customized extra information, you should re-implement\n",
+      " |      this method in your own modules. Both single-line and multi-line\n",
+      " |      strings are acceptable.\n",
+      " |  \n",
+      " |  float(self: ~T) -> ~T\n",
+      " |      Casts all floating point parameters and buffers to ``float`` datatype.\n",
+      " |      \n",
+      " |      .. note::\n",
+      " |          This method modifies the module in-place.\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          Module: self\n",
+      " |  \n",
+      " |  get_buffer(self, target: str) -> 'Tensor'\n",
+      " |      Returns the buffer given by ``target`` if it exists,\n",
+      " |      otherwise throws an error.\n",
+      " |      \n",
+      " |      See the docstring for ``get_submodule`` for a more detailed\n",
+      " |      explanation of this method's functionality as well as how to\n",
+      " |      correctly specify ``target``.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          target: The fully-qualified string name of the buffer\n",
+      " |              to look for. (See ``get_submodule`` for how to specify a\n",
+      " |              fully-qualified string.)\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          torch.Tensor: The buffer referenced by ``target``\n",
+      " |      \n",
+      " |      Raises:\n",
+      " |          AttributeError: If the target string references an invalid\n",
+      " |              path or resolves to something that is not a\n",
+      " |              buffer\n",
+      " |  \n",
+      " |  get_extra_state(self) -> Any\n",
+      " |      Returns any extra state to include in the module's state_dict.\n",
+      " |      Implement this and a corresponding :func:`set_extra_state` for your module\n",
+      " |      if you need to store extra state. This function is called when building the\n",
+      " |      module's `state_dict()`.\n",
+      " |      \n",
+      " |      Note that extra state should be picklable to ensure working serialization\n",
+      " |      of the state_dict. We only provide provide backwards compatibility guarantees\n",
+      " |      for serializing Tensors; other objects may break backwards compatibility if\n",
+      " |      their serialized pickled form changes.\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          object: Any extra state to store in the module's state_dict\n",
+      " |  \n",
+      " |  get_parameter(self, target: str) -> 'Parameter'\n",
+      " |      Returns the parameter given by ``target`` if it exists,\n",
+      " |      otherwise throws an error.\n",
+      " |      \n",
+      " |      See the docstring for ``get_submodule`` for a more detailed\n",
+      " |      explanation of this method's functionality as well as how to\n",
+      " |      correctly specify ``target``.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          target: The fully-qualified string name of the Parameter\n",
+      " |              to look for. (See ``get_submodule`` for how to specify a\n",
+      " |              fully-qualified string.)\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          torch.nn.Parameter: The Parameter referenced by ``target``\n",
+      " |      \n",
+      " |      Raises:\n",
+      " |          AttributeError: If the target string references an invalid\n",
+      " |              path or resolves to something that is not an\n",
+      " |              ``nn.Parameter``\n",
+      " |  \n",
+      " |  get_submodule(self, target: str) -> 'Module'\n",
+      " |      Returns the submodule given by ``target`` if it exists,\n",
+      " |      otherwise throws an error.\n",
+      " |      \n",
+      " |      For example, let's say you have an ``nn.Module`` ``A`` that\n",
+      " |      looks like this:\n",
+      " |      \n",
+      " |      .. code-block:: text\n",
+      " |      \n",
+      " |          A(\n",
+      " |              (net_b): Module(\n",
+      " |                  (net_c): Module(\n",
+      " |                      (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))\n",
+      " |                  )\n",
+      " |                  (linear): Linear(in_features=100, out_features=200, bias=True)\n",
+      " |              )\n",
+      " |          )\n",
+      " |      \n",
+      " |      (The diagram shows an ``nn.Module`` ``A``. ``A`` has a nested\n",
+      " |      submodule ``net_b``, which itself has two submodules ``net_c``\n",
+      " |      and ``linear``. ``net_c`` then has a submodule ``conv``.)\n",
+      " |      \n",
+      " |      To check whether or not we have the ``linear`` submodule, we\n",
+      " |      would call ``get_submodule(\"net_b.linear\")``. To check whether\n",
+      " |      we have the ``conv`` submodule, we would call\n",
+      " |      ``get_submodule(\"net_b.net_c.conv\")``.\n",
+      " |      \n",
+      " |      The runtime of ``get_submodule`` is bounded by the degree\n",
+      " |      of module nesting in ``target``. A query against\n",
+      " |      ``named_modules`` achieves the same result, but it is O(N) in\n",
+      " |      the number of transitive modules. So, for a simple check to see\n",
+      " |      if some submodule exists, ``get_submodule`` should always be\n",
+      " |      used.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          target: The fully-qualified string name of the submodule\n",
+      " |              to look for. (See above example for how to specify a\n",
+      " |              fully-qualified string.)\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          torch.nn.Module: The submodule referenced by ``target``\n",
+      " |      \n",
+      " |      Raises:\n",
+      " |          AttributeError: If the target string references an invalid\n",
+      " |              path or resolves to something that is not an\n",
+      " |              ``nn.Module``\n",
+      " |  \n",
+      " |  half(self: ~T) -> ~T\n",
+      " |      Casts all floating point parameters and buffers to ``half`` datatype.\n",
+      " |      \n",
+      " |      .. note::\n",
+      " |          This method modifies the module in-place.\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          Module: self\n",
+      " |  \n",
+      " |  ipu(self: ~T, device: Union[int, torch.device, NoneType] = None) -> ~T\n",
+      " |      Moves all model parameters and buffers to the IPU.\n",
+      " |      \n",
+      " |      This also makes associated parameters and buffers different objects. So\n",
+      " |      it should be called before constructing optimizer if the module will\n",
+      " |      live on IPU while being optimized.\n",
+      " |      \n",
+      " |      .. note::\n",
+      " |          This method modifies the module in-place.\n",
+      " |      \n",
+      " |      Arguments:\n",
+      " |          device (int, optional): if specified, all parameters will be\n",
+      " |              copied to that device\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          Module: self\n",
+      " |  \n",
+      " |  load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True)\n",
+      " |      Copies parameters and buffers from :attr:`state_dict` into\n",
+      " |      this module and its descendants. If :attr:`strict` is ``True``, then\n",
+      " |      the keys of :attr:`state_dict` must exactly match the keys returned\n",
+      " |      by this module's :meth:`~torch.nn.Module.state_dict` function.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          state_dict (dict): a dict containing parameters and\n",
+      " |              persistent buffers.\n",
+      " |          strict (bool, optional): whether to strictly enforce that the keys\n",
+      " |              in :attr:`state_dict` match the keys returned by this module's\n",
+      " |              :meth:`~torch.nn.Module.state_dict` function. Default: ``True``\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:\n",
+      " |              * **missing_keys** is a list of str containing the missing keys\n",
+      " |              * **unexpected_keys** is a list of str containing the unexpected keys\n",
+      " |      \n",
+      " |      Note:\n",
+      " |          If a parameter or buffer is registered as ``None`` and its corresponding key\n",
+      " |          exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a\n",
+      " |          ``RuntimeError``.\n",
+      " |  \n",
+      " |  modules(self) -> Iterator[ForwardRef('Module')]\n",
+      " |      Returns an iterator over all modules in the network.\n",
+      " |      \n",
+      " |      Yields:\n",
+      " |          Module: a module in the network\n",
+      " |      \n",
+      " |      Note:\n",
+      " |          Duplicate modules are returned only once. In the following\n",
+      " |          example, ``l`` will be returned only once.\n",
+      " |      \n",
+      " |      Example::\n",
+      " |      \n",
+      " |          >>> l = nn.Linear(2, 2)\n",
+      " |          >>> net = nn.Sequential(l, l)\n",
+      " |          >>> for idx, m in enumerate(net.modules()):\n",
+      " |          ...     print(idx, '->', m)\n",
+      " |      \n",
+      " |          0 -> Sequential(\n",
+      " |            (0): Linear(in_features=2, out_features=2, bias=True)\n",
+      " |            (1): Linear(in_features=2, out_features=2, bias=True)\n",
+      " |          )\n",
+      " |          1 -> Linear(in_features=2, out_features=2, bias=True)\n",
+      " |  \n",
+      " |  named_buffers(self, prefix: str = '', recurse: bool = True, remove_duplicate: bool = True) -> Iterator[Tuple[str, torch.Tensor]]\n",
+      " |      Returns an iterator over module buffers, yielding both the\n",
+      " |      name of the buffer as well as the buffer itself.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          prefix (str): prefix to prepend to all buffer names.\n",
+      " |          recurse (bool, optional): if True, then yields buffers of this module\n",
+      " |              and all submodules. Otherwise, yields only buffers that\n",
+      " |              are direct members of this module. Defaults to True.\n",
+      " |          remove_duplicate (bool, optional): whether to remove the duplicated buffers in the result. Defaults to True.\n",
+      " |      \n",
+      " |      Yields:\n",
+      " |          (str, torch.Tensor): Tuple containing the name and buffer\n",
+      " |      \n",
+      " |      Example::\n",
+      " |      \n",
+      " |          >>> # xdoctest: +SKIP(\"undefined vars\")\n",
+      " |          >>> for name, buf in self.named_buffers():\n",
+      " |          >>>     if name in ['running_var']:\n",
+      " |          >>>         print(buf.size())\n",
+      " |  \n",
+      " |  named_children(self) -> Iterator[Tuple[str, ForwardRef('Module')]]\n",
+      " |      Returns an iterator over immediate children modules, yielding both\n",
+      " |      the name of the module as well as the module itself.\n",
+      " |      \n",
+      " |      Yields:\n",
+      " |          (str, Module): Tuple containing a name and child module\n",
+      " |      \n",
+      " |      Example::\n",
+      " |      \n",
+      " |          >>> # xdoctest: +SKIP(\"undefined vars\")\n",
+      " |          >>> for name, module in model.named_children():\n",
+      " |          >>>     if name in ['conv4', 'conv5']:\n",
+      " |          >>>         print(module)\n",
+      " |  \n",
+      " |  named_modules(self, memo: Optional[Set[ForwardRef('Module')]] = None, prefix: str = '', remove_duplicate: bool = True)\n",
+      " |      Returns an iterator over all modules in the network, yielding\n",
+      " |      both the name of the module as well as the module itself.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          memo: a memo to store the set of modules already added to the result\n",
+      " |          prefix: a prefix that will be added to the name of the module\n",
+      " |          remove_duplicate: whether to remove the duplicated module instances in the result\n",
+      " |              or not\n",
+      " |      \n",
+      " |      Yields:\n",
+      " |          (str, Module): Tuple of name and module\n",
+      " |      \n",
+      " |      Note:\n",
+      " |          Duplicate modules are returned only once. In the following\n",
+      " |          example, ``l`` will be returned only once.\n",
+      " |      \n",
+      " |      Example::\n",
+      " |      \n",
+      " |          >>> l = nn.Linear(2, 2)\n",
+      " |          >>> net = nn.Sequential(l, l)\n",
+      " |          >>> for idx, m in enumerate(net.named_modules()):\n",
+      " |          ...     print(idx, '->', m)\n",
+      " |      \n",
+      " |          0 -> ('', Sequential(\n",
+      " |            (0): Linear(in_features=2, out_features=2, bias=True)\n",
+      " |            (1): Linear(in_features=2, out_features=2, bias=True)\n",
+      " |          ))\n",
+      " |          1 -> ('0', Linear(in_features=2, out_features=2, bias=True))\n",
+      " |  \n",
+      " |  named_parameters(self, prefix: str = '', recurse: bool = True, remove_duplicate: bool = True) -> Iterator[Tuple[str, torch.nn.parameter.Parameter]]\n",
+      " |      Returns an iterator over module parameters, yielding both the\n",
+      " |      name of the parameter as well as the parameter itself.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          prefix (str): prefix to prepend to all parameter names.\n",
+      " |          recurse (bool): if True, then yields parameters of this module\n",
+      " |              and all submodules. Otherwise, yields only parameters that\n",
+      " |              are direct members of this module.\n",
+      " |          remove_duplicate (bool, optional): whether to remove the duplicated\n",
+      " |              parameters in the result. Defaults to True.\n",
+      " |      \n",
+      " |      Yields:\n",
+      " |          (str, Parameter): Tuple containing the name and parameter\n",
+      " |      \n",
+      " |      Example::\n",
+      " |      \n",
+      " |          >>> # xdoctest: +SKIP(\"undefined vars\")\n",
+      " |          >>> for name, param in self.named_parameters():\n",
+      " |          >>>     if name in ['bias']:\n",
+      " |          >>>         print(param.size())\n",
+      " |  \n",
+      " |  parameters(self, recurse: bool = True) -> Iterator[torch.nn.parameter.Parameter]\n",
+      " |      Returns an iterator over module parameters.\n",
+      " |      \n",
+      " |      This is typically passed to an optimizer.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          recurse (bool): if True, then yields parameters of this module\n",
+      " |              and all submodules. Otherwise, yields only parameters that\n",
+      " |              are direct members of this module.\n",
+      " |      \n",
+      " |      Yields:\n",
+      " |          Parameter: module parameter\n",
+      " |      \n",
+      " |      Example::\n",
+      " |      \n",
+      " |          >>> # xdoctest: +SKIP(\"undefined vars\")\n",
+      " |          >>> for param in model.parameters():\n",
+      " |          >>>     print(type(param), param.size())\n",
+      " |          <class 'torch.Tensor'> (20L,)\n",
+      " |          <class 'torch.Tensor'> (20L, 1L, 5L, 5L)\n",
+      " |  \n",
+      " |  register_backward_hook(self, hook: Callable[[ForwardRef('Module'), Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, Tuple[torch.Tensor, ...], torch.Tensor]]) -> torch.utils.hooks.RemovableHandle\n",
+      " |      Registers a backward hook on the module.\n",
+      " |      \n",
+      " |      This function is deprecated in favor of :meth:`~torch.nn.Module.register_full_backward_hook` and\n",
+      " |      the behavior of this function will change in future versions.\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
+      " |              a handle that can be used to remove the added hook by calling\n",
+      " |              ``handle.remove()``\n",
+      " |  \n",
+      " |  register_buffer(self, name: str, tensor: Optional[torch.Tensor], persistent: bool = True) -> None\n",
+      " |      Adds a buffer to the module.\n",
+      " |      \n",
+      " |      This is typically used to register a buffer that should not to be\n",
+      " |      considered a model parameter. For example, BatchNorm's ``running_mean``\n",
+      " |      is not a parameter, but is part of the module's state. Buffers, by\n",
+      " |      default, are persistent and will be saved alongside parameters. This\n",
+      " |      behavior can be changed by setting :attr:`persistent` to ``False``. The\n",
+      " |      only difference between a persistent buffer and a non-persistent buffer\n",
+      " |      is that the latter will not be a part of this module's\n",
+      " |      :attr:`state_dict`.\n",
+      " |      \n",
+      " |      Buffers can be accessed as attributes using given names.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          name (str): name of the buffer. The buffer can be accessed\n",
+      " |              from this module using the given name\n",
+      " |          tensor (Tensor or None): buffer to be registered. If ``None``, then operations\n",
+      " |              that run on buffers, such as :attr:`cuda`, are ignored. If ``None``,\n",
+      " |              the buffer is **not** included in the module's :attr:`state_dict`.\n",
+      " |          persistent (bool): whether the buffer is part of this module's\n",
+      " |              :attr:`state_dict`.\n",
+      " |      \n",
+      " |      Example::\n",
+      " |      \n",
+      " |          >>> # xdoctest: +SKIP(\"undefined vars\")\n",
+      " |          >>> self.register_buffer('running_mean', torch.zeros(num_features))\n",
+      " |  \n",
+      " |  register_forward_hook(self, hook: Union[Callable[[~T, Tuple[Any, ...], Any], Optional[Any]], Callable[[~T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]]], *, prepend: bool = False, with_kwargs: bool = False) -> torch.utils.hooks.RemovableHandle\n",
+      " |      Registers a forward hook on the module.\n",
+      " |      \n",
+      " |      The hook will be called every time after :func:`forward` has computed an output.\n",
+      " |      \n",
+      " |      If ``with_kwargs`` is ``False`` or not specified, the input contains only\n",
+      " |      the positional arguments given to the module. Keyword arguments won't be\n",
+      " |      passed to the hooks and only to the ``forward``. The hook can modify the\n",
+      " |      output. It can modify the input inplace but it will not have effect on\n",
+      " |      forward since this is called after :func:`forward` is called. The hook\n",
+      " |      should have the following signature::\n",
+      " |      \n",
+      " |          hook(module, args, output) -> None or modified output\n",
+      " |      \n",
+      " |      If ``with_kwargs`` is ``True``, the forward hook will be passed the\n",
+      " |      ``kwargs`` given to the forward function and be expected to return the\n",
+      " |      output possibly modified. The hook should have the following signature::\n",
+      " |      \n",
+      " |          hook(module, args, kwargs, output) -> None or modified output\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          hook (Callable): The user defined hook to be registered.\n",
+      " |          prepend (bool): If ``True``, the provided ``hook`` will be fired\n",
+      " |              before all existing ``forward`` hooks on this\n",
+      " |              :class:`torch.nn.modules.Module`. Otherwise, the provided\n",
+      " |              ``hook`` will be fired after all existing ``forward`` hooks on\n",
+      " |              this :class:`torch.nn.modules.Module`. Note that global\n",
+      " |              ``forward`` hooks registered with\n",
+      " |              :func:`register_module_forward_hook` will fire before all hooks\n",
+      " |              registered by this method.\n",
+      " |              Default: ``False``\n",
+      " |          with_kwargs (bool): If ``True``, the ``hook`` will be passed the\n",
+      " |              kwargs given to the forward function.\n",
+      " |              Default: ``False``\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
+      " |              a handle that can be used to remove the added hook by calling\n",
+      " |              ``handle.remove()``\n",
+      " |  \n",
+      " |  register_forward_pre_hook(self, hook: Union[Callable[[~T, Tuple[Any, ...]], Optional[Any]], Callable[[~T, Tuple[Any, ...], Dict[str, Any]], Optional[Tuple[Any, Dict[str, Any]]]]], *, prepend: bool = False, with_kwargs: bool = False) -> torch.utils.hooks.RemovableHandle\n",
+      " |      Registers a forward pre-hook on the module.\n",
+      " |      \n",
+      " |      The hook will be called every time before :func:`forward` is invoked.\n",
+      " |      \n",
+      " |      \n",
+      " |      If ``with_kwargs`` is false or not specified, the input contains only\n",
+      " |      the positional arguments given to the module. Keyword arguments won't be\n",
+      " |      passed to the hooks and only to the ``forward``. The hook can modify the\n",
+      " |      input. User can either return a tuple or a single modified value in the\n",
+      " |      hook. We will wrap the value into a tuple if a single value is returned\n",
+      " |      (unless that value is already a tuple). The hook should have the\n",
+      " |      following signature::\n",
+      " |      \n",
+      " |          hook(module, args) -> None or modified input\n",
+      " |      \n",
+      " |      If ``with_kwargs`` is true, the forward pre-hook will be passed the\n",
+      " |      kwargs given to the forward function. And if the hook modifies the\n",
+      " |      input, both the args and kwargs should be returned. The hook should have\n",
+      " |      the following signature::\n",
+      " |      \n",
+      " |          hook(module, args, kwargs) -> None or a tuple of modified input and kwargs\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          hook (Callable): The user defined hook to be registered.\n",
+      " |          prepend (bool): If true, the provided ``hook`` will be fired before\n",
+      " |              all existing ``forward_pre`` hooks on this\n",
+      " |              :class:`torch.nn.modules.Module`. Otherwise, the provided\n",
+      " |              ``hook`` will be fired after all existing ``forward_pre`` hooks\n",
+      " |              on this :class:`torch.nn.modules.Module`. Note that global\n",
+      " |              ``forward_pre`` hooks registered with\n",
+      " |              :func:`register_module_forward_pre_hook` will fire before all\n",
+      " |              hooks registered by this method.\n",
+      " |              Default: ``False``\n",
+      " |          with_kwargs (bool): If true, the ``hook`` will be passed the kwargs\n",
+      " |              given to the forward function.\n",
+      " |              Default: ``False``\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
+      " |              a handle that can be used to remove the added hook by calling\n",
+      " |              ``handle.remove()``\n",
+      " |  \n",
+      " |  register_full_backward_hook(self, hook: Callable[[ForwardRef('Module'), Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, Tuple[torch.Tensor, ...], torch.Tensor]], prepend: bool = False) -> torch.utils.hooks.RemovableHandle\n",
+      " |      Registers a backward hook on the module.\n",
+      " |      \n",
+      " |      The hook will be called every time the gradients with respect to a module\n",
+      " |      are computed, i.e. the hook will execute if and only if the gradients with\n",
+      " |      respect to module outputs are computed. The hook should have the following\n",
+      " |      signature::\n",
+      " |      \n",
+      " |          hook(module, grad_input, grad_output) -> tuple(Tensor) or None\n",
+      " |      \n",
+      " |      The :attr:`grad_input` and :attr:`grad_output` are tuples that contain the gradients\n",
+      " |      with respect to the inputs and outputs respectively. The hook should\n",
+      " |      not modify its arguments, but it can optionally return a new gradient with\n",
+      " |      respect to the input that will be used in place of :attr:`grad_input` in\n",
+      " |      subsequent computations. :attr:`grad_input` will only correspond to the inputs given\n",
+      " |      as positional arguments and all kwarg arguments are ignored. Entries\n",
+      " |      in :attr:`grad_input` and :attr:`grad_output` will be ``None`` for all non-Tensor\n",
+      " |      arguments.\n",
+      " |      \n",
+      " |      For technical reasons, when this hook is applied to a Module, its forward function will\n",
+      " |      receive a view of each Tensor passed to the Module. Similarly the caller will receive a view\n",
+      " |      of each Tensor returned by the Module's forward function.\n",
+      " |      \n",
+      " |      .. warning ::\n",
+      " |          Modifying inputs or outputs inplace is not allowed when using backward hooks and\n",
+      " |          will raise an error.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          hook (Callable): The user-defined hook to be registered.\n",
+      " |          prepend (bool): If true, the provided ``hook`` will be fired before\n",
+      " |              all existing ``backward`` hooks on this\n",
+      " |              :class:`torch.nn.modules.Module`. Otherwise, the provided\n",
+      " |              ``hook`` will be fired after all existing ``backward`` hooks on\n",
+      " |              this :class:`torch.nn.modules.Module`. Note that global\n",
+      " |              ``backward`` hooks registered with\n",
+      " |              :func:`register_module_full_backward_hook` will fire before\n",
+      " |              all hooks registered by this method.\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
+      " |              a handle that can be used to remove the added hook by calling\n",
+      " |              ``handle.remove()``\n",
+      " |  \n",
+      " |  register_full_backward_pre_hook(self, hook: Callable[[ForwardRef('Module'), Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, Tuple[torch.Tensor, ...], torch.Tensor]], prepend: bool = False) -> torch.utils.hooks.RemovableHandle\n",
+      " |      Registers a backward pre-hook on the module.\n",
+      " |      \n",
+      " |      The hook will be called every time the gradients for the module are computed.\n",
+      " |      The hook should have the following signature::\n",
+      " |      \n",
+      " |          hook(module, grad_output) -> Tensor or None\n",
+      " |      \n",
+      " |      The :attr:`grad_output` is a tuple. The hook should\n",
+      " |      not modify its arguments, but it can optionally return a new gradient with\n",
+      " |      respect to the output that will be used in place of :attr:`grad_output` in\n",
+      " |      subsequent computations. Entries in :attr:`grad_output` will be ``None`` for\n",
+      " |      all non-Tensor arguments.\n",
+      " |      \n",
+      " |      For technical reasons, when this hook is applied to a Module, its forward function will\n",
+      " |      receive a view of each Tensor passed to the Module. Similarly the caller will receive a view\n",
+      " |      of each Tensor returned by the Module's forward function.\n",
+      " |      \n",
+      " |      .. warning ::\n",
+      " |          Modifying inputs inplace is not allowed when using backward hooks and\n",
+      " |          will raise an error.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          hook (Callable): The user-defined hook to be registered.\n",
+      " |          prepend (bool): If true, the provided ``hook`` will be fired before\n",
+      " |              all existing ``backward_pre`` hooks on this\n",
+      " |              :class:`torch.nn.modules.Module`. Otherwise, the provided\n",
+      " |              ``hook`` will be fired after all existing ``backward_pre`` hooks\n",
+      " |              on this :class:`torch.nn.modules.Module`. Note that global\n",
+      " |              ``backward_pre`` hooks registered with\n",
+      " |              :func:`register_module_full_backward_pre_hook` will fire before\n",
+      " |              all hooks registered by this method.\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
+      " |              a handle that can be used to remove the added hook by calling\n",
+      " |              ``handle.remove()``\n",
+      " |  \n",
+      " |  register_load_state_dict_post_hook(self, hook)\n",
+      " |      Registers a post hook to be run after module's ``load_state_dict``\n",
+      " |      is called.\n",
+      " |      \n",
+      " |      It should have the following signature::\n",
+      " |          hook(module, incompatible_keys) -> None\n",
+      " |      \n",
+      " |      The ``module`` argument is the current module that this hook is registered\n",
+      " |      on, and the ``incompatible_keys`` argument is a ``NamedTuple`` consisting\n",
+      " |      of attributes ``missing_keys`` and ``unexpected_keys``. ``missing_keys``\n",
+      " |      is a ``list`` of ``str`` containing the missing keys and\n",
+      " |      ``unexpected_keys`` is a ``list`` of ``str`` containing the unexpected keys.\n",
+      " |      \n",
+      " |      The given incompatible_keys can be modified inplace if needed.\n",
+      " |      \n",
+      " |      Note that the checks performed when calling :func:`load_state_dict` with\n",
+      " |      ``strict=True`` are affected by modifications the hook makes to\n",
+      " |      ``missing_keys`` or ``unexpected_keys``, as expected. Additions to either\n",
+      " |      set of keys will result in an error being thrown when ``strict=True``, and\n",
+      " |      clearing out both missing and unexpected keys will avoid an error.\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
+      " |              a handle that can be used to remove the added hook by calling\n",
+      " |              ``handle.remove()``\n",
+      " |  \n",
+      " |  register_module(self, name: str, module: Optional[ForwardRef('Module')]) -> None\n",
+      " |      Alias for :func:`add_module`.\n",
+      " |  \n",
+      " |  register_parameter(self, name: str, param: Optional[torch.nn.parameter.Parameter]) -> None\n",
+      " |      Adds a parameter to the module.\n",
+      " |      \n",
+      " |      The parameter can be accessed as an attribute using given name.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          name (str): name of the parameter. The parameter can be accessed\n",
+      " |              from this module using the given name\n",
+      " |          param (Parameter or None): parameter to be added to the module. If\n",
+      " |              ``None``, then operations that run on parameters, such as :attr:`cuda`,\n",
+      " |              are ignored. If ``None``, the parameter is **not** included in the\n",
+      " |              module's :attr:`state_dict`.\n",
+      " |  \n",
+      " |  register_state_dict_pre_hook(self, hook)\n",
+      " |      These hooks will be called with arguments: ``self``, ``prefix``,\n",
+      " |      and ``keep_vars`` before calling ``state_dict`` on ``self``. The registered\n",
+      " |      hooks can be used to perform pre-processing before the ``state_dict``\n",
+      " |      call is made.\n",
+      " |  \n",
+      " |  requires_grad_(self: ~T, requires_grad: bool = True) -> ~T\n",
+      " |      Change if autograd should record operations on parameters in this\n",
+      " |      module.\n",
+      " |      \n",
+      " |      This method sets the parameters' :attr:`requires_grad` attributes\n",
+      " |      in-place.\n",
+      " |      \n",
+      " |      This method is helpful for freezing part of the module for finetuning\n",
+      " |      or training parts of a model individually (e.g., GAN training).\n",
+      " |      \n",
+      " |      See :ref:`locally-disable-grad-doc` for a comparison between\n",
+      " |      `.requires_grad_()` and several similar mechanisms that may be confused with it.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          requires_grad (bool): whether autograd should record operations on\n",
+      " |                                parameters in this module. Default: ``True``.\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          Module: self\n",
+      " |  \n",
+      " |  set_extra_state(self, state: Any)\n",
+      " |      This function is called from :func:`load_state_dict` to handle any extra state\n",
+      " |      found within the `state_dict`. Implement this function and a corresponding\n",
+      " |      :func:`get_extra_state` for your module if you need to store extra state within its\n",
+      " |      `state_dict`.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          state (dict): Extra state from the `state_dict`\n",
+      " |  \n",
+      " |  share_memory(self: ~T) -> ~T\n",
+      " |      See :meth:`torch.Tensor.share_memory_`\n",
+      " |  \n",
+      " |  state_dict(self, *args, destination=None, prefix='', keep_vars=False)\n",
+      " |      Returns a dictionary containing references to the whole state of the module.\n",
+      " |      \n",
+      " |      Both parameters and persistent buffers (e.g. running averages) are\n",
+      " |      included. Keys are corresponding parameter and buffer names.\n",
+      " |      Parameters and buffers set to ``None`` are not included.\n",
+      " |      \n",
+      " |      .. note::\n",
+      " |          The returned object is a shallow copy. It contains references\n",
+      " |          to the module's parameters and buffers.\n",
+      " |      \n",
+      " |      .. warning::\n",
+      " |          Currently ``state_dict()`` also accepts positional arguments for\n",
+      " |          ``destination``, ``prefix`` and ``keep_vars`` in order. However,\n",
+      " |          this is being deprecated and keyword arguments will be enforced in\n",
+      " |          future releases.\n",
+      " |      \n",
+      " |      .. warning::\n",
+      " |          Please avoid the use of argument ``destination`` as it is not\n",
+      " |          designed for end-users.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          destination (dict, optional): If provided, the state of module will\n",
+      " |              be updated into the dict and the same object is returned.\n",
+      " |              Otherwise, an ``OrderedDict`` will be created and returned.\n",
+      " |              Default: ``None``.\n",
+      " |          prefix (str, optional): a prefix added to parameter and buffer\n",
+      " |              names to compose the keys in state_dict. Default: ``''``.\n",
+      " |          keep_vars (bool, optional): by default the :class:`~torch.Tensor` s\n",
+      " |              returned in the state dict are detached from autograd. If it's\n",
+      " |              set to ``True``, detaching will not be performed.\n",
+      " |              Default: ``False``.\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          dict:\n",
+      " |              a dictionary containing a whole state of the module\n",
+      " |      \n",
+      " |      Example::\n",
+      " |      \n",
+      " |          >>> # xdoctest: +SKIP(\"undefined vars\")\n",
+      " |          >>> module.state_dict().keys()\n",
+      " |          ['bias', 'weight']\n",
+      " |  \n",
+      " |  to(self, *args, **kwargs)\n",
+      " |      Moves and/or casts the parameters and buffers.\n",
+      " |      \n",
+      " |      This can be called as\n",
+      " |      \n",
+      " |      .. function:: to(device=None, dtype=None, non_blocking=False)\n",
+      " |         :noindex:\n",
+      " |      \n",
+      " |      .. function:: to(dtype, non_blocking=False)\n",
+      " |         :noindex:\n",
+      " |      \n",
+      " |      .. function:: to(tensor, non_blocking=False)\n",
+      " |         :noindex:\n",
+      " |      \n",
+      " |      .. function:: to(memory_format=torch.channels_last)\n",
+      " |         :noindex:\n",
+      " |      \n",
+      " |      Its signature is similar to :meth:`torch.Tensor.to`, but only accepts\n",
+      " |      floating point or complex :attr:`dtype`\\ s. In addition, this method will\n",
+      " |      only cast the floating point or complex parameters and buffers to :attr:`dtype`\n",
+      " |      (if given). The integral parameters and buffers will be moved\n",
+      " |      :attr:`device`, if that is given, but with dtypes unchanged. When\n",
+      " |      :attr:`non_blocking` is set, it tries to convert/move asynchronously\n",
+      " |      with respect to the host if possible, e.g., moving CPU Tensors with\n",
+      " |      pinned memory to CUDA devices.\n",
+      " |      \n",
+      " |      See below for examples.\n",
+      " |      \n",
+      " |      .. note::\n",
+      " |          This method modifies the module in-place.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          device (:class:`torch.device`): the desired device of the parameters\n",
+      " |              and buffers in this module\n",
+      " |          dtype (:class:`torch.dtype`): the desired floating point or complex dtype of\n",
+      " |              the parameters and buffers in this module\n",
+      " |          tensor (torch.Tensor): Tensor whose dtype and device are the desired\n",
+      " |              dtype and device for all parameters and buffers in this module\n",
+      " |          memory_format (:class:`torch.memory_format`): the desired memory\n",
+      " |              format for 4D parameters and buffers in this module (keyword\n",
+      " |              only argument)\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          Module: self\n",
+      " |      \n",
+      " |      Examples::\n",
+      " |      \n",
+      " |          >>> # xdoctest: +IGNORE_WANT(\"non-deterministic\")\n",
+      " |          >>> linear = nn.Linear(2, 2)\n",
+      " |          >>> linear.weight\n",
+      " |          Parameter containing:\n",
+      " |          tensor([[ 0.1913, -0.3420],\n",
+      " |                  [-0.5113, -0.2325]])\n",
+      " |          >>> linear.to(torch.double)\n",
+      " |          Linear(in_features=2, out_features=2, bias=True)\n",
+      " |          >>> linear.weight\n",
+      " |          Parameter containing:\n",
+      " |          tensor([[ 0.1913, -0.3420],\n",
+      " |                  [-0.5113, -0.2325]], dtype=torch.float64)\n",
+      " |          >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)\n",
+      " |          >>> gpu1 = torch.device(\"cuda:1\")\n",
+      " |          >>> linear.to(gpu1, dtype=torch.half, non_blocking=True)\n",
+      " |          Linear(in_features=2, out_features=2, bias=True)\n",
+      " |          >>> linear.weight\n",
+      " |          Parameter containing:\n",
+      " |          tensor([[ 0.1914, -0.3420],\n",
+      " |                  [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')\n",
+      " |          >>> cpu = torch.device(\"cpu\")\n",
+      " |          >>> linear.to(cpu)\n",
+      " |          Linear(in_features=2, out_features=2, bias=True)\n",
+      " |          >>> linear.weight\n",
+      " |          Parameter containing:\n",
+      " |          tensor([[ 0.1914, -0.3420],\n",
+      " |                  [-0.5112, -0.2324]], dtype=torch.float16)\n",
+      " |      \n",
+      " |          >>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)\n",
+      " |          >>> linear.weight\n",
+      " |          Parameter containing:\n",
+      " |          tensor([[ 0.3741+0.j,  0.2382+0.j],\n",
+      " |                  [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)\n",
+      " |          >>> linear(torch.ones(3, 2, dtype=torch.cdouble))\n",
+      " |          tensor([[0.6122+0.j, 0.1150+0.j],\n",
+      " |                  [0.6122+0.j, 0.1150+0.j],\n",
+      " |                  [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)\n",
+      " |  \n",
+      " |  to_empty(self: ~T, *, device: Union[str, torch.device]) -> ~T\n",
+      " |      Moves the parameters and buffers to the specified device without copying storage.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          device (:class:`torch.device`): The desired device of the parameters\n",
+      " |              and buffers in this module.\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          Module: self\n",
+      " |  \n",
+      " |  train(self: ~T, mode: bool = True) -> ~T\n",
+      " |      Sets the module in training mode.\n",
+      " |      \n",
+      " |      This has any effect only on certain modules. See documentations of\n",
+      " |      particular modules for details of their behaviors in training/evaluation\n",
+      " |      mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,\n",
+      " |      etc.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          mode (bool): whether to set training mode (``True``) or evaluation\n",
+      " |                       mode (``False``). Default: ``True``.\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          Module: self\n",
+      " |  \n",
+      " |  type(self: ~T, dst_type: Union[torch.dtype, str]) -> ~T\n",
+      " |      Casts all parameters and buffers to :attr:`dst_type`.\n",
+      " |      \n",
+      " |      .. note::\n",
+      " |          This method modifies the module in-place.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          dst_type (type or string): the desired type\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          Module: self\n",
+      " |  \n",
+      " |  xpu(self: ~T, device: Union[int, torch.device, NoneType] = None) -> ~T\n",
+      " |      Moves all model parameters and buffers to the XPU.\n",
+      " |      \n",
+      " |      This also makes associated parameters and buffers different objects. So\n",
+      " |      it should be called before constructing optimizer if the module will\n",
+      " |      live on XPU while being optimized.\n",
+      " |      \n",
+      " |      .. note::\n",
+      " |          This method modifies the module in-place.\n",
+      " |      \n",
+      " |      Arguments:\n",
+      " |          device (int, optional): if specified, all parameters will be\n",
+      " |              copied to that device\n",
+      " |      \n",
+      " |      Returns:\n",
+      " |          Module: self\n",
+      " |  \n",
+      " |  zero_grad(self, set_to_none: bool = True) -> None\n",
+      " |      Sets gradients of all model parameters to zero. See similar function\n",
+      " |      under :class:`torch.optim.Optimizer` for more context.\n",
+      " |      \n",
+      " |      Args:\n",
+      " |          set_to_none (bool): instead of setting to zero, set the grads to None.\n",
+      " |              See :meth:`torch.optim.Optimizer.zero_grad` for details.\n",
+      " |  \n",
+      " |  ----------------------------------------------------------------------\n",
+      " |  Data descriptors inherited from torch.nn.modules.module.Module:\n",
+      " |  \n",
+      " |  __dict__\n",
+      " |      dictionary for instance variables (if defined)\n",
+      " |  \n",
+      " |  __weakref__\n",
+      " |      list of weak references to the object (if defined)\n",
+      " |  \n",
+      " |  ----------------------------------------------------------------------\n",
+      " |  Data and other attributes inherited from torch.nn.modules.module.Module:\n",
+      " |  \n",
+      " |  T_destination = ~T_destination\n",
+      " |  \n",
+      " |  call_super_init = False\n",
+      " |  \n",
+      " |  dump_patches = False\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "help(eva02_model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2f5ac1a7-6f1b-4417-8a67-1b2e32d385dd",
+   "metadata": {},
+   "source": [
+    "# DETR"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "id": "5c3ade1b-18ea-4368-abd9-53be1fdfb610",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2023-08-28 01:51:14,033] [INFO] [real_accelerator.py:133:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "The `max_size` parameter is deprecated and will be removed in v4.26. Please specify in `size['longest_edge'] instead`.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from transformers import DetrImageProcessor, DetrForObjectDetection\n",
+    "import torch\n",
+    "from PIL import Image\n",
+    "import requests\n",
+    "\n",
+    "url = \"http://images.cocodataset.org/val2017/000000039769.jpg\"\n",
+    "image = Image.open(requests.get(url, stream=True).raw)\n",
+    "\n",
+    "processor = DetrImageProcessor.from_pretrained(\"facebook/detr-resnet-50\", cache_dir='/fsx/proj-fmri/shared/cache')\n",
+    "model = DetrForObjectDetection.from_pretrained(\"facebook/detr-resnet-50\", cache_dir='/fsx/proj-fmri/shared/cache')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 34,
+   "id": "1d5aa2d7-4868-4751-8d90-7c52be028cd9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "inputs = processor(images=image, return_tensors=\"pt\")\n",
+    "outputs = model(**inputs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 35,
+   "id": "ae6bafc6-cee4-4e59-b7ba-12efc2a65b74",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Detected remote with confidence 0.998 at location [40.16, 70.81, 175.55, 117.98]\n",
+      "Detected remote with confidence 0.996 at location [333.24, 72.55, 368.33, 187.66]\n",
+      "Detected couch with confidence 0.995 at location [-0.02, 1.15, 639.73, 473.76]\n",
+      "Detected cat with confidence 0.999 at location [13.24, 52.05, 314.02, 470.93]\n",
+      "Detected cat with confidence 0.999 at location [345.4, 23.85, 640.37, 368.72]\n"
+     ]
+    }
+   ],
+   "source": [
+    "# convert outputs (bounding boxes and class logits) to COCO API\n",
+    "# let's only keep detections with score > 0.9\n",
+    "target_sizes = torch.tensor([image.size[::-1]])\n",
+    "results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]\n",
+    "\n",
+    "for score, label, box in zip(results[\"scores\"], results[\"labels\"], results[\"boxes\"]):\n",
+    "    box = [round(i, 2) for i in box.tolist()]\n",
+    "    print(\n",
+    "            f\"Detected {model.config.id2label[label.item()]} with confidence \"\n",
+    "            f\"{round(score.item(), 3)} at location {box}\"\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "id": "6dcc5934-79d4-4062-8b32-e42b3ebcdc0f",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "DetrImageProcessor {\n",
+       "  \"do_normalize\": true,\n",
+       "  \"do_pad\": true,\n",
+       "  \"do_rescale\": true,\n",
+       "  \"do_resize\": true,\n",
+       "  \"feature_extractor_type\": \"DetrFeatureExtractor\",\n",
+       "  \"format\": \"coco_detection\",\n",
+       "  \"image_mean\": [\n",
+       "    0.485,\n",
+       "    0.456,\n",
+       "    0.406\n",
+       "  ],\n",
+       "  \"image_processor_type\": \"DetrImageProcessor\",\n",
+       "  \"image_std\": [\n",
+       "    0.229,\n",
+       "    0.224,\n",
+       "    0.225\n",
+       "  ],\n",
+       "  \"resample\": 2,\n",
+       "  \"rescale_factor\": 0.00392156862745098,\n",
+       "  \"size\": {\n",
+       "    \"longest_edge\": 1333,\n",
+       "    \"shortest_edge\": 800\n",
+       "  }\n",
+       "}"
+      ]
+     },
+     "execution_count": 36,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "processor"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "db1d89cc-b432-473e-af69-d81c435ac731",
+   "metadata": {},
+   "source": [
+    "# CLIPSeg"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "id": "15db14d1-ee4d-4429-9286-054c4498293b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation\n",
+    "\n",
+    "processor = CLIPSegProcessor.from_pretrained(\"CIDAS/clipseg-rd16\",cache_dir='/fsx/proj-fmri/shared/cache')\n",
+    "model = CLIPSegForImageSegmentation.from_pretrained(\"CIDAS/clipseg-rd16\",cache_dir='/fsx/proj-fmri/shared/cache')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 38,
+   "id": "4aa225d4-5a3b-4dbb-ae57-dea2872ff492",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "AttributeError",
+     "evalue": "'JpegImageFile' object has no attribute 'shape'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[38], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mimage\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mshape\u001b[49m\n",
+      "\u001b[0;31mAttributeError\u001b[0m: 'JpegImageFile' object has no attribute 'shape'"
+     ]
+    }
+   ],
+   "source": [
+    "image.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ad7e2daf-0c7c-4fec-b29e-9ba47a037c6b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from PIL import Image\n",
+    "import requests\n",
+    "import h5py\n",
+    "\n",
+    "# url = \"https://unsplash.com/photos/8Nc_oQsc2qQ/download?ixid=MnwxMjA3fDB8MXxhbGx8fHx8fHx8fHwxNjcxMjAwNzI0&force=true&w=640\"\n",
+    "# image = Image.open(requests.get(url, stream=True).raw)\n",
+    "\n",
+    "image_path = \"/fsx/proj-fmri/shared/mindeyev2_dataset/coco_images_224_float16.hdf5\"\n",
+    "with h5py.File(image_path, 'r') as file:\n",
+    "    image = file['images'][0]\n",
+    "image = np.moveaxis(image, 0, -1).astype(np.float32)\n",
+    "plt.imshow(image)\n",
+    "\n",
+    "prompts = [\"person\",\"animal\",\"object\",\"background\"]\n",
+    "import torch\n",
+    "\n",
+    "# Rescale to [0, 255]\n",
+    "array = (image * 255).astype(np.uint8)\n",
+    "\n",
+    "# Convert to PIL image\n",
+    "image = Image.fromarray(array)\n",
+    "\n",
+    "inputs = processor(text=prompts, images=[image] * len(prompts), padding=\"max_length\", return_tensors=\"pt\")\n",
+    "# predict\n",
+    "with torch.no_grad():\n",
+    "  outputs = model(**inputs)\n",
+    "preds = outputs.logits.unsqueeze(1)\n",
+    "print(preds.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "131eb5b7-2f16-4a79-8402-edc1a1d8c348",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "preds = ((preds[0] + preds[1] + preds[2] + preds[-1].max() - preds[-1]) / 4)[None]\n",
+    "preds.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e2bf99e7-064d-4c22-997f-aa1a35dbab82",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "_, ax = plt.subplots(1, len(prompts) + 1, figsize=(3*(len(prompts) + 1), 4))\n",
+    "[a.axis('off') for a in ax.flatten()]\n",
+    "ax[0].imshow(image)\n",
+    "[ax[i+1].imshow(torch.sigmoid(preds[i][0])) for i in range(1)];\n",
+    "# [ax[i+1].text(0, -15, prompt) for i, prompt in enumerate(prompts)];"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b58b926f-a2b2-423b-b367-18808cf6b4f7",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/deepspeed_config_stage1.json

@@ -0,0 +1 @@
+{"bf16": {"enabled": false}, "fp16": {"enabled": true}, "zero_optimization": {"stage": 1, "contiguous_gradients": true, "stage3_gather_16bit_weights_on_model_save": true, "stage3_max_live_parameters": 1000000000.0, "stage3_max_reuse_distance": 1000000000.0, "stage3_prefetch_bucket_size": 10000000.0, "stage3_param_persistence_threshold": 100000.0, "reduce_bucket_size": 10000000.0, "sub_group_size": 1000000000.0, "offload_optimizer": {"device": "none", "nvme_path": "/scratch", "pin_memory": true}, "offload_param": {"device": "none", "nvme_path": "/scratch", "buffer_size": 4000000000.0, "pin_memory": true}}, "aio": {"block_size": 26214400, "queue_depth": 32, "thread_count": 1, "single_submit": false, "overlap_events": true}, "gradient_accumulation_steps": 1, "gradient_clipping": 1.0, "steps_per_print": 20000, "train_batch_size": 8, "train_micro_batch_size_per_gpu": 8, "wall_clock_breakdown": false, "zero_allow_untested_optimizer": true}

src/deepspeed_config_stage2.json

@@ -0,0 +1 @@
+{"bf16": {"enabled": false}, "fp16": {"enabled": true}, "zero_optimization": {"stage": 2, "contiguous_gradients": true, "stage3_gather_16bit_weights_on_model_save": true, "stage3_max_live_parameters": 1000000000.0, "stage3_max_reuse_distance": 1000000000.0, "stage3_prefetch_bucket_size": 10000000.0, "stage3_param_persistence_threshold": 100000.0, "reduce_bucket_size": 10000000.0, "sub_group_size": 1000000000.0, "offload_optimizer": {"device": "cpu", "nvme_path": "/scratch", "pin_memory": true}, "offload_param": {"device": "none", "nvme_path": "/scratch", "buffer_size": 4000000000.0, "pin_memory": true}}, "aio": {"block_size": 26214400, "queue_depth": 32, "thread_count": 1, "single_submit": false, "overlap_events": true}, "gradient_accumulation_steps": 1, "gradient_clipping": 1.0, "steps_per_print": 20000, "train_batch_size": 256, "train_micro_batch_size_per_gpu": 32, "wall_clock_breakdown": false, "zero_allow_untested_optimizer": true}

src/deepspeed_config_stage2_cpuoffload.json

@@ -0,0 +1,44 @@
+{
+    "bf16": {
+        "enabled": false
+    },
+    "fp16": {
+        "enabled": true
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "contiguous_gradients": true,
+        "stage3_gather_16bit_weights_on_model_save": true,
+        "stage3_max_live_parameters": 1000000000.0,
+        "stage3_max_reuse_distance": 1000000000.0,
+        "stage3_prefetch_bucket_size": 10000000.0,
+        "stage3_param_persistence_threshold": 100000.0,
+        "reduce_bucket_size": 10000000.0,
+        "sub_group_size": 1000000000.0,
+        "offload_optimizer": {
+            "device": "cpu",
+            "nvme_path": "/scratch",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "none",
+            "nvme_path": "/scratch",
+            "buffer_size": 4000000000.0,
+            "pin_memory": true
+        }
+    },
+    "aio": {
+        "block_size": 26214400,
+        "queue_depth": 32,
+        "thread_count": 1,
+        "single_submit": false,
+        "overlap_events": true
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 1.0,
+    "steps_per_print": 20000,
+    "train_batch_size": 256,
+    "train_micro_batch_size_per_gpu": 32,
+    "wall_clock_breakdown": false,
+    "zero_allow_untested_optimizer": true
+}

src/deepspeed_config_stage3.json

@@ -0,0 +1 @@
+{"fp16": {"enabled": true}, "optimizer": {"type": "AdamW", "params": {"lr": "auto", "weight_decay": "auto"}}, "scheduler": {"type": "WarmupDecayLR", "params": {"warmup_min_lr": "auto", "warmup_max_lr": "auto", "warmup_num_steps": "auto", "total_num_steps": "auto"}}, "zero_optimization": {"stage": 3, "offload_optimizer": {"device": "cpu", "pin_memory": true}, "offload_param": {"device": "cpu", "pin_memory": true}, "overlap_comm": true, "contiguous_gradients": true, "reduce_bucket_size": "auto", "stage3_prefetch_bucket_size": "auto", "stage3_param_persistence_threshold": "auto", "sub_group_size": 1000000000.0, "stage3_max_live_parameters": 1000000000.0, "stage3_max_reuse_distance": 1000000000.0, "stage3_gather_16bit_weights_on_model_save": "auto"}, "gradient_accumulation_steps": 1, "gradient_clipping": "auto", "steps_per_print": 2000, "train_batch_size": 256, "train_micro_batch_size_per_gpu": 32, "wall_clock_breakdown": false, "bf16": {"enabled": false}}

src/huggingface_to_s3.ipynb

@@ -0,0 +1,422 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "cf698d59-1cc2-4859-9c43-9a5d4d924ee1",
+   "metadata": {},
+   "source": [
+    "# Transfer huggingface mindeyev2 dataset to Stability aws s3"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "94c7404c-7a0f-4508-a630-954bc9af11fa",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/shared1000.npy -O /fsx/proj-fmri/shared/mindeyev2_dataset/shared1000.npy\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/betas_all_subj01.hdf5 -O /fsx/proj-fmri/shared/mindeyev2_dataset/betas_all_subj01.hdf5\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/betas_all_subj02.hdf5 -O /fsx/proj-fmri/shared/mindeyev2_dataset/betas_all_subj02.hdf5\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/betas_all_subj03.hdf5 -O /fsx/proj-fmri/shared/mindeyev2_dataset/betas_all_subj03.hdf5\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/betas_all_subj04.hdf5 -O /fsx/proj-fmri/shared/mindeyev2_dataset/betas_all_subj04.hdf5\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/betas_all_subj05.hdf5 -O /fsx/proj-fmri/shared/mindeyev2_dataset/betas_all_subj05.hdf5\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/betas_all_subj06.hdf5 -O /fsx/proj-fmri/shared/mindeyev2_dataset/betas_all_subj06.hdf5\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/betas_all_subj07.hdf5 -O /fsx/proj-fmri/shared/mindeyev2_dataset/betas_all_subj07.hdf5\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/betas_all_subj08.hdf5 -O /fsx/proj-fmri/shared/mindeyev2_dataset/betas_all_subj08.hdf5\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/coco_images_224_float16.hdf5 -O /fsx/proj-fmri/shared/mindeyev2_dataset/coco_images_224_float16.hdf5\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/COCO_73k_subj_indices.hdf5 -O /fsx/proj-fmri/shared/mindeyev2_dataset/COCO_73k_subj_indices.hdf5\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/0.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/1.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/1.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/2.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/2.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/3.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/3.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/4.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/4.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/5.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/5.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/6.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/6.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/7.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/7.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/8.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/8.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/9.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/9.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/10.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/10.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/11.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/11.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/12.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/12.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/13.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/13.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/14.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/14.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/15.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/15.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/16.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/16.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/17.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/17.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/18.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/18.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/19.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/19.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/20.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/20.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/21.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/21.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/22.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/22.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/23.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/23.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/24.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/24.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/25.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/25.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/26.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/26.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/27.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/27.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/28.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/28.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/29.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/29.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/30.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/30.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/31.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/31.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/32.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/32.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/33.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/33.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/34.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/34.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/35.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/35.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/train/36.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/36.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj01/test/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/test/0.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/0.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/1.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/1.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/2.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/2.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/3.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/3.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/4.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/4.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/5.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/5.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/6.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/6.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/7.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/7.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/8.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/8.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/9.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/9.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/10.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/10.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/11.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/11.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/12.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/12.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/13.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/13.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/14.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/14.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/15.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/15.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/16.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/16.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/17.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/17.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/18.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/18.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/19.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/19.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/20.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/20.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/21.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/21.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/22.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/22.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/23.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/23.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/24.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/24.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/25.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/25.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/26.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/26.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/27.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/27.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/28.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/28.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/29.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/29.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/30.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/30.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/31.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/31.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/32.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/32.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/33.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/33.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/34.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/34.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/35.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/35.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/train/36.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/train/36.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj02/test/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj02/test/0.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/0.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/1.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/1.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/2.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/2.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/3.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/3.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/4.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/4.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/5.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/5.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/6.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/6.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/7.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/7.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/8.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/8.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/9.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/9.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/10.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/10.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/11.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/11.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/12.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/12.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/13.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/13.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/14.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/14.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/15.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/15.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/16.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/16.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/17.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/17.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/18.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/18.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/19.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/19.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/20.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/20.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/21.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/21.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/22.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/22.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/23.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/23.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/24.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/24.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/25.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/25.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/26.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/26.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/27.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/27.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/28.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/28.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/29.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/29.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/30.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/30.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/31.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/31.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/32.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/32.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/33.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/33.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/34.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/34.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/35.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/35.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/train/36.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/train/36.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj03/test/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj03/test/0.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/0.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/1.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/1.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/2.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/2.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/3.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/3.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/4.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/4.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/5.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/5.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/6.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/6.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/7.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/7.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/8.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/8.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/9.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/9.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/10.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/10.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/11.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/11.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/12.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/12.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/13.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/13.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/14.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/14.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/15.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/15.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/16.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/16.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/17.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/17.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/18.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/18.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/19.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/19.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/20.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/20.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/21.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/21.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/22.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/22.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/23.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/23.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/24.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/24.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/25.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/25.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/26.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/26.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/27.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/27.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/28.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/28.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/29.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/29.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/30.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/30.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/31.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/31.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/32.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/32.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/33.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/33.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/34.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/34.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/35.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/35.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/train/36.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/train/36.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj04/test/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj04/test/0.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/0.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/1.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/1.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/2.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/2.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/3.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/3.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/4.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/4.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/5.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/5.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/6.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/6.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/7.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/7.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/8.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/8.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/9.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/9.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/10.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/10.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/11.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/11.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/12.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/12.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/13.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/13.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/14.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/14.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/15.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/15.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/16.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/16.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/17.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/17.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/18.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/18.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/19.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/19.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/20.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/20.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/21.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/21.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/22.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/22.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/23.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/23.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/24.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/24.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/25.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/25.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/26.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/26.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/27.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/27.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/28.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/28.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/29.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/29.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/30.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/30.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/31.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/31.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/32.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/32.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/33.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/33.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/34.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/34.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/35.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/35.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/train/36.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/train/36.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj05/test/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj05/test/0.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/0.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/1.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/1.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/2.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/2.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/3.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/3.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/4.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/4.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/5.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/5.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/6.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/6.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/7.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/7.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/8.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/8.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/9.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/9.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/10.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/10.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/11.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/11.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/12.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/12.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/13.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/13.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/14.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/14.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/15.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/15.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/16.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/16.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/17.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/17.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/18.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/18.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/19.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/19.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/20.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/20.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/21.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/21.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/22.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/22.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/23.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/23.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/24.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/24.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/25.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/25.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/26.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/26.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/27.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/27.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/28.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/28.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/29.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/29.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/30.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/30.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/31.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/31.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/32.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/32.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/33.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/33.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/34.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/34.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/35.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/35.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/train/36.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/train/36.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj06/test/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj06/test/0.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/0.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/1.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/1.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/2.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/2.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/3.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/3.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/4.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/4.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/5.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/5.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/6.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/6.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/7.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/7.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/8.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/8.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/9.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/9.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/10.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/10.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/11.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/11.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/12.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/12.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/13.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/13.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/14.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/14.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/15.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/15.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/16.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/16.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/17.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/17.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/18.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/18.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/19.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/19.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/20.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/20.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/21.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/21.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/22.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/22.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/23.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/23.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/24.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/24.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/25.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/25.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/26.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/26.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/27.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/27.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/28.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/28.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/29.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/29.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/30.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/30.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/31.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/31.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/32.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/32.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/33.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/33.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/34.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/34.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/35.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/35.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/train/36.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/train/36.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj07/test/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj07/test/0.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/0.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/1.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/1.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/2.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/2.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/3.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/3.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/4.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/4.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/5.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/5.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/6.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/6.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/7.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/7.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/8.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/8.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/9.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/9.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/10.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/10.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/11.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/11.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/12.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/12.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/13.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/13.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/14.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/14.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/15.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/15.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/16.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/16.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/17.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/17.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/18.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/18.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/19.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/19.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/20.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/20.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/21.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/21.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/22.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/22.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/23.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/23.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/24.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/24.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/25.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/25.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/26.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/26.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/27.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/27.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/28.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/28.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/29.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/29.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/30.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/30.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/31.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/31.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/32.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/32.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/33.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/33.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/34.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/34.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/35.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/35.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/train/36.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/train/36.tar\n",
+      "wget --show-progress https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/wds/subj08/test/0.tar -O /fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj08/test/0.tar\n",
+      "aws s3 sync /scratch/mindeyev2_dataset s3://proj-fmri/mindeyev2_dataset --region us-west-2\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "# from subprocess import call\n",
+    "# PS Note: it's faster to print the wget statements and then manually copy paste all them into terminal than to use subprocess call()\n",
+    "tmp = '/fsx/proj-fmri/shared/mindeyev2_dataset/' #'/scratch/mindeyev2_dataset/'\n",
+    "\n",
+    "hf_base_link = 'https://huggingface.co/datasets/pscotti/mindeyev2/resolve/main/'\n",
+    "\n",
+    "os.makedirs(tmp,exist_ok=True)\n",
+    "\n",
+    "files = [\n",
+    "    \"shared1000.npy\",\n",
+    "    \"betas_all_subj01.hdf5\",\n",
+    "    \"betas_all_subj02.hdf5\",\n",
+    "    \"betas_all_subj03.hdf5\",\n",
+    "    \"betas_all_subj04.hdf5\",\n",
+    "    \"betas_all_subj05.hdf5\",\n",
+    "    \"betas_all_subj06.hdf5\",\n",
+    "    \"betas_all_subj07.hdf5\",\n",
+    "    \"betas_all_subj08.hdf5\",\n",
+    "    \"coco_images_224_float16.hdf5\",\n",
+    "    \"COCO_73k_subj_indices.hdf5\",\n",
+    "]\n",
+    "\n",
+    "for f in files:    \n",
+    "    command = f\"wget --show-progress {hf_base_link}{f} -O {tmp}{f}\"\n",
+    "    print(command)\n",
+    "    # call(command,shell=True)\n",
+    "\n",
+    "for sub in range(1,9):\n",
+    "    subject = f'subj0{sub}'\n",
+    "\n",
+    "    tmp_fol = f'{tmp}wds/{subject}/'\n",
+    "    os.makedirs(tmp_fol,exist_ok=True)\n",
+    "    os.makedirs(tmp_fol+'train',exist_ok=True)\n",
+    "    os.makedirs(tmp_fol+'test',exist_ok=True)\n",
+    "\n",
+    "    for i in range(37):\n",
+    "        link = f'train/{i}.tar'\n",
+    "        command = f\"wget --show-progress {hf_base_link}wds/{subject}/{link} -O {tmp}wds/{subject}/{link}\"\n",
+    "        print(command)\n",
+    "        # call(command,shell=True)\n",
+    "\n",
+    "    link = f'test/0.tar'\n",
+    "    command = f\"wget --show-progress {hf_base_link}wds/{subject}/{link} -O {tmp}wds/{subject}/{link}\"\n",
+    "    print(command)\n",
+    "    # call(command,shell=True)\n",
+    "\n",
+    "command = \"aws s3 sync /scratch/mindeyev2_dataset s3://proj-fmri/mindeyev2_dataset --region us-west-2\"\n",
+    "print(command)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "30966082-59c2-411c-9b2e-4f4e3f9eb0f3",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/models.py

@@ -0,0 +1,210 @@
+import os
+import numpy as np
+from torchvision import transforms
+import torch
+import torch.nn as nn
+import PIL
+import clip
+import open_clip
+from functools import partial
+import random
+import json
+
+# class BrainMLP(nn.Module):
+#     def __init__(self, out_dim=257*768, in_dim=15724, clip_size=768, h=4096):
+#         super().__init__()
+#         self.lin0 = nn.Sequential(
+#             nn.Linear(in_dim, h, bias=False),
+#             nn.LayerNorm(h),
+#             nn.GELU(inplace=True),
+#             nn.Dropout(0.5))
+#         self.mlp = nn.ModuleList([
+#             nn.Sequential(
+#                 nn.Linear(h, h),
+#                 nn.LayerNorm(h),
+#                 nn.GELU(inplace=True),
+#                 nn.Dropout(0.15)
+#             ) for _ in range(4)])
+#         self.lin1 = nn.Linear(h, out_dim, bias=True)
+#         self.proj = nn.Sequential(
+#             nn.LayerNorm(clip_size),
+#             nn.GELU(),
+#             nn.Linear(clip_size, 2048),
+#             nn.LayerNorm(2048),
+#             nn.GELU(),
+#             nn.Linear(2048, 2048),
+#             nn.LayerNorm(2048),
+#             nn.GELU(),
+#             nn.Linear(2048, clip_size))
+#     def forward(self, x):
+#         x = self.lin0(x)
+#         residual = x
+#         for res_block in range(self.n_blocks):
+#             x = self.mlp[res_block](x)
+#             x += residual
+#             residual = x
+#         diffusion_prior_input = self.lin1(x.reshape(len(x), -1))
+#         disjointed_clip_fmri = self.proj(diffusion_prior_input.reshape(
+#                                         len(x),-1, self.clip_size))
+#         return diffusion_prior_input, disjointed_clip_fmri
+
+
+
+class Clipper(torch.nn.Module):
+    def __init__(self, clip_variant, clamp_embs=False, norm_embs=False,
+                 hidden_state=False, device=torch.device('cpu')):
+        super().__init__()
+        assert clip_variant in ("RN50", "ViT-L/14", "ViT-B/32", "RN50x64"), \
+            "clip_variant must be one of RN50, ViT-L/14, ViT-B/32, RN50x64"
+        print(clip_variant, device)
+        
+        if clip_variant=="ViT-L/14" and hidden_state:
+            # from transformers import CLIPVisionModelWithProjection
+            # image_encoder = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-large-patch14",cache_dir="/fsx/proj-medarc/fmri/cache")
+            from transformers import CLIPVisionModelWithProjection
+            sd_cache_dir = '/fsx/proj-fmri/shared/cache/models--shi-labs--versatile-diffusion/snapshots/2926f8e11ea526b562cd592b099fcf9c2985d0b7'
+            image_encoder = CLIPVisionModelWithProjection.from_pretrained(sd_cache_dir, subfolder='image_encoder').eval()
+            image_encoder = image_encoder.to(device)
+            for param in image_encoder.parameters():
+                param.requires_grad = False # dont need to calculate gradients
+            self.image_encoder = image_encoder
+        elif hidden_state:
+            raise Exception("hidden_state embeddings only works with ViT-L/14 right now")
+        
+        clip_model, preprocess = clip.load(clip_variant, device=device)
+        clip_model.eval() # dont want to train model
+        for param in clip_model.parameters():
+            param.requires_grad = False # dont need to calculate gradients
+            
+        self.clip = clip_model
+        self.clip_variant = clip_variant
+        if clip_variant == "RN50x64":
+            self.clip_size = (448,448)
+        else:
+            self.clip_size = (224,224)
+            
+        preproc = transforms.Compose([
+            transforms.Resize(size=self.clip_size[0], interpolation=transforms.InterpolationMode.BICUBIC),
+            transforms.CenterCrop(size=self.clip_size),
+            transforms.Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711))
+        ])
+        self.preprocess = preproc
+        self.hidden_state = hidden_state
+        self.mean = np.array([0.48145466, 0.4578275, 0.40821073])
+        self.std = np.array([0.26862954, 0.26130258, 0.27577711])
+        self.normalize = transforms.Normalize(self.mean, self.std)
+        self.denormalize = transforms.Normalize((-self.mean / self.std).tolist(), (1.0 / self.std).tolist())
+        self.clamp_embs = clamp_embs
+        self.norm_embs = norm_embs
+        self.device= device
+        
+        def versatile_normalize_embeddings(encoder_output):
+            embeds = encoder_output.last_hidden_state
+            embeds = image_encoder.vision_model.post_layernorm(embeds)
+            embeds = image_encoder.visual_projection(embeds)
+            return embeds
+        self.versatile_normalize_embeddings = versatile_normalize_embeddings
+
+    def resize_image(self, image):
+        # note: antialias should be False if planning to use Pinkney's Image Variation SD model
+        return transforms.Resize(self.clip_size)(image.to(self.device))
+
+    def embed_image(self, image):
+        """Expects images in -1 to 1 range"""
+        if self.hidden_state:
+            # clip_emb = self.preprocess((image/1.5+.25).to(self.device)) # for some reason the /1.5+.25 prevents oversaturation
+            clip_emb = self.preprocess((image).to(self.device))
+            clip_emb = self.image_encoder(clip_emb)
+            clip_emb = self.versatile_normalize_embeddings(clip_emb)
+        else:
+            clip_emb = self.preprocess(image.to(self.device))
+            clip_emb = self.clip.encode_image(clip_emb)
+        # input is now in CLIP space, but mind-reader preprint further processes embeddings:
+        if self.clamp_embs:
+            clip_emb = torch.clamp(clip_emb, -1.5, 1.5)
+        if self.norm_embs:
+            if self.hidden_state:        
+                # normalize all tokens by cls token's norm
+                clip_emb = clip_emb / torch.norm(clip_emb[:, 0], dim=-1).reshape(-1, 1, 1)
+            else:
+                clip_emb = nn.functional.normalize(clip_emb, dim=-1)
+        return clip_emb
+
+    def embed_text(self, text_samples):
+        clip_text = clip.tokenize(text_samples).to(self.device)
+        clip_text = self.clip.encode_text(clip_text)
+        if self.clamp_embs:
+            clip_text = torch.clamp(clip_text, -1.5, 1.5)
+        if self.norm_embs:
+            clip_text = nn.functional.normalize(clip_text, dim=-1)
+        return clip_text
+
+    def embed_curated_annotations(self, annots):
+        for i,b in enumerate(annots):
+            t = ''
+            while t == '':
+                rand = torch.randint(5,(1,1))[0][0]
+                t = b[0,rand]
+            if i==0:
+                txt = np.array(t)
+            else:
+                txt = np.vstack((txt,t))
+        txt = txt.flatten()
+        return self.embed_text(txt)
+    
+class OpenClipper(torch.nn.Module):
+    def __init__(self, clip_variant, norm_embs=False, device=torch.device('cpu')):
+        super().__init__()
+        print(clip_variant, device)
+        assert clip_variant == 'ViT-H-14' # not setup for other models yet
+                
+        clip_model, _, preprocess = open_clip.create_model_and_transforms('ViT-H-14', 
+                                        pretrained='laion2b_s32b_b79k', device=device)
+        clip_model.eval() # dont want to train model
+        for param in clip_model.parameters():
+            param.requires_grad = False # dont need to calculate gradients
+            
+        # overwrite preprocess to accept torch inputs instead of PIL Image
+        preprocess = transforms.Compose([
+                transforms.Resize(224, interpolation=transforms.InterpolationMode.BICUBIC, antialias=None),
+                transforms.CenterCrop(224),
+                transforms.Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711))
+        ])
+        
+        tokenizer = open_clip.get_tokenizer('ViT-H-14')
+            
+        self.clip = clip_model
+        self.norm_embs = norm_embs
+        self.preprocess = preprocess
+        self.tokenizer = tokenizer
+        self.device = device
+        
+    def embed_image(self, image):
+        """Expects images in -1 to 1 range"""
+        image = self.preprocess(image).to(self.device)
+        with torch.no_grad(), torch.cuda.amp.autocast():
+            image_features = self.clip.encode_image(image)
+        if self.norm_embs:
+            image_features = nn.functional.normalize(image_features, dim=-1)
+        return image_features
+
+    def embed_text(self, text_samples):
+        text = self.tokenizer(text_samples).to(self.device)
+        with torch.no_grad(), torch.cuda.amp.autocast():
+            text_features = self.clip.encode_text(text)
+        if self.norm_embs:
+            text_features = nn.functional.normalize(text_features, dim=-1)
+        return text_features
+
+    def embed_curated_annotations(self, annots):
+        for i,b in enumerate(annots):
+            t = ''
+            while t == '':
+                rand = torch.randint(5,(1,1))[0][0]
+                t = b[0,rand]
+            if i==0:
+                txt = np.array(t)
+            else:
+                txt = np.vstack((txt,t))
+        txt = txt.flatten()
+        return self.embed_text(txt)

src/setup.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+# Commands to setup a new conda environment and install all the necessary packages
+# See the environment.yaml file for "conda env export > environment.yaml" after running this.
+
+set -e
+
+conda create -n fmri python=3.10.8 -y
+conda activate fmri
+
+conda install numpy matplotlib tqdm scikit-image jupyterlab -y
+
+pip install accelerate webdataset clip pandas matplotlib ftfy regex kornia umap-learn h5py
+pip install torchvision==0.15.2 torch==2.0.1
+pip install diffusers
+pip install deepspeed

src/train2-tryal.ipynb

@@ -0,0 +1,2409 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "f4d95fac-ac1d-473c-ab96-650f76e6aaf5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# # Code to convert this notebook to .py if you want to run it via command line or with Slurm\n",
+    "# from subprocess import call\n",
+    "# command = \"jupyter nbconvert Train.ipynb --to python\"\n",
+    "# call(command,shell=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b0f0f4f3",
+   "metadata": {},
+   "source": [
+    "# Import packages & functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "5bad764b-45c1-45ce-a716-8d055e09821a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2023-11-19 16:32:39,711] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import sys\n",
+    "import json\n",
+    "import argparse\n",
+    "import numpy as np\n",
+    "import math\n",
+    "from einops import rearrange\n",
+    "import time\n",
+    "import random\n",
+    "import h5py\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "import webdataset as wds\n",
+    "import gc\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torchvision import transforms\n",
+    "from torchvision.transforms import ToPILImage   #CHANGED (added)\n",
+    "\n",
+    "from accelerate import Accelerator, DeepSpeedPlugin\n",
+    "\n",
+    "# tf32 data type is faster than standard float32\n",
+    "torch.backends.cuda.matmul.allow_tf32 = True\n",
+    "\n",
+    "# custom functions #\n",
+    "import utils\n",
+    "\n",
+    "global_batch_size = 128 #128"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "cc5d2e32-6027-4a19-bef4-5ca068db35bb",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "LOCAL RANK  0\n"
+     ]
+    }
+   ],
+   "source": [
+    "### Multi-GPU config ###\n",
+    "local_rank = os.getenv('RANK')\n",
+    "if local_rank is None: \n",
+    "    local_rank = 0\n",
+    "else:\n",
+    "    local_rank = int(local_rank)\n",
+    "print(\"LOCAL RANK \", local_rank)  \n",
+    "\n",
+    "num_devices = torch.cuda.device_count()\n",
+    "if num_devices==0: num_devices = 1\n",
+    "\n",
+    "accelerator = Accelerator(split_batches=False)\n",
+    "\n",
+    "### UNCOMMENT BELOW STUFF TO USE DEEPSPEED (also comment out the immediately above \"accelerator = \" line) ###\n",
+    "\n",
+    "# if num_devices <= 1 and utils.is_interactive():\n",
+    "#     # can emulate a distributed environment for deepspeed to work in jupyter notebook\n",
+    "#     os.environ[\"MASTER_ADDR\"] = \"localhost\"\n",
+    "#     os.environ[\"MASTER_PORT\"] = str(np.random.randint(10000)+9000)\n",
+    "#     os.environ[\"RANK\"] = \"0\"\n",
+    "#     os.environ[\"LOCAL_RANK\"] = \"0\"\n",
+    "#     os.environ[\"WORLD_SIZE\"] = \"1\"\n",
+    "#     os.environ[\"GLOBAL_BATCH_SIZE\"] = str(global_batch_size) # set this to your batch size!\n",
+    "#     global_batch_size = os.environ[\"GLOBAL_BATCH_SIZE\"]\n",
+    "\n",
+    "# # alter the deepspeed config according to your global and local batch size\n",
+    "# if local_rank == 0:\n",
+    "#     with open('deepspeed_config_stage2.json', 'r') as file:\n",
+    "#         config = json.load(file)\n",
+    "#     config['train_batch_size'] = int(os.environ[\"GLOBAL_BATCH_SIZE\"])\n",
+    "#     config['train_micro_batch_size_per_gpu'] = int(os.environ[\"GLOBAL_BATCH_SIZE\"]) // num_devices\n",
+    "#     with open('deepspeed_config_stage2.json', 'w') as file:\n",
+    "#         json.dump(config, file)\n",
+    "# else:\n",
+    "#     # give some time for the local_rank=0 gpu to prep new deepspeed config file\n",
+    "#     time.sleep(10)\n",
+    "# deepspeed_plugin = DeepSpeedPlugin(\"deepspeed_config_stage2.json\")\n",
+    "# accelerator = Accelerator(split_batches=False, deepspeed_plugin=deepspeed_plugin)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "b767ab6f-d4a9-47a5-b3bf-f56bf6760c0c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "PID of this process = 2370606\n",
+      "device: cuda\n",
+      "Distributed environment: NO\n",
+      "Num processes: 1\n",
+      "Process index: 0\n",
+      "Local process index: 0\n",
+      "Device: cuda\n",
+      "\n",
+      "Mixed precision type: no\n",
+      "\n",
+      "distributed = False num_devices = 1 local rank = 0 world size = 1\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(\"PID of this process =\",os.getpid())\n",
+    "device = accelerator.device\n",
+    "print(\"device:\",device)\n",
+    "num_workers = num_devices\n",
+    "print(accelerator.state)\n",
+    "world_size = accelerator.state.num_processes\n",
+    "distributed = not accelerator.state.distributed_type == 'NO'\n",
+    "print(\"distributed =\",distributed, \"num_devices =\", num_devices, \"local rank =\", local_rank, \"world size =\", world_size)\n",
+    "print = accelerator.print # only print if local_rank=0"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9018b82b-c054-4463-9527-4b0c2a75bda6",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "# Configurations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "2b61fec7-72a0-4b67-86da-1375f1d9fbd3",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "['--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset', '--model_name=captions', '--subj=1', '--batch_size=128', '--n_samples_save=0', '--max_lr=3e-1', '--mixup_pct=.66', '--num_epochs=30', '--ckpt_interval=999', '--no-use_image_aug']\n"
+     ]
+    }
+   ],
+   "source": [
+    "# if running this interactively, can specify jupyter_args here for argparser to use\n",
+    "if utils.is_interactive():\n",
+    "    # Example use\n",
+    "    jupyter_args = f\"--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset \\\n",
+    "                    --model_name=captions \\\n",
+    "                    --subj=1 --batch_size={global_batch_size} --n_samples_save=0 \\\n",
+    "                    --max_lr=3e-1 --mixup_pct=.66 --num_epochs=30 --ckpt_interval=999 --no-use_image_aug\"\n",
+    "    #max_lr=3e-5 originally\n",
+    "    jupyter_args = jupyter_args.split()\n",
+    "    print(jupyter_args)\n",
+    "    \n",
+    "    from IPython.display import clear_output # function to clear print outputs in cell\n",
+    "    %load_ext autoreload \n",
+    "    # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions\n",
+    "    %autoreload 2 "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "2028bdf0-2f41-46d9-b6e7-86b870dbf16c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "global batch_size 128\n",
+      "batch_size 128\n"
+     ]
+    }
+   ],
+   "source": [
+    "parser = argparse.ArgumentParser(description=\"Model Training Configuration\")\n",
+    "parser.add_argument(\n",
+    "    \"--model_name\", type=str, default=\"testing\",\n",
+    "    help=\"name of model, used for ckpt saving and wandb logging (if enabled)\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--data_path\", type=str, default=\"/fsx/proj-fmri/shared/natural-scenes-dataset\",\n",
+    "    help=\"Path to where NSD data is stored / where to download it to\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--subj\",type=int, default=1, choices=[1,2,5,7],\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--batch_size\", type=int, default=32,\n",
+    "    help=\"Batch size can be increased by 10x if only training v2c and not diffusion diffuser\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--wandb_log\",action=argparse.BooleanOptionalAction,default=False,\n",
+    "    help=\"whether to log to wandb\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--resume_from_ckpt\",action=argparse.BooleanOptionalAction,default=False,\n",
+    "    help=\"if not using wandb and want to resume from a ckpt\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--wandb_project\",type=str,default=\"stability\",\n",
+    "    help=\"wandb project name\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--mixup_pct\",type=float,default=.33,\n",
+    "    help=\"proportion of way through training when to switch from BiMixCo to SoftCLIP\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--use_image_aug\",action=argparse.BooleanOptionalAction,default=True,\n",
+    "    help=\"whether to use image augmentation\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--num_epochs\",type=int,default=240,\n",
+    "    help=\"number of epochs of training\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--lr_scheduler_type\",type=str,default='cycle',choices=['cycle','linear'],\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--ckpt_saving\",action=argparse.BooleanOptionalAction,default=True,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--ckpt_interval\",type=int,default=5,\n",
+    "    help=\"save backup ckpt and reconstruct every x epochs\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--seed\",type=int,default=42,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--max_lr\",type=float,default=3e-4,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--n_samples_save\",type=int,default=0,choices=[0,1],\n",
+    "    help=\"Number of reconstructions for monitoring progress, 0 will speed up training\",\n",
+    ")\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    args = parser.parse_args(jupyter_args)\n",
+    "else:\n",
+    "    args = parser.parse_args()\n",
+    "\n",
+    "# create global variables without the args prefix\n",
+    "for attribute_name in vars(args).keys():\n",
+    "    globals()[attribute_name] = getattr(args, attribute_name)\n",
+    "\n",
+    "print(\"global batch_size\", batch_size)\n",
+    "batch_size = int(batch_size / num_devices)\n",
+    "print(\"batch_size\", batch_size)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "60cd7f2c-37fd-426b-a0c6-633e51bc4c4d",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "outdir = os.path.abspath(f'../train_logs/{model_name}')\n",
+    "if not os.path.exists(outdir):\n",
+    "    os.makedirs(outdir,exist_ok=True)\n",
+    "if use_image_aug:\n",
+    "    import kornia\n",
+    "    from kornia.augmentation.container import AugmentationSequential\n",
+    "    img_augment = AugmentationSequential(\n",
+    "        kornia.augmentation.RandomResizedCrop((224,224), (0.6,1), p=0.3),\n",
+    "        kornia.augmentation.Resize((224, 224)),\n",
+    "        kornia.augmentation.RandomHorizontalFlip(p=0.3),\n",
+    "        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.3),\n",
+    "        kornia.augmentation.RandomGrayscale(p=0.3),\n",
+    "        same_on_batch=False,\n",
+    "        data_keys=[\"input\"],\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "e7807ba9-02b6-4bc0-873c-69869abe4091",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "wandb_log = False"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "42d13c25-1369-4c49-81d4-83d713586096",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "# Prep data, models, and dataloaders"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1c023f24-5233-4a15-a2f5-78487b3a8546",
+   "metadata": {},
+   "source": [
+    "## Dataloader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "81084834-035f-4465-ad59-59e6b806a2f5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/{0..36}.tar\n",
+      "/fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/test/0.tar\n"
+     ]
+    }
+   ],
+   "source": [
+    "if subj==1:\n",
+    "    num_train = 24958\n",
+    "    num_test = 2770\n",
+    "test_batch_size = num_test\n",
+    "\n",
+    "def my_split_by_node(urls): return urls\n",
+    "    \n",
+    "train_url = f\"{data_path}/wds/subj0{subj}/train/\" + \"{0..36}.tar\"\n",
+    "print(train_url)\n",
+    "\n",
+    "train_data = wds.WebDataset(train_url,resampled=False,nodesplitter=my_split_by_node)\\\n",
+    "                    .shuffle(750, initial=1500, rng=random.Random(42))\\\n",
+    "                    .decode(\"torch\")\\\n",
+    "                    .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n",
+    "                    .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n",
+    "train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)\n",
+    "\n",
+    "test_url = f\"{data_path}/wds/subj0{subj}/test/\" + \"0.tar\"\n",
+    "print(test_url)\n",
+    "\n",
+    "test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\\\n",
+    "                    .shuffle(750, initial=1500, rng=random.Random(42))\\\n",
+    "                    .decode(\"torch\")\\\n",
+    "                    .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n",
+    "                    .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n",
+    "test_dl = torch.utils.data.DataLoader(test_data, batch_size=test_batch_size, shuffle=False, drop_last=False, pin_memory=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "203b060a-2dd2-4c35-929b-c576be82eb52",
+   "metadata": {},
+   "source": [
+    "### check dataloaders are working"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "e7a9c68c-c3c9-4080-bd99-067c4486dc37",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# test_indices = []\n",
+    "# test_images = []\n",
+    "# for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):\n",
+    "#     test_indices = np.append(test_indices, behav[:,0,5].numpy())\n",
+    "#     test_images = np.append(test_images, behav[:,0,0].numpy())\n",
+    "# test_indices = test_indices.astype(np.int16)\n",
+    "# print(test_i, (test_i+1) * test_batch_size, len(test_indices))\n",
+    "# print(\"---\\n\")\n",
+    "\n",
+    "# train_indices = []\n",
+    "# train_images = []\n",
+    "# for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):\n",
+    "#     train_indices = np.append(train_indices, behav[:,0,5].long().numpy())\n",
+    "#     train_images = np.append(train_images, behav[:,0,0].numpy())\n",
+    "# train_indices = train_indices.astype(np.int16)\n",
+    "# print(train_i, (train_i+1) * batch_size, len(train_indices))\n",
+    "\n",
+    "# # train_images = np.hstack((train_images, test_images))\n",
+    "# # print(\"WARNING: ADDED TEST IMAGES TO TRAIN IMAGES\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "45fad12c-f9fb-4408-8fd4-9bca324ad634",
+   "metadata": {},
+   "source": [
+    "## Load data and images"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "039dd330-7339-4f88-8f00-45f95e47baa0",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "subj01 betas loaded into memory\n",
+      "voxels torch.Size([27750, 15729])\n",
+      "images torch.Size([73000, 3, 224, 224])\n"
+     ]
+    }
+   ],
+   "source": [
+    "# load betas\n",
+    "f = h5py.File(f'{data_path}/betas_all_subj0{subj}.hdf5', 'r')\n",
+    "voxels = f['betas'][:]\n",
+    "print(f\"subj0{subj} betas loaded into memory\")\n",
+    "voxels = torch.Tensor(voxels).to(\"cpu\").half()\n",
+    "if subj==1:\n",
+    "    voxels = torch.hstack((voxels, torch.zeros((len(voxels), 5))))\n",
+    "print(\"voxels\", voxels.shape)\n",
+    "num_voxels = voxels.shape[-1]\n",
+    "\n",
+    "# load orig images\n",
+    "f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')\n",
+    "images = f['images'][:]\n",
+    "images = torch.Tensor(images).to(\"cpu\").half()\n",
+    "print(\"images\", images.shape)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "10ec4517-dbdf-4ece-98f6-4714d5de4e15",
+   "metadata": {},
+   "source": [
+    "## Load models"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "48d6160e-1ee8-4da7-a755-9dbb452a6fa5",
+   "metadata": {},
+   "source": [
+    "### CLIP image embeddings  model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "795e2885-bd07-4e27-bed7-181473c06df9",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import transformers\n",
+    "from transformers import Blip2Processor, Blip2ForConditionalGeneration\n",
+    "\n",
+    "from PIL import Image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "b0420dc0-199e-4c1a-857d-b1747058b467",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "ViT-L/14 cuda:0\n"
+     ]
+    }
+   ],
+   "source": [
+    "from models import Clipper\n",
+    "clip_model = Clipper(\"ViT-L/14\", device=torch.device(f\"cuda:{local_rank}\"), hidden_state=True, norm_embs=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "23428fb7-2955-4295-bea1-447cebf9f72e",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [01:08<00:00, 34.47s/it]\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'from lavis.models import load_model_and_preprocess\\nfrom lavis.models import model_zoo\\nblip2_model, vis_processors, _ = load_model_and_preprocess(\\n            name=\"blip2_t5\", model_type=\"pretrain_flant5xl_vitL\", is_eval=True, device=device)\\n\\nclip_seq_dim = 257\\nclip_emb_dim = 1024\\nhidden_dim = 4096'"
+      ]
+     },
+     "execution_count": 14,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "cache_blip2 = \"/fsx/proj-fmri/shared/cache/models--Salesforce--blip2-opt-2.7b/snapshots/6e723d92ee91ebcee4ba74d7017632f11ff4217b\"\n",
+    "\n",
+    "b2_processor = Blip2Processor.from_pretrained(cache_blip2)\n",
+    "b2_model = Blip2ForConditionalGeneration.from_pretrained(cache_blip2, torch_dtype=torch.float16, device_map=\"auto\")\n",
+    "\n",
+    "#Load in blip2 as well\n",
+    "\"\"\"from lavis.models import load_model_and_preprocess\n",
+    "from lavis.models import model_zoo\n",
+    "blip2_model, vis_processors, _ = load_model_and_preprocess(\n",
+    "            name=\"blip2_t5\", model_type=\"pretrain_flant5xl_vitL\", is_eval=True, device=device)\n",
+    "\n",
+    "clip_seq_dim = 257\n",
+    "clip_emb_dim = 1024\n",
+    "hidden_dim = 4096\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 74,
+   "id": "b06f3de2-a8da-4ba0-94f0-99096f738d55",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "def embed_images_b2(images):\n",
+    "    images = (images * 255).type(torch.uint8)\n",
+    "    with torch.no_grad():\n",
+    "        inputs_processed = b2_processor(images, return_tensors=\"pt\").to(\"cuda\", torch.float16)\n",
+    "        enc_imgs = b2_model.vision_model.forward(inputs_processed['pixel_values'])\n",
+    "        return enc_imgs.last_hidden_state.detach(), inputs_processed\n",
+    "\n",
+    "def embeds_to_captions_b2(embeds, sample = False, temp = 0.9):\n",
+    "    with torch.no_grad():\n",
+    "        input_ids = None #inputs['input_ids']\n",
+    "        attention_mask = None\n",
+    "        batch_size = embeds.shape[0]\n",
+    "        image_embeds = embeds\n",
+    "        image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device)\n",
+    "\n",
+    "        query_tokens = b2_model.query_tokens.expand(image_embeds.shape[0], -1, -1)\n",
+    "        query_outputs = b2_model.qformer(\n",
+    "            query_embeds=query_tokens,\n",
+    "            encoder_hidden_states=image_embeds,\n",
+    "            encoder_attention_mask=image_attention_mask,\n",
+    "            return_dict=True,\n",
+    "        )\n",
+    "        query_output = query_outputs.last_hidden_state\n",
+    "\n",
+    "        language_model_inputs = b2_model.language_projection(query_output)\n",
+    "        language_attention_mask = torch.ones(\n",
+    "            language_model_inputs.size()[:-1], dtype=torch.long, device=language_model_inputs.device\n",
+    "        )\n",
+    "        if input_ids is None:\n",
+    "            input_ids = (\n",
+    "                torch.LongTensor([[b2_model.config.text_config.bos_token_id]])\n",
+    "                .repeat(batch_size, 1)\n",
+    "                .to(image_embeds.device)\n",
+    "            )\n",
+    "        if attention_mask is None:\n",
+    "            attention_mask = torch.ones_like(input_ids)\n",
+    "        attention_mask = torch.cat([language_attention_mask, attention_mask.to(language_attention_mask.device)], dim=1)\n",
+    "\n",
+    "        # concatenate query embeddings with prompt embeddings\n",
+    "        inputs_embeds = b2_model.get_input_embeddings()(input_ids)\n",
+    "        inputs_embeds = torch.cat([language_model_inputs, inputs_embeds.to(language_model_inputs.device)], dim=1)\n",
+    "\n",
+    "        outputs = b2_model.language_model.generate(\n",
+    "            inputs_embeds=inputs_embeds,\n",
+    "            attention_mask=attention_mask,\n",
+    "            temperature=temp,\n",
+    "            do_sample = sample\n",
+    "        )\n",
+    "        text = b2_processor.batch_decode(outputs, skip_special_tokens=True)\n",
+    "        \n",
+    "        return outputs, text\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 73,
+   "id": "51b29638-2c81-4e9f-b06d-525fdbac44b1",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([[    2,  6209,    14,    10,   205,   425,    13,    10,  7297,  1280,\n",
+       "             9,   418,   116,  1437,    38, 10728,    33,   117,  1114,    99]],\n",
+       "       device='cuda:0')"
+      ]
+     },
+     "execution_count": 73,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "b2_model.language_model.generate(do_sample = True, temperature=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "ec0a34d3-76e0-4a47-a9ab-6131ab2ccecd",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "image_test = images[1:20].permute(0,2,3,1)\n",
+    "#raw_image = Image.open('/fsx/proj-fmri/shared/controlNetData/target/img_t1.jpg').convert('RGB')\n",
+    "# Convert the image to a NumPy array\n",
+    "#image_test = np.array(raw_image)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "e04876a4-45c7-4015-8255-8574c8f50f14",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "\"import matplotlib.pyplot as plt\\n# Plotting one of the images (taking the first image as an example)\\nimg_to_plot = inputs_rec['pixel_values'][-1]\\n\\n# Transpose the image for correct display (PyTorch: [C, H, W], Matplotlib: [H, W, C])\\nimg_to_plot = img_to_plot.permute(1, 2, 0).to(torch.float32).to('cpu')\\nprint(img_to_plot.shape)\\n\\nplt.imshow(img_to_plot)\\nplt.show()\""
+      ]
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "\"\"\"import matplotlib.pyplot as plt\n",
+    "# Plotting one of the images (taking the first image as an example)\n",
+    "img_to_plot = inputs_rec['pixel_values'][-1]\n",
+    "\n",
+    "# Transpose the image for correct display (PyTorch: [C, H, W], Matplotlib: [H, W, C])\n",
+    "img_to_plot = img_to_plot.permute(1, 2, 0).to(torch.float32).to('cpu')\n",
+    "print(img_to_plot.shape)\n",
+    "\n",
+    "plt.imshow(img_to_plot)\n",
+    "plt.show()\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "328a17d0-593b-4d1e-812a-10a3b6efea6a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "embeds_test, inputs_rec = embed_images_b2(image_test)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "abe5f8a8-fca9-4083-8596-a913bdb57de7",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "#inputs_rec['pixel_values'].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "c5f3ca7e-b880-421e-b354-7b6c3df565e9",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "#out = b2_model.generate(**inputs_rec)\n",
+    "#print(b2_processor.decode(out[0], skip_special_tokens=True).strip())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "fb462016-78d7-46ea-8058-0d608f17ea65",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/transformers/generation/utils.py:1260: UserWarning: Using the model-agnostic default `max_length` (=20) to control thegeneration length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
+      "  warnings.warn(\n"
+     ]
+    }
+   ],
+   "source": [
+    "outputs_test, text_test = embeds_to_captions_b2(embeds_test)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "6a95fcdf-db87-4c02-9728-09f85605fb1c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['a cat sitting on a toilet seat\\n',\n",
+       " 'a person cutting a pizza on a cutting board\\n',\n",
+       " 'a sandwich and a drink on a table\\n',\n",
+       " 'a man crossing the street in front of a truck\\n',\n",
+       " 'a giraffe standing in front of trees\\n',\n",
+       " 'three men standing together\\n',\n",
+       " 'a bird standing on a rock next to a body of water\\n',\n",
+       " 'two men sitting on a street corner in asia\\n',\n",
+       " 'a woman and two children playing tennis on a court\\n',\n",
+       " 'a tall brick building with a clock on the side\\n',\n",
+       " 'a train is on the tracks\\n',\n",
+       " 'a man and woman in the water with a surfboard\\n',\n",
+       " 'a living room with a desk and a chair\\n',\n",
+       " 'a group of men on a basketball court\\n',\n",
+       " 'a man holding an umbrella\\n',\n",
+       " 'a man in a red shirt\\n',\n",
+       " 'a group of people holding cell phones and wine glasses\\n',\n",
+       " 'a laptop computer sitting on a table in front of a television\\n',\n",
+       " 'a baseball player is swinging a bat on a field\\n']"
+      ]
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "text_test"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "9ac69fbd-55db-435b-bed6-5ae9186450e3",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "#inputss['pixel_values'].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "id": "0524f498-c8da-4e8a-8970-d75d2d0f6b8b",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "#image_test.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "5417541b-49eb-4e43-a3e2-d937d9653e04",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "max_lr = 1e-4"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "id": "da0ce190-1b3e-4c12-9e9f-91cbc076d044",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "clip_seq_dim = 257 #blip2 image encoder shapes\n",
+    "clip_emb_dim = 1408 #blip2 image encoder shapes\n",
+    "hidden_dim = 2048"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5b79bd38-6990-4504-8d45-4a68d57d8885",
+   "metadata": {},
+   "source": [
+    "### SD VAE (blurry images)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 40,
+   "id": "01baff79-8114-482b-b115-6f05aa8ad691",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "83,653,863 total\n",
+      "0 trainable\n"
+     ]
+    }
+   ],
+   "source": [
+    "from diffusers import AutoencoderKL\n",
+    "autoenc = AutoencoderKL.from_pretrained(\"madebyollin/sdxl-vae-fp16-fix\", torch_dtype=torch.float16, cache_dir=\"/fsx/proj-fmri/shared/cache\")\n",
+    "# autoenc.load_state_dict(torch.load('../train_logs/sdxl_vae_normed/best.pth')[\"model_state_dict\"])\n",
+    "autoenc.eval()\n",
+    "autoenc.requires_grad_(False)\n",
+    "autoenc.to(device)\n",
+    "utils.count_params(autoenc)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "260e5e4a-f697-4b2c-88fc-01f6a54886c0",
+   "metadata": {},
+   "source": [
+    "### MindEye modules"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 41,
+   "id": "c44c271b-173f-472e-b059-a2eda0f4c4c5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "MindEyeModule()"
+      ]
+     },
+     "execution_count": 41,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "class MindEyeModule(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super(MindEyeModule, self).__init__()\n",
+    "    def forward(self, x):\n",
+    "        return x\n",
+    "        \n",
+    "model = MindEyeModule()\n",
+    "model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 42,
+   "id": "038a5d61-4769-40b9-a004-f4e7b5b38bb0",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "32,215,040 total\n",
+      "32,215,040 trainable\n",
+      "param counts:\n",
+      "32,215,040 total\n",
+      "32,215,040 trainable\n",
+      "torch.Size([2, 1, 15729]) torch.Size([2, 1, 2048])\n"
+     ]
+    }
+   ],
+   "source": [
+    "class RidgeRegression(torch.nn.Module):\n",
+    "    # make sure to add weight_decay when initializing optimizer\n",
+    "    def __init__(self, input_size, out_features): \n",
+    "        super(RidgeRegression, self).__init__()\n",
+    "        self.out_features = out_features\n",
+    "        self.linear = torch.nn.Linear(input_size, out_features)\n",
+    "    def forward(self, x):\n",
+    "        return self.linear(x)\n",
+    "        \n",
+    "model.ridge = RidgeRegression(voxels.shape[1], out_features=hidden_dim)\n",
+    "utils.count_params(model.ridge)\n",
+    "utils.count_params(model)\n",
+    "\n",
+    "b = torch.randn((2,1,voxels.shape[1]))\n",
+    "print(b.shape, model.ridge(b).shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 43,
+   "id": "3602c333-d029-465c-8fb4-c3ccffdba6fd",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "772,419,072 total\n",
+      "772,419,072 trainable\n",
+      "param counts:\n",
+      "804,634,112 total\n",
+      "804,634,112 trainable\n",
+      "torch.Size([4, 2048])\n",
+      "torch.Size([4, 257, 1408])\n"
+     ]
+    }
+   ],
+   "source": [
+    "from functools import partial\n",
+    "from diffusers.models.vae import Decoder\n",
+    "class BrainNetwork(nn.Module):\n",
+    "    def __init__(self, out_dim=768, in_dim=15724, clip_size=768, h=4096, n_blocks=4, norm_type='ln', act_first=False, drop=.15, blurry_dim=16):\n",
+    "        super().__init__()\n",
+    "        self.blurry_dim = blurry_dim\n",
+    "        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)\n",
+    "        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU\n",
+    "        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)\n",
+    "        self.lin0 = nn.Linear(in_dim, h)\n",
+    "        self.mlp = nn.ModuleList([\n",
+    "            nn.Sequential(\n",
+    "                nn.Linear(h, h),\n",
+    "                *[item() for item in act_and_norm],\n",
+    "                nn.Dropout(drop)\n",
+    "            ) for _ in range(n_blocks)\n",
+    "        ])\n",
+    "        self.lin1 = nn.Linear(h, out_dim, bias=True)\n",
+    "        # self.blin1 = nn.Linear(out_dim, blurry_dim, bias=True)\n",
+    "        self.n_blocks = n_blocks\n",
+    "        self.clip_size = clip_size\n",
+    "        self.clip_proj = nn.Sequential(\n",
+    "            nn.LayerNorm(clip_size),\n",
+    "            nn.GELU(),\n",
+    "            nn.Linear(clip_size, 2048),\n",
+    "            nn.LayerNorm(2048),\n",
+    "            nn.GELU(),\n",
+    "            nn.Linear(2048, 2048),\n",
+    "            nn.LayerNorm(2048),\n",
+    "            nn.GELU(),\n",
+    "            nn.Linear(2048, clip_size)\n",
+    "        )\n",
+    "        # self.upsampler = Decoder(\n",
+    "        #         in_channels=64,\n",
+    "        #         out_channels=4,\n",
+    "        #         up_block_types=[\"UpDecoderBlock2D\",\"UpDecoderBlock2D\",\"UpDecoderBlock2D\"],\n",
+    "        #         block_out_channels=[64, 128, 256],\n",
+    "        #         layers_per_block=1,\n",
+    "        #     )\n",
+    "        \n",
+    "    def forward(self, x):\n",
+    "        x = self.lin0(x)\n",
+    "        residual = x\n",
+    "        for res_block in range(self.n_blocks):\n",
+    "            x = self.mlp[res_block](x)\n",
+    "            x += residual\n",
+    "            residual = x\n",
+    "        x = x.reshape(len(x), -1)\n",
+    "        x = self.lin1(x)\n",
+    "        # b = self.blin1(x)\n",
+    "        # b = self.upsampler(b.reshape(len(b), -1, 7, 7))\n",
+    "        c = self.clip_proj(x.reshape(len(x), -1, self.clip_size))\n",
+    "        # return c, b\n",
+    "        return c\n",
+    "\n",
+    "model.backbone = BrainNetwork(h=2048, in_dim=hidden_dim, clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, blurry_dim=64*7*7) \n",
+    "utils.count_params(model.backbone)\n",
+    "utils.count_params(model)\n",
+    "\n",
+    "b = torch.randn((4,hidden_dim))\n",
+    "print(b.shape)\n",
+    "clip_ = model.backbone(b)\n",
+    "print(clip_.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 44,
+   "id": "e14d0482-dc42-43b9-9ce1-953c32f2c9c1",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Done with model preparations!\n",
+      "param counts:\n",
+      "804,634,112 total\n",
+      "804,634,112 trainable\n"
+     ]
+    }
+   ],
+   "source": [
+    "no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n",
+    "opt_grouped_parameters = [\n",
+    "    {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},\n",
+    "    {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},\n",
+    "    {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},\n",
+    "]\n",
+    "\n",
+    "optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr, betas=(0.9, 0.95))\n",
+    "\n",
+    "if lr_scheduler_type == 'linear':\n",
+    "    lr_scheduler = torch.optim.lr_scheduler.LinearLR(\n",
+    "        optimizer,\n",
+    "        total_iters=int(num_epochs*(num_train*num_devices//batch_size)),\n",
+    "        last_epoch=-1\n",
+    "    )\n",
+    "elif lr_scheduler_type == 'cycle':\n",
+    "    total_steps=int(num_epochs*(num_train*num_devices//batch_size))\n",
+    "    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(\n",
+    "        optimizer, \n",
+    "        max_lr=max_lr,\n",
+    "        total_steps=total_steps,\n",
+    "        final_div_factor=1000,\n",
+    "        last_epoch=-1, pct_start=2/num_epochs\n",
+    "    )\n",
+    "    \n",
+    "def save_ckpt(tag):    \n",
+    "    ckpt_path = outdir+f'/{tag}.pth'\n",
+    "    print(f'saving {ckpt_path}',flush=True)\n",
+    "    unwrapped_model = accelerator.unwrap_model(model)\n",
+    "    try:\n",
+    "        torch.save({\n",
+    "            'epoch': epoch,\n",
+    "            'model_state_dict': unwrapped_model.state_dict(),\n",
+    "            'optimizer_state_dict': optimizer.state_dict(),\n",
+    "            'lr_scheduler': lr_scheduler.state_dict(),\n",
+    "            'train_losses': losses,\n",
+    "            'test_losses': test_losses,\n",
+    "            'lrs': lrs,\n",
+    "            }, ckpt_path)\n",
+    "    except:\n",
+    "        print(\"Couldn't save... moving on to prevent crashing.\")\n",
+    "    del unwrapped_model\n",
+    "        \n",
+    "print(\"\\nDone with model preparations!\")\n",
+    "utils.count_params(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "983f458b-35b8-49f2-b6db-80296cece730",
+   "metadata": {},
+   "source": [
+    "# Weights and Biases"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "id": "0a25a662-daa8-4de9-9233-8364800fcb6b",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "wandb mindeyev2 run captions\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mckadirt\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "wandb_config:\n",
+      " {'model_name': 'captions', 'batch_size': 128, 'num_epochs': 30, 'use_image_aug': False, 'max_lr': 0.0001, 'lr_scheduler_type': 'cycle', 'mixup_pct': 0.66, 'num_train': 24958, 'num_test': 2770, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1}\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "wandb version 0.16.0 is available!  To upgrade, please run:\n",
+       " $ pip install wandb --upgrade"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Tracking run with wandb version 0.15.5"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Run data is saved locally in <code>/fsx/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20231119_163615-o1xwsqre</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Syncing run <strong><a href='https://stability.wandb.io/ckadirt/mindeyev2/runs/o1xwsqre' target=\"_blank\">captions</a></strong> to <a href='https://stability.wandb.io/ckadirt/mindeyev2' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View project at <a href='https://stability.wandb.io/ckadirt/mindeyev2' target=\"_blank\">https://stability.wandb.io/ckadirt/mindeyev2</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View run at <a href='https://stability.wandb.io/ckadirt/mindeyev2/runs/o1xwsqre' target=\"_blank\">https://stability.wandb.io/ckadirt/mindeyev2/runs/o1xwsqre</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# params for wandb\n",
+    "if local_rank==0 and True: # only use main process for wandb logging\n",
+    "    import wandb\n",
+    "    \n",
+    "    wandb_project = 'mindeyev2'\n",
+    "    wandb_run = model_name\n",
+    "    wandb_notes = ''\n",
+    "    \n",
+    "    print(f\"wandb {wandb_project} run {wandb_run}\")\n",
+    "    wandb.login(host='https://stability.wandb.io')#, relogin=True)\n",
+    "    wandb_config = {\n",
+    "      \"model_name\": model_name,\n",
+    "      \"batch_size\": batch_size,\n",
+    "      \"num_epochs\": num_epochs,\n",
+    "      \"use_image_aug\": use_image_aug,\n",
+    "      \"max_lr\": max_lr,\n",
+    "      \"lr_scheduler_type\": lr_scheduler_type,\n",
+    "      \"mixup_pct\": mixup_pct,\n",
+    "      \"num_train\": num_train,\n",
+    "      \"num_test\": num_test,\n",
+    "      \"seed\": seed,\n",
+    "      \"distributed\": distributed,\n",
+    "      \"num_devices\": num_devices,\n",
+    "      \"world_size\": world_size,\n",
+    "    }\n",
+    "    print(\"wandb_config:\\n\",wandb_config)\n",
+    "    if False: # wandb_auto_resume\n",
+    "        print(\"wandb_id:\",model_name)\n",
+    "        wandb.init(\n",
+    "            id = model_name,\n",
+    "            project=wandb_project,\n",
+    "            name=wandb_run,\n",
+    "            config=wandb_config,\n",
+    "            notes=wandb_notes,\n",
+    "            resume=\"allow\",\n",
+    "        )\n",
+    "    else:\n",
+    "        wandb.init(\n",
+    "            project=wandb_project,\n",
+    "            name=wandb_run,\n",
+    "            config=wandb_config,\n",
+    "            notes=wandb_notes,\n",
+    "        )\n",
+    "else:\n",
+    "    wandb_log = False"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "id": "4e5de216-5318-4b45-ac02-113f03105adc",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ff0000; text-decoration-color: #ff0000\">╭──────────────────────────────────────────────────────────────────────────────────────────────────╮</span>\n",
+       "<span style=\"color: #ff0000; text-decoration-color: #ff0000\">│</span> n++                                                                                              <span style=\"color: #ff0000; text-decoration-color: #ff0000\">│</span>\n",
+       "<span style=\"color: #ff0000; text-decoration-color: #ff0000\">│</span>    <span style=\"color: #ff0000; text-decoration-color: #ff0000; font-weight: bold\">▲</span>                                                                                             <span style=\"color: #ff0000; text-decoration-color: #ff0000\">│</span>\n",
+       "<span style=\"color: #ff0000; text-decoration-color: #ff0000\">╰───────────────────────────────────────────────────────────────────────────────���──────────────────╯</span>\n",
+       "<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-weight: bold\">SyntaxError: </span>invalid syntax\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[91m╭──────────────────────────────────────────────────────────────────────────────────────────────────╮\u001b[0m\n",
+       "\u001b[91m│\u001b[0m n++                                                                                              \u001b[91m│\u001b[0m\n",
+       "\u001b[91m│\u001b[0m    \u001b[1;91m▲\u001b[0m                                                                                             \u001b[91m│\u001b[0m\n",
+       "\u001b[91m╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n",
+       "\u001b[1;91mSyntaxError: \u001b[0minvalid syntax\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5b0ae095-3203-4eb8-8606-acc2db6ccf20",
+   "metadata": {},
+   "source": [
+    "# More custom functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 34,
+   "id": "827ead88-7eb3-47cc-82da-31565063b927",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# using the same preprocessing as was used in MindEye + BrainDiffuser\n",
+    "pixcorr_preprocess = transforms.Compose([\n",
+    "    transforms.Resize(425, interpolation=transforms.InterpolationMode.BILINEAR),\n",
+    "])\n",
+    "def pixcorr(images,brains):\n",
+    "    # Flatten images while keeping the batch dimension\n",
+    "    all_images_flattened = pixcorr_preprocess(images).reshape(len(images), -1)\n",
+    "    all_brain_recons_flattened = pixcorr_preprocess(brains).view(len(brains), -1)\n",
+    "    corrmean = torch.diag(utils.batchwise_pearson_correlation(all_images_flattened, all_brain_recons_flattened)).mean()\n",
+    "    return corrmean"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d5690151-2131-4918-b750-e869cbd1a8a8",
+   "metadata": {},
+   "source": [
+    "# Main"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 51,
+   "id": "12de6387-6e18-4e4b-b5ce-a847d625330a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "epoch = 0\n",
+    "losses, test_losses, lrs = [], [], []\n",
+    "best_test_loss = 1e9\n",
+    "soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))\n",
+    "\n",
+    "# Optionally resume from checkpoint #\n",
+    "if resume_from_ckpt:\n",
+    "    print(\"\\n---resuming from last.pth ckpt---\\n\")\n",
+    "    try:\n",
+    "        checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')\n",
+    "    except:\n",
+    "        print('last.pth failed... trying last_backup.pth')\n",
+    "        checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')\n",
+    "    epoch = checkpoint['epoch']\n",
+    "    print(\"Epoch\",epoch)\n",
+    "    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
+    "    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n",
+    "    diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])\n",
+    "    del checkpoint\n",
+    "elif wandb_log:\n",
+    "    if wandb.run.resumed:\n",
+    "        print(\"\\n---resuming from last.pth ckpt---\\n\")\n",
+    "        try:\n",
+    "            checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')\n",
+    "        except:\n",
+    "            print('last.pth failed... trying last_backup.pth')\n",
+    "            checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')\n",
+    "        epoch = checkpoint['epoch']\n",
+    "        print(\"Epoch\",epoch)\n",
+    "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
+    "        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n",
+    "        diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])\n",
+    "        del checkpoint\n",
+    "torch.cuda.empty_cache()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "id": "b4755749-2d99-4e98-ad98-3df661746058",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "checkpoint = torch.load('/fsx/proj-fmri/ckadirt/MindEyeV2/train_logs/caption_clip_0.5_bz/last.pth', map_location='cpu')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 45,
+   "id": "cd3dc793-5a20-4b48-959c-bc64430c8c02",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<All keys matched successfully>"
+      ]
+     },
+     "execution_count": 45,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model.load_state_dict(checkpoint['model_state_dict'])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 46,
+   "id": "0faa2c6a-00da-4b66-b5e5-8c4864768805",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "MindEyeModule(\n",
+       "  (ridge): RidgeRegression(\n",
+       "    (linear): Linear(in_features=15729, out_features=2048, bias=True)\n",
+       "  )\n",
+       "  (backbone): BrainNetwork(\n",
+       "    (lin0): Linear(in_features=2048, out_features=2048, bias=True)\n",
+       "    (mlp): ModuleList(\n",
+       "      (0-3): 4 x Sequential(\n",
+       "        (0): Linear(in_features=2048, out_features=2048, bias=True)\n",
+       "        (1): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
+       "        (2): GELU(approximate='none')\n",
+       "        (3): Dropout(p=0.15, inplace=False)\n",
+       "      )\n",
+       "    )\n",
+       "    (lin1): Linear(in_features=2048, out_features=361856, bias=True)\n",
+       "    (clip_proj): Sequential(\n",
+       "      (0): LayerNorm((1408,), eps=1e-05, elementwise_affine=True)\n",
+       "      (1): GELU(approximate='none')\n",
+       "      (2): Linear(in_features=1408, out_features=2048, bias=True)\n",
+       "      (3): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
+       "      (4): GELU(approximate='none')\n",
+       "      (5): Linear(in_features=2048, out_features=2048, bias=True)\n",
+       "      (6): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
+       "      (7): GELU(approximate='none')\n",
+       "      (8): Linear(in_features=2048, out_features=1408, bias=True)\n",
+       "    )\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 46,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 47,
+   "id": "99f09f76-4481-4133-b09a-a22b10dbc0c4",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "model, optimizer, train_dl, test_dl, lr_scheduler = accelerator.prepare(\n",
+    "model, optimizer, train_dl, test_dl, lr_scheduler\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bfeeda32-82ca-4364-bce1-eaa41b4f3e25",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "\"\"\"transform = transforms.Compose(\n",
+    "            [\n",
+    "                transforms.Resize(\n",
+    "                    (224, 224),\n",
+    "                ),\n",
+    "                transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),\n",
+    "            ]\n",
+    "        )\n",
+    "\n",
+    "def tensor_2_embed(image):    \n",
+    "    image_for_blip2 = transform(image)\n",
+    "    \n",
+    "    #Generate embeddings\n",
+    "    with blip2_model.maybe_autocast():\n",
+    "        blip2_target = blip2_model.ln_vision(blip2_model.visual_encoder(image_for_blip2))\n",
+    "                \n",
+    "    return blip2_target\n",
+    "\n",
+    "def embed_2_caption(image_embeds, model):\n",
+    "    image_embeds = image_embeds.float()\n",
+    "    image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(\n",
+    "                image.device)\n",
+    "\n",
+    "    query_tokens = model.query_tokens.expand(image_embeds.shape[0], -1, -1)\n",
+    "    query_output = model.Qformer.bert(\n",
+    "                query_embeds=query_tokens,\n",
+    "                encoder_hidden_states=image_embeds,\n",
+    "                encoder_attention_mask=image_atts,\n",
+    "                return_dict=True)\n",
+    "\n",
+    "    inputs_t5 = model.t5_proj(query_output.last_hidden_state)\n",
+    "    atts_t5 = torch.ones(inputs_t5.size()[:-1], dtype=torch.long).to(image.device)\n",
+    "    prompt = model.prompt\n",
+    "    input_tokens = model.t5_tokenizer(\n",
+    "                prompt, padding=\"longest\", return_tensors=\"pt\"\n",
+    "            ).to(image.device)\n",
+    "    encoder_atts = torch.cat([atts_t5, input_tokens.attention_mask], dim=1)\n",
+    "    \n",
+    "    with model.maybe_autocast(dtype=torch.bfloat16):\n",
+    "            inputs_embeds = model.t5_model.encoder.embed_tokens(input_tokens.input_ids)\n",
+    "            inputs_embeds = torch.cat([inputs_t5, inputs_embeds], dim=1)\n",
+    "\n",
+    "            outputs = model.t5_model.generate(\n",
+    "                inputs_embeds=inputs_embeds,\n",
+    "                attention_mask=encoder_atts)\n",
+    "            output_text = model.t5_tokenizer.batch_decode(\n",
+    "                outputs, skip_special_tokens=True)\n",
+    "    \n",
+    "    return output_text\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 48,
+   "id": "636b4684-df9a-4e29-8683-86fb035ba690",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "wandb_log = False"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 49,
+   "id": "0847b380-2edb-4a56-9b33-fdc4c0c3f8d3",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "predicted_embeddings = None"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 52,
+   "id": "60be0d5f-3e94-4612-9373-61b53d836393",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "captions starting with epoch 0 / 30\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                    | 0/30 [00:17<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "===Finished!===\n",
+      "\n",
+      "saving /fsx/proj-fmri/ckadirt/MindEyeV2/train_logs/captions/last.pth\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(f\"{model_name} starting with epoch {epoch} / {num_epochs}\")\n",
+    "progress_bar = tqdm(range(epoch,num_epochs), ncols=1200, disable=(local_rank!=0))\n",
+    "test_image, test_voxel = None, None\n",
+    "mse = nn.MSELoss()\n",
+    "for epoch in progress_bar:\n",
+    "    model.train()\n",
+    "    \n",
+    "    fwd_percent_correct = 0.\n",
+    "    bwd_percent_correct = 0.\n",
+    "    test_fwd_percent_correct = 0.\n",
+    "    test_bwd_percent_correct = 0.\n",
+    "\n",
+    "    loss_clip_total = 0.\n",
+    "    loss_blurry_total = 0.\n",
+    "    test_loss_clip_total = 0.\n",
+    "    test_loss_blurry_total = 0.\n",
+    "\n",
+    "    blurry_pixcorr = 0.\n",
+    "    test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1\n",
+    "    \n",
+    "    \"\"\"for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):\n",
+    "        if epoch == 0:\n",
+    "            lrs.append(0)\n",
+    "            break\n",
+    "        with torch.cuda.amp.autocast():\n",
+    "            optimizer.zero_grad()\n",
+    "\n",
+    "            voxel = voxels[behav[:,0,5].cpu().long()].to(device)\n",
+    "            \n",
+    "            image = images[behav[:,0,0].cpu().long()].to(device).float()\n",
+    "\n",
+    "            # blurry_image_enc = autoenc.encode(image).latent_dist.mode()\n",
+    "            \n",
+    "            if use_image_aug: image = img_augment(image)\n",
+    "            # clip_target = clip_model.embed_image(image)\n",
+    "            clip_target = embed_images_b2(image)[0].to(device) #####CHANGED\n",
+    "            assert not torch.any(torch.isnan(clip_target))\n",
+    "  \n",
+    "            if epoch < int(mixup_pct * num_epochs):\n",
+    "                voxel, perm, betas, select = utils.mixco(voxel)\n",
+    "\n",
+    "            voxel_ridge = model.ridge(voxel)\n",
+    "            \n",
+    "            # clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)\n",
+    "            clip_voxels = model.backbone(voxel_ridge)\n",
+    "            \n",
+    "            clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)\n",
+    "            clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)\n",
+    "\n",
+    "            if epoch < int(mixup_pct * num_epochs):                \n",
+    "                loss_clip = utils.mixco_nce(\n",
+    "                     clip_voxels_norm,\n",
+    "                     clip_target_norm,\n",
+    "                     temp=.006, \n",
+    "                     perm=perm, betas=betas, select=select)\n",
+    "            else:\n",
+    "                epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]\n",
+    "                loss_clip = utils.soft_clip_loss(\n",
+    "                     clip_voxels_norm,\n",
+    "                     clip_target_norm,\n",
+    "                     temp=epoch_temp)\n",
+    "            \n",
+    "            loss_mse= mse(clip_voxels, clip_target)\n",
+    "\n",
+    "            # loss_blurry = mse(blurry_image_enc_, blurry_image_enc) \n",
+    "\n",
+    "            loss_clip_total += loss_clip.item()\n",
+    "            # loss_blurry_total += loss_blurry.item()\n",
+    "\n",
+    "            # loss = loss_blurry + loss_clip\n",
+    "            loss = 0.7 * loss_clip + 0.3 * loss_mse\n",
+    "            if (train_i % 10 == 0):\n",
+    "                print(train_i, loss)\n",
+    "            # print(batch_size)\n",
+    "            utils.check_loss(loss)\n",
+    "            accelerator.backward(loss)\n",
+    "            optimizer.step()\n",
+    "    \n",
+    "            losses.append(loss.item())\n",
+    "            lrs.append(optimizer.param_groups[0]['lr'])\n",
+    "    \n",
+    "            # forward and backward top 1 accuracy        \n",
+    "            labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) \n",
+    "            fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)\n",
+    "            bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)\n",
+    "\n",
+    "            # with torch.no_grad():\n",
+    "            #     # only doing pixcorr eval on a subset (8) of the samples per batch because its costly & slow to compute autoenc.decode()\n",
+    "            #     random_samps = np.random.choice(np.arange(len(voxel)), size=8, replace=False)\n",
+    "            #     blurry_recon_images = autoenc.decode(blurry_image_enc_[random_samps]).sample.clamp(0,1)\n",
+    "            #     blurry_pixcorr += pixcorr(image[random_samps], blurry_recon_images)\n",
+    "\n",
+    "            if lr_scheduler_type is not None:\n",
+    "                lr_scheduler.step()\"\"\"\n",
+    "    \n",
+    "    model.eval()\n",
+    "    for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):\n",
+    "        with torch.cuda.amp.autocast():\n",
+    "            with torch.no_grad():   \n",
+    "                # all test samples should be loaded per batch such that test_i should never exceed 0\n",
+    "                if len(behav) != num_test: print(\"!\",len(behav),num_test)\n",
+    "                \n",
+    "                ## Average same-image repeats ##\n",
+    "                if test_image is None:\n",
+    "                    voxel = voxels[behav[:,0,5].cpu().long()].to(device)\n",
+    "                    \n",
+    "                    image = behav[:,0,0].cpu().long()\n",
+    "                    \n",
+    "                    unique_image, sort_indices = torch.unique(image, return_inverse=True)\n",
+    "                    for im in unique_image:\n",
+    "                        locs = torch.where(im == image)[0]\n",
+    "                        if test_image is None:\n",
+    "                            test_image = images[im][None]\n",
+    "                            test_voxel = torch.mean(voxel[locs],axis=0)[None]\n",
+    "                        else:\n",
+    "                            test_image = torch.vstack((test_image, images[im][None]))\n",
+    "                            test_voxel = torch.vstack((test_voxel, torch.mean(voxel[locs],axis=0)[None]))\n",
+    "    \n",
+    "                # sample of batch_size\n",
+    "                random_indices = torch.arange(len(test_voxel))[:batch_size] #torch.randperm(len(test_voxel))[:300]\n",
+    "                voxel = test_voxel[random_indices].to(device)\n",
+    "                image = test_image[random_indices].to(device)\n",
+    "                assert len(image) == batch_size\n",
+    "    \n",
+    "                # blurry_image_enc = autoenc.encode(image).latent_dist.mode()\n",
+    "        \n",
+    "                # clip_target = clip_model.embed_image(image.float())\n",
+    "                clip_target = embed_images_b2(image)[0].to(device) #####CHANGED\n",
+    "    \n",
+    "                voxel_ridge = model.ridge(voxel)\n",
+    "                \n",
+    "                # clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)\n",
+    "                clip_voxels = model.backbone(voxel_ridge)\n",
+    "                \n",
+    "                predicted_embeddings = clip_voxels\n",
+    "                break\n",
+    "                \n",
+    "                clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)\n",
+    "                clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)\n",
+    "        \n",
+    "                # loss_clip = utils.soft_clip_loss(\n",
+    "                #     clip_voxels_norm,\n",
+    "                #     clip_target_norm,\n",
+    "                #     temp=.006)\n",
+    "                \n",
+    "                loss_clip = mse(clip_voxels, clip_target)\n",
+    "\n",
+    "                # loss_blurry = mse(blurry_image_enc_, blurry_image_enc)\n",
+    "                \n",
+    "                # loss = loss_blurry + loss_clip\n",
+    "                loss = loss_clip\n",
+    "                \n",
+    "                utils.check_loss(loss)\n",
+    "        \n",
+    "                test_losses.append(loss.item())\n",
+    "        \n",
+    "                # forward and backward top 1 accuracy        \n",
+    "                labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) \n",
+    "                test_fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)\n",
+    "                test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)\n",
+    "\n",
+    "                # # halving the batch size because the decoder is computationally heavy\n",
+    "                # blurry_recon_images = autoenc.decode(blurry_image_enc_[:len(voxel)//2]).sample.clamp(0,1)\n",
+    "                # blurry_recon_images = torch.vstack((blurry_recon_images, autoenc.decode(blurry_image_enc_[len(voxel)//2:]).sample.clamp(0,1)))\n",
+    "                # test_blurry_pixcorr += pixcorr(image, blurry_recon_images)\n",
+    "\n",
+    "                #Find captions and print next to images\n",
+    "                #caption1 = embed_2_caption(clip_voxels[[0]], blip2_model)\n",
+    "                #caption2 = embed_2_caption(clip_voxels[[1]], blip2_model)\n",
+    "\n",
+    "                #true_embed1 = tensor_2_embed(image[[0]])\n",
+    "                #true_embed2 = tensor_2_embed(image[[1]])\n",
+    "\n",
+    "                # print(clip_voxels[[0]].shape)\n",
+    "                # print(true_embed1.shape)\n",
+    "                \n",
+    "                #true_caption1 = embed_2_caption(true_embed1, blip2_model)\n",
+    "                #true_caption2 = embed_2_caption(true_embed2, blip2_model)\n",
+    "                \n",
+    "                # transform blurry recon latents to images and plot it\n",
+    "                #fig, axes = plt.subplots(2, 2, figsize=(8, 4))\n",
+    "                #axes[0,0].imshow(utils.torch_to_Image(image[[0]]))\n",
+    "                #axes[0,1].imshow(utils.torch_to_Image(image[[1]]))\n",
+    "                #axes[0,0].axis('off'); axes[0,1].axis('off'); axes[1,0].axis('off'); axes[1,1].axis('off')\n",
+    "                #axes[0,0].set_title(caption1)\n",
+    "                #axes[0,1].set_title(caption2)\n",
+    "                #axes[1,0].set_title(true_caption1)\n",
+    "                #axes[1,1].set_title(true_caption2)\n",
+    "\n",
+    "                #plt.show()\n",
+    "                \n",
+    "                # # transform blurry recon latents to images and plot it\n",
+    "                # fig, axes = plt.subplots(1, 4, figsize=(8, 4))\n",
+    "                # axes[0].imshow(utils.torch_to_Image(image[[0]]))\n",
+    "                # axes[1].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[0]]).sample.clamp(0,1)))\n",
+    "                # axes[2].imshow(utils.torch_to_Image(image[[1]]))\n",
+    "                # axes[3].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[1]]).sample.clamp(0,1)))\n",
+    "                # axes[0].axis('off'); axes[1].axis('off'); axes[2].axis('off'); axes[3].axis('off')\n",
+    "                # axes[0].set_title(caption1)\n",
+    "                # axes[3].set_title(caption2)\n",
+    "                # plt.show()\n",
+    "                \n",
+    "    break\n",
+    "    if local_rank==0:      \n",
+    "        # if utils.is_interactive(): clear_output(wait=True)\n",
+    "        assert (test_i+1) == 1\n",
+    "        logs = {\"train/loss\": np.mean(losses[-(train_i+1):]),\n",
+    "            \"test/loss\": np.mean(test_losses[-(test_i+1):]),\n",
+    "            \"train/lr\": lrs[-1],\n",
+    "            \"train/num_steps\": len(losses),\n",
+    "            \"test/num_steps\": len(test_losses),\n",
+    "            \"train/fwd_pct_correct\": fwd_percent_correct / (train_i + 1),\n",
+    "            \"train/bwd_pct_correct\": bwd_percent_correct / (train_i + 1),\n",
+    "            \"test/test_fwd_pct_correct\": test_fwd_percent_correct / (test_i + 1),\n",
+    "            \"test/test_bwd_pct_correct\": test_bwd_percent_correct / (test_i + 1),\n",
+    "            \"train/loss_clip_total\": loss_clip_total / (train_i + 1),\n",
+    "            \"train/loss_blurry_total\": loss_blurry_total / (train_i + 1),\n",
+    "            \"test/loss_clip_total\": test_loss_clip_total / (test_i + 1),\n",
+    "            \"test/loss_blurry_total\": test_loss_blurry_total / (test_i + 1),\n",
+    "            \"train/blurry_pixcorr\": blurry_pixcorr / (train_i + 1),\n",
+    "            \"test/blurry_pixcorr\": test_blurry_pixcorr / (test_i + 1),\n",
+    "            }\n",
+    "        progress_bar.set_postfix(**logs)\n",
+    "                          \n",
+    "        fig, axes = plt.subplots(1, 8, figsize=(10, 4))\n",
+    "        jj=-1\n",
+    "        for j in [0,1,2,3,4,5,6,7]:\n",
+    "            jj+=1\n",
+    "            axes[jj].imshow(utils.torch_to_Image(image[j]))\n",
+    "            axes[jj].axis('off')\n",
+    "\n",
+    "        if wandb_log:\n",
+    "            generated_captions = embeds_to_captions_b2(clip_voxels[0:8])\n",
+    "            print(generated_captions[1])\n",
+    "            logs[f\"test/recons\"] = wandb.Image(fig, caption=f\"epoch{epoch:03d}\" + \"\\n\".join(generated_captions[1]))\n",
+    "            plt.close()\n",
+    "        # Save model checkpoint and reconstruct\n",
+    "        if epoch % ckpt_interval == 0:\n",
+    "            if not utils.is_interactive():\n",
+    "                save_ckpt(f'last')\n",
+    "                \n",
+    "        if wandb_log: wandb.log(logs)\n",
+    "\n",
+    "    # wait for other GPUs to catch up if needed\n",
+    "    accelerator.wait_for_everyone()\n",
+    "    torch.cuda.empty_cache()\n",
+    "    gc.collect()\n",
+    "\n",
+    "print(\"\\n===Finished!===\\n\")\n",
+    "if ckpt_saving:\n",
+    "    save_ckpt(f'last')\n",
+    "if not utils.is_interactive():\n",
+    "    sys.exit(0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 54,
+   "id": "f5b47c76-a97a-48ee-b4b3-051c17aebac4",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([128, 257, 1408])"
+      ]
+     },
+     "execution_count": 54,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "predicted_embeddings.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 55,
+   "id": "92d0029f-079f-4710-bf43-bc9e3fd08d5e",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/transformers/generation/utils.py:1260: UserWarning: Using the model-agnostic default `max_length` (=20) to control thegeneration length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "['a group of people are sitting around a table\\n', 'a man is holding a glass of water in front of a television\\n', 'a man is riding a skateboard on a hill\\n', 'a group of people standing around a bike\\n', 'a building with a sign that says \"the house\"\\n', 'a plate of food with vegetables and meat\\n', 'a white cup with a small bottle of wine\\n', 'a group of people playing baseball and one is holding a ball\\n']\n"
+     ]
+    }
+   ],
+   "source": [
+    "generated_captions = embeds_to_captions_b2(predicted_embeddings[0:8])\n",
+    "print(generated_captions[1])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 75,
+   "id": "88750a6d-0b61-4943-a7e5-1d675bbb4f8f",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "['a group of people are sitting at a table with food and drinks\\n', 'a man in a kitchen with a large screen\\n', 'a man on a surfboard with his legs in the air\\n', 'a group of people are standing on the beach in front of a boat\\n', 'a building with a sign that says \"home of the person\"\\n', 'a vegetable salad with a variety of vegetables and other ingredients\\n', 'a white cup with a small amount of coffee and a bottle of wine\\n', 'a group of people playing baseball and soccer\\n']\n"
+     ]
+    }
+   ],
+   "source": [
+    "generated_captions = embeds_to_captions_b2(predicted_embeddings[0:8], sample = True, temp = 0.3)\n",
+    "print(generated_captions[1])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 95,
+   "id": "d99e7583-0f26-41c1-8035-a1aa3b1c2d55",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "def concatenate_lists_any_depth(list1, list2):\n",
+    "    \"\"\"\n",
+    "    Concatenates two lists of potentially varying depths, forming a new list of lists.\n",
+    "\n",
+    "    Args:\n",
+    "    list1 (list): The first list to concatenate. Elements can be of any type.\n",
+    "    list2 (list): The second list to concatenate. Elements can be of any type.\n",
+    "\n",
+    "    Returns:\n",
+    "    list: A new list containing lists of elements from the original lists.\n",
+    "    \"\"\"\n",
+    "    # Ensure that both lists have the same length\n",
+    "    if len(list1) != len(list2):\n",
+    "        raise ValueError(\"Lists must be of the same length\")\n",
+    "\n",
+    "    concatenated_list = []\n",
+    "\n",
+    "    for a, b in zip(list1, list2):\n",
+    "        # If the elements are not lists, convert them to lists\n",
+    "        if not isinstance(a, list):\n",
+    "            a = [a]\n",
+    "        if not isinstance(b, list):\n",
+    "            b = [b]\n",
+    "\n",
+    "        # Concatenate the lists\n",
+    "        concatenated_list.append(a + b)\n",
+    "\n",
+    "    return concatenated_list"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 96,
+   "id": "ed8167ea-a3ab-438a-aa85-f1309047199c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "def sample_several(embeddings, num=10, temp=0.3):\n",
+    "    # embeddings shape = batch, 257, 1408\n",
+    "    results = None  # Initialize results as None\n",
+    "\n",
+    "    for i in range(num):  # Iterate from 0 to num-1\n",
+    "        if results is None:\n",
+    "            # For the first iteration, assign the results directly\n",
+    "            results = embeds_to_captions_b2(embeddings, sample=True, temp=temp)[1]\n",
+    "        else:\n",
+    "            # For subsequent iterations, combine the new results with the existing ones\n",
+    "            new_results = embeds_to_captions_b2(embeddings, sample=True, temp=temp)[1]\n",
+    "            results = concatenate_lists_any_depth(results, new_results)\n",
+    "\n",
+    "    return results  # Return the combined results\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 77,
+   "id": "6700e130-8ae4-4475-a5b4-972fd8b9717a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "['a group of people sitting on a bench in front of a building\\n', 'a woman is using a computer to make a video\\n', 'a man in a black shirt is sitting on a surfboard\\n', 'a group of people on the beach with a bike and some other things\\n', 'a large building with a sign that says \"the old farmhouse\"\\n', 'a plate with many different types of vegetables\\n', 'a white cup with a bottle of wine and a small bottle of wine\\n', 'a group of people are playing baseball in a field\\n']\n"
+     ]
+    }
+   ],
+   "source": [
+    "generated_captions = embeds_to_captions_b2(predicted_embeddings[0:8], sample = True, temp = 0.3)\n",
+    "print(generated_captions[1])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 99,
+   "id": "f0e111e3-6134-4a63-a6d7-17b3441be8c8",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[['people are sitting at a table with a bunch of chairs\\n',\n",
+       "  'several people in the yard with some food\\n',\n",
+       "  'people sitting on a bench near a water fountain\\n',\n",
+       "  'a group of people are sitting around a table\\n',\n",
+       "  'a group of people in a room with several people in the foreground\\n',\n",
+       "  'a group of people sitting around a table with food\\n',\n",
+       "  'the people in the background are sitting on the edge of a table\\n',\n",
+       "  'beverages and food are served at a family picnic\\n',\n",
+       "  'a group of people eating in a restaurant\\n',\n",
+       "  'a group of people sitting around a table\\n',\n",
+       "  'people are sitting at a table next to a tree\\n',\n",
+       "  'people are sitting around a table with a lot of food\\n'],\n",
+       " ['a person is holding a newspaper in a restaurant\\n',\n",
+       "  'the man is holding a cup of coffee in front of a television\\n',\n",
+       "  'a woman is preparing to cook in a kitchen\\n',\n",
+       "  'a man working in an office setting with a computer and a man in a chair\\n',\n",
+       "  'a person is using a smartphone in a restaurant\\n',\n",
+       "  'a man is holding a glass of water in front of a television\\n',\n",
+       "  'a man in a kitchen with a knife and a cup of coffee\\n',\n",
+       "  'the kitchen at the new york times\\n',\n",
+       "  'a man is holding a knife and cutting a piece of pizza\\n',\n",
+       "  'a man is reading a book while another is working on a computer\\n',\n",
+       "  'a person is using a computer to make a presentation\\n',\n",
+       "  'a man is holding up a box of food\\n'],\n",
+       " ['a man in a suit and a woman wearing a helmet on a surfboard\\n',\n",
+       "  'a person is on the ground while holding onto a skateboard\\n',\n",
+       "  'a man in a beach chair riding a skateboard\\n',\n",
+       "  'a woman is standing on a surfboard in the ocean while holding a skateboard\\n',\n",
+       "  'a man is riding on a surfboard\\n',\n",
+       "  'a man on his knees in a surfboard with his leg up\\n',\n",
+       "  'a man is doing a trick on a skateboard\\n',\n",
+       "  'a man is riding a skateboard on a wave\\n',\n",
+       "  'a person is sitting on a surfboard while another person is riding on it\\n',\n",
+       "  'a man in a jumpsuit is holding onto a surfboard\\n',\n",
+       "  'a man is jumping on a surfboard while another is sitting on it\\n',\n",
+       "  'a man is sitting on a surfboard while he is riding\\n'],\n",
+       " ['a picture of a man riding a bike next to a bike\\n',\n",
+       "  'a group of people standing on a street with a bike\\n',\n",
+       "  'people are sitting around a picnic table and a bike is being ridden\\n',\n",
+       "  'a group of people are on a beach with a bike and a car\\n',\n",
+       "  \"the world's largest boat race is underway in the bay of britain\\n\",\n",
+       "  'a man and his bike standing on the side of a road\\n',\n",
+       "  'a motorcycle is sitting on top of a hill with a boat and a bicycle\\n',\n",
+       "  'a bunch of people are standing around a large park\\n',\n",
+       "  'a group of people standing around a table with bicycles\\n',\n",
+       "  'a man with his bike and helmet in the air\\n',\n",
+       "  'the sun is shining brightly and there are people walking around\\n',\n",
+       "  'a group of people standing on a beach next to a boat\\n'],\n",
+       " ['the home has a large yellow sign\\n',\n",
+       "  'the building has two small windows and a sign\\n',\n",
+       "  'a view of a home with a building in the background\\n',\n",
+       "  'the house has been built in the style of a traditional english cottage\\n',\n",
+       "  'the house is on the corner of a street\\n',\n",
+       "  'the home is an old style building with a white door\\n',\n",
+       "  'the house is in a residential area with many buildings\\n',\n",
+       "  'the building is white and has a red roof\\n',\n",
+       "  'the old building is now a park and recreation center\\n',\n",
+       "  'a large building with a lot of windows and a lot of people\\n',\n",
+       "  'a large house with a white door and a blue sign\\n',\n",
+       "  'the house is in a residential area with a front and back door\\n'],\n",
+       " ['the vegetables are arranged in a square shape on the table\\n',\n",
+       "  'a plate full of vegetables and fruit with a knife and fork\\n',\n",
+       "  'a plate of various vegetables with a knife\\n',\n",
+       "  'a plate with several different types of food\\n',\n",
+       "  'a plate of food with various vegetables and meat\\n',\n",
+       "  'a picture of some vegetables and a plate of food\\n',\n",
+       "  'a close up of several types of food on a table\\n',\n",
+       "  'a plate of food with a variety of vegetables\\n',\n",
+       "  'a large plate with many different types of food\\n',\n",
+       "  'a plate of vegetables and meat on a table\\n',\n",
+       "  'a plate with lots of different types of vegetables\\n',\n",
+       "  'a close up of some food with a knife\\n'],\n",
+       " ['a white cup with a green tea bag and a small bottle of alcohol\\n',\n",
+       "  'a bottle of wine with two glasses and a spoon\\n',\n",
+       "  'the chocolate bar is sitting next to the bottle of wine\\n',\n",
+       "  'a white and black cup and a bottle of wine\\n',\n",
+       "  'a white cup sitting next to some drinks\\n',\n",
+       "  'a bottle of wine and a bottle of champagne on a table\\n',\n",
+       "  'a white cup with two pills and a small bottle of wine\\n',\n",
+       "  'a bottle of wine and a cup of coffee next to a bottle of wine\\n',\n",
+       "  'a bottle of wine and a bottle of beer in a glass\\n',\n",
+       "  'a bottle of wine, a bottle of beer and a wine bottle\\n',\n",
+       "  'a bottle of wine and a cup with some food\\n',\n",
+       "  'a glass of wine and a pair of glasses on a table\\n'],\n",
+       " ['a group of people in white and blue uniforms playing baseball\\n',\n",
+       "  'a group of people playing baseball in a field\\n',\n",
+       "  'a group of people playing a game of baseball\\n',\n",
+       "  'a group of people standing on a field and one is holding a tennis ball\\n',\n",
+       "  'a group of people in uniform playing baseball\\n',\n",
+       "  'two men and a woman in the middle of a game\\n',\n",
+       "  'a group of people playing baseball in the grass\\n',\n",
+       "  'a group of men and women are playing baseball\\n',\n",
+       "  'july 15th, 2011 - june 20th, 2012 - june 17th,',\n",
+       "  'the team is playing soccer and one is holding a ball\\n',\n",
+       "  'people are playing baseball with each other and one is holding a ball\\n',\n",
+       "  'the women are laughing and the man is running\\n']]"
+      ]
+     },
+     "execution_count": 99,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "several = sample_several(predicted_embeddings[0:8], num = 12, temp = 0.5)\n",
+    "several"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 100,
+   "id": "7ced031a-f259-4797-afd7-876fa62cdcfd",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[['a group of people are sitting around a table\\n',\n",
+       "  'a group of people are sitting around a table with food\\n',\n",
+       "  'a group of people sitting at a table with food\\n',\n",
+       "  'a group of people are sitting on the ground in front of a table\\n',\n",
+       "  'a group of people sitting around a table with a person and a dog\\n',\n",
+       "  'a group of people are sitting on the ground and eating\\n',\n",
+       "  'the group is sitting around a table with food\\n',\n",
+       "  'people are sitting around a table with food\\n',\n",
+       "  'a group of people sitting around a table with food\\n',\n",
+       "  'the people are eating in front of a table\\n',\n",
+       "  'a group of people are sitting on a bench in a field\\n',\n",
+       "  'a group of people are sitting on a bench\\n'],\n",
+       " ['a man is using a computer and a phone\\n',\n",
+       "  'a person in a kitchen with a large screen\\n',\n",
+       "  'a man is preparing food in a kitchen\\n',\n",
+       "  'a man is standing in front of a computer and a woman is sitting behind him\\n',\n",
+       "  'a man is using a computer to play a game\\n',\n",
+       "  'a man is using a computer to play a game\\n',\n",
+       "  'a man in a kitchen with a large television\\n',\n",
+       "  'a man is holding a glass of water in front of a television\\n',\n",
+       "  'the man is holding a bottle of water and a glass\\n',\n",
+       "  'a man is using a computer to make a video\\n',\n",
+       "  'a man is serving food at a restaurant\\n',\n",
+       "  'a man is holding a drink in his hand\\n'],\n",
+       " ['a man with a skateboard is riding on a wave\\n',\n",
+       "  'a man is riding a skateboard on a hill\\n',\n",
+       "  'a man is riding a skateboard on a hill\\n',\n",
+       "  'a person is sitting on a surfboard while another person is riding on it\\n',\n",
+       "  'a man is riding a surfboard on a wave\\n',\n",
+       "  'a man with a skateboard is on top of a hill\\n',\n",
+       "  'a person in a surfboard is riding a wave\\n',\n",
+       "  'a man on a surfboard is riding on a wave\\n',\n",
+       "  'a man in a suit and a woman in a bikini are playing on a surf board\\n',\n",
+       "  'a man is riding a skateboard while wearing a helmet\\n',\n",
+       "  'a man on the surf board with his legs in the air\\n',\n",
+       "  'a man in a suit is playing a game with a skateboard\\n'],\n",
+       " ['a group of people standing on a beach with a bike\\n',\n",
+       "  'a group of people standing on a beach with a bike\\n',\n",
+       "  'a group of people standing on a road with a bike and a car\\n',\n",
+       "  'a group of people in the water with two bikes\\n',\n",
+       "  'the bike is in the middle of the road and there are two people on the side of the',\n",
+       "  'a group of people standing around a car with a bike\\n',\n",
+       "  'a man is standing on a bike with a skateboard\\n',\n",
+       "  'a group of people riding bicycles on a road\\n',\n",
+       "  'a bicycle is in the middle of a field with a person on it\\n',\n",
+       "  'a man is standing on a bicycle with a helmet and a skateboard\\n',\n",
+       "  'a photo of a bicycle with a man on it\\n',\n",
+       "  'a group of people riding bicycles on a road\\n'],\n",
+       " ['a building with a sign that says \"the old man\"\\n',\n",
+       "  'a house with a sign that says \"the house that james bond built\"\\n',\n",
+       "  'a building with a sign that says \"the house\"\\n',\n",
+       "  'a house with a sign that says \"museum\"\\n',\n",
+       "  'a building with a sign that says \"the home of the person\"\\n',\n",
+       "  'a building with a sign that says \"the museum of american history\"\\n',\n",
+       "  'a white building with a sign on the side\\n',\n",
+       "  'a brown house with a white roof and a green sign\\n',\n",
+       "  'a house with a large sign on the side\\n',\n",
+       "  'a building with a sign that says \"the building\"\\n',\n",
+       "  'the building is in the middle of the street\\n',\n",
+       "  'the front of an old building with a sign\\n'],\n",
+       " ['a plate of different types of vegetables and meat\\n',\n",
+       "  'a close up of some vegetables and meat\\n',\n",
+       "  'a plate with a variety of different foods on it\\n',\n",
+       "  'a plate of vegetables and meat with a green border\\n',\n",
+       "  'a plate of vegetables with a variety of toppings\\n',\n",
+       "  'a plate of food with different types of vegetables\\n',\n",
+       "  'a plate of food with various vegetables and meat\\n',\n",
+       "  'a plate of vegetables with some green leaves on it\\n',\n",
+       "  'a bunch of vegetables and mushrooms on a plate\\n',\n",
+       "  'a bunch of vegetables and fruit on a table\\n',\n",
+       "  'a plate of vegetables and other items on a table\\n',\n",
+       "  'a close up of some vegetables and meat\\n'],\n",
+       " ['a white cup with a spoon and a spoon\\n',\n",
+       "  'a bottle of wine and a bottle of champagne\\n',\n",
+       "  'a white cup with a small bottle and a small bottle of wine\\n',\n",
+       "  'a white cup with a small bottle of wine and a small bottle of water\\n',\n",
+       "  'the bottle is open and the bottle is next to a cup\\n',\n",
+       "  'the white cup with a small bottle of wine and a small bottle of wine\\n',\n",
+       "  'a white cup with a black handle and a pair of scissors\\n',\n",
+       "  'a bottle of wine and a bottle of wine glasses\\n',\n",
+       "  'a bottle of wine and a bottle of champagne\\n',\n",
+       "  'a white and black cup with a small spoon next to it\\n',\n",
+       "  'a white cup with a small bottle of wine\\n',\n",
+       "  'a white cup with a spoon and a bottle of wine\\n'],\n",
+       " ['a group of people playing baseball and soccer\\n',\n",
+       "  'a group of people are playing baseball in the grass\\n',\n",
+       "  'a group of people playing baseball and running\\n',\n",
+       "  'a group of people playing baseball and soccer\\n',\n",
+       "  'a group of people playing soccer on a field\\n',\n",
+       "  'a group of people are playing baseball in the grass\\n',\n",
+       "  'a group of people playing baseball with a man in the background\\n',\n",
+       "  'a group of people playing baseball and one is holding a ball\\n',\n",
+       "  'a group of people playing baseball in front of a field\\n',\n",
+       "  'a group of people playing baseball on a field\\n',\n",
+       "  'a group of people playing baseball with one person in the background\\n',\n",
+       "  'a group of people are playing baseball and one is holding a ball\\n']]"
+      ]
+     },
+     "execution_count": 100,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "several = sample_several(predicted_embeddings[0:8], num = 12, temp = 0.3)\n",
+    "several"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "93e87fde-815d-4452-9915-f5f5dacf7c2a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "plt.plot(losses)\n",
+    "plt.show()\n",
+    "plt.plot(test_losses)\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ccfccd4f-764d-4624-842c-f931676eb43b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print('test')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f1a60e19-c440-4c9c-a634-30186209012f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def tensor_2_embed_old(tensor):\n",
+    "    embed_array = torch.zeros((tensor.shape[0],257, 1024)) \n",
+    "    to_pil = ToPILImage()\n",
+    "    for sample in range(tensor.shape[0]):\n",
+    "        PIL_image = to_pil(tensor[sample])\n",
+    "        image_for_blip2 = vis_processors[\"eval\"](PIL_image).unsqueeze(0).to(device)\n",
+    "        #Generate embeddings\n",
+    "        with blip2_model.maybe_autocast():\n",
+    "                blip2_target = blip2_model.ln_vision(blip2_model.visual_encoder(image_for_blip2))\n",
+    "                embed_array[sample] = blip2_target\n",
+    "                \n",
+    "    return embed_array"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d39ddada-47f7-4111-92fa-0dd98e8a83d6",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ec8ed96a-61fa-4c20-8da2-fcd9d0a2ed38",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6228eb1a-e8e7-4500-b7bc-d0c57bcac4c6",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.8"
+  },
+  "toc": {
+   "base_numbering": 1,
+   "nav_menu": {},
+   "number_sections": true,
+   "sideBar": true,
+   "skip_h1_title": false,
+   "title_cell": "Table of Contents",
+   "title_sidebar": "Contents",
+   "toc_cell": false,
+   "toc_position": {
+    "height": "calc(100% - 180px)",
+    "left": "10px",
+    "top": "150px",
+    "width": "165px"
+   },
+   "toc_section_display": true,
+   "toc_window_display": true
+  },
+  "toc-autonumbering": true,
+  "vscode": {
+   "interpreter": {
+    "hash": "62aae01ef0cf7b6af841ab1c8ce59175c4332e693ab3d00bc32ceffb78a35376"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/train2.ipynb

@@ -0,0 +1,1856 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "f4d95fac-ac1d-473c-ab96-650f76e6aaf5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# # Code to convert this notebook to .py if you want to run it via command line or with Slurm\n",
+    "# from subprocess import call\n",
+    "# command = \"jupyter nbconvert Train.ipynb --to python\"\n",
+    "# call(command,shell=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b0f0f4f3",
+   "metadata": {},
+   "source": [
+    "# Import packages & functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "5bad764b-45c1-45ce-a716-8d055e09821a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2023-11-19 16:32:39,711] [INFO] [real_accelerator.py:110:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import sys\n",
+    "import json\n",
+    "import argparse\n",
+    "import numpy as np\n",
+    "import math\n",
+    "from einops import rearrange\n",
+    "import time\n",
+    "import random\n",
+    "import h5py\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "import webdataset as wds\n",
+    "import gc\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torchvision import transforms\n",
+    "from torchvision.transforms import ToPILImage   #CHANGED (added)\n",
+    "\n",
+    "from accelerate import Accelerator, DeepSpeedPlugin\n",
+    "\n",
+    "# tf32 data type is faster than standard float32\n",
+    "torch.backends.cuda.matmul.allow_tf32 = True\n",
+    "\n",
+    "# custom functions #\n",
+    "import utils\n",
+    "\n",
+    "global_batch_size = 128 #128"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "cc5d2e32-6027-4a19-bef4-5ca068db35bb",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "LOCAL RANK  0\n"
+     ]
+    }
+   ],
+   "source": [
+    "### Multi-GPU config ###\n",
+    "local_rank = os.getenv('RANK')\n",
+    "if local_rank is None: \n",
+    "    local_rank = 0\n",
+    "else:\n",
+    "    local_rank = int(local_rank)\n",
+    "print(\"LOCAL RANK \", local_rank)  \n",
+    "\n",
+    "num_devices = torch.cuda.device_count()\n",
+    "if num_devices==0: num_devices = 1\n",
+    "\n",
+    "accelerator = Accelerator(split_batches=False)\n",
+    "\n",
+    "### UNCOMMENT BELOW STUFF TO USE DEEPSPEED (also comment out the immediately above \"accelerator = \" line) ###\n",
+    "\n",
+    "# if num_devices <= 1 and utils.is_interactive():\n",
+    "#     # can emulate a distributed environment for deepspeed to work in jupyter notebook\n",
+    "#     os.environ[\"MASTER_ADDR\"] = \"localhost\"\n",
+    "#     os.environ[\"MASTER_PORT\"] = str(np.random.randint(10000)+9000)\n",
+    "#     os.environ[\"RANK\"] = \"0\"\n",
+    "#     os.environ[\"LOCAL_RANK\"] = \"0\"\n",
+    "#     os.environ[\"WORLD_SIZE\"] = \"1\"\n",
+    "#     os.environ[\"GLOBAL_BATCH_SIZE\"] = str(global_batch_size) # set this to your batch size!\n",
+    "#     global_batch_size = os.environ[\"GLOBAL_BATCH_SIZE\"]\n",
+    "\n",
+    "# # alter the deepspeed config according to your global and local batch size\n",
+    "# if local_rank == 0:\n",
+    "#     with open('deepspeed_config_stage2.json', 'r') as file:\n",
+    "#         config = json.load(file)\n",
+    "#     config['train_batch_size'] = int(os.environ[\"GLOBAL_BATCH_SIZE\"])\n",
+    "#     config['train_micro_batch_size_per_gpu'] = int(os.environ[\"GLOBAL_BATCH_SIZE\"]) // num_devices\n",
+    "#     with open('deepspeed_config_stage2.json', 'w') as file:\n",
+    "#         json.dump(config, file)\n",
+    "# else:\n",
+    "#     # give some time for the local_rank=0 gpu to prep new deepspeed config file\n",
+    "#     time.sleep(10)\n",
+    "# deepspeed_plugin = DeepSpeedPlugin(\"deepspeed_config_stage2.json\")\n",
+    "# accelerator = Accelerator(split_batches=False, deepspeed_plugin=deepspeed_plugin)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "b767ab6f-d4a9-47a5-b3bf-f56bf6760c0c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "PID of this process = 2370606\n",
+      "device: cuda\n",
+      "Distributed environment: NO\n",
+      "Num processes: 1\n",
+      "Process index: 0\n",
+      "Local process index: 0\n",
+      "Device: cuda\n",
+      "\n",
+      "Mixed precision type: no\n",
+      "\n",
+      "distributed = False num_devices = 1 local rank = 0 world size = 1\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(\"PID of this process =\",os.getpid())\n",
+    "device = accelerator.device\n",
+    "print(\"device:\",device)\n",
+    "num_workers = num_devices\n",
+    "print(accelerator.state)\n",
+    "world_size = accelerator.state.num_processes\n",
+    "distributed = not accelerator.state.distributed_type == 'NO'\n",
+    "print(\"distributed =\",distributed, \"num_devices =\", num_devices, \"local rank =\", local_rank, \"world size =\", world_size)\n",
+    "print = accelerator.print # only print if local_rank=0"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9018b82b-c054-4463-9527-4b0c2a75bda6",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "# Configurations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "2b61fec7-72a0-4b67-86da-1375f1d9fbd3",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "['--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset', '--model_name=captions', '--subj=1', '--batch_size=128', '--n_samples_save=0', '--max_lr=3e-1', '--mixup_pct=.66', '--num_epochs=30', '--ckpt_interval=999', '--no-use_image_aug']\n"
+     ]
+    }
+   ],
+   "source": [
+    "# if running this interactively, can specify jupyter_args here for argparser to use\n",
+    "if utils.is_interactive():\n",
+    "    # Example use\n",
+    "    jupyter_args = f\"--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset \\\n",
+    "                    --model_name=captions \\\n",
+    "                    --subj=1 --batch_size={global_batch_size} --n_samples_save=0 \\\n",
+    "                    --max_lr=3e-1 --mixup_pct=.66 --num_epochs=30 --ckpt_interval=999 --no-use_image_aug\"\n",
+    "    #max_lr=3e-5 originally\n",
+    "    jupyter_args = jupyter_args.split()\n",
+    "    print(jupyter_args)\n",
+    "    \n",
+    "    from IPython.display import clear_output # function to clear print outputs in cell\n",
+    "    %load_ext autoreload \n",
+    "    # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions\n",
+    "    %autoreload 2 "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "2028bdf0-2f41-46d9-b6e7-86b870dbf16c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "global batch_size 128\n",
+      "batch_size 128\n"
+     ]
+    }
+   ],
+   "source": [
+    "parser = argparse.ArgumentParser(description=\"Model Training Configuration\")\n",
+    "parser.add_argument(\n",
+    "    \"--model_name\", type=str, default=\"testing\",\n",
+    "    help=\"name of model, used for ckpt saving and wandb logging (if enabled)\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--data_path\", type=str, default=\"/fsx/proj-fmri/shared/natural-scenes-dataset\",\n",
+    "    help=\"Path to where NSD data is stored / where to download it to\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--subj\",type=int, default=1, choices=[1,2,5,7],\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--batch_size\", type=int, default=32,\n",
+    "    help=\"Batch size can be increased by 10x if only training v2c and not diffusion diffuser\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--wandb_log\",action=argparse.BooleanOptionalAction,default=False,\n",
+    "    help=\"whether to log to wandb\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--resume_from_ckpt\",action=argparse.BooleanOptionalAction,default=False,\n",
+    "    help=\"if not using wandb and want to resume from a ckpt\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--wandb_project\",type=str,default=\"stability\",\n",
+    "    help=\"wandb project name\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--mixup_pct\",type=float,default=.33,\n",
+    "    help=\"proportion of way through training when to switch from BiMixCo to SoftCLIP\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--use_image_aug\",action=argparse.BooleanOptionalAction,default=True,\n",
+    "    help=\"whether to use image augmentation\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--num_epochs\",type=int,default=240,\n",
+    "    help=\"number of epochs of training\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--lr_scheduler_type\",type=str,default='cycle',choices=['cycle','linear'],\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--ckpt_saving\",action=argparse.BooleanOptionalAction,default=True,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--ckpt_interval\",type=int,default=5,\n",
+    "    help=\"save backup ckpt and reconstruct every x epochs\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--seed\",type=int,default=42,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--max_lr\",type=float,default=3e-4,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--n_samples_save\",type=int,default=0,choices=[0,1],\n",
+    "    help=\"Number of reconstructions for monitoring progress, 0 will speed up training\",\n",
+    ")\n",
+    "\n",
+    "if utils.is_interactive():\n",
+    "    args = parser.parse_args(jupyter_args)\n",
+    "else:\n",
+    "    args = parser.parse_args()\n",
+    "\n",
+    "# create global variables without the args prefix\n",
+    "for attribute_name in vars(args).keys():\n",
+    "    globals()[attribute_name] = getattr(args, attribute_name)\n",
+    "\n",
+    "print(\"global batch_size\", batch_size)\n",
+    "batch_size = int(batch_size / num_devices)\n",
+    "print(\"batch_size\", batch_size)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "60cd7f2c-37fd-426b-a0c6-633e51bc4c4d",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "outdir = os.path.abspath(f'../train_logs/{model_name}')\n",
+    "if not os.path.exists(outdir):\n",
+    "    os.makedirs(outdir,exist_ok=True)\n",
+    "if use_image_aug:\n",
+    "    import kornia\n",
+    "    from kornia.augmentation.container import AugmentationSequential\n",
+    "    img_augment = AugmentationSequential(\n",
+    "        kornia.augmentation.RandomResizedCrop((224,224), (0.6,1), p=0.3),\n",
+    "        kornia.augmentation.Resize((224, 224)),\n",
+    "        kornia.augmentation.RandomHorizontalFlip(p=0.3),\n",
+    "        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.3),\n",
+    "        kornia.augmentation.RandomGrayscale(p=0.3),\n",
+    "        same_on_batch=False,\n",
+    "        data_keys=[\"input\"],\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "e7807ba9-02b6-4bc0-873c-69869abe4091",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "wandb_log = False"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "42d13c25-1369-4c49-81d4-83d713586096",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "# Prep data, models, and dataloaders"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1c023f24-5233-4a15-a2f5-78487b3a8546",
+   "metadata": {},
+   "source": [
+    "## Dataloader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "81084834-035f-4465-ad59-59e6b806a2f5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/train/{0..36}.tar\n",
+      "/fsx/proj-fmri/shared/mindeyev2_dataset/wds/subj01/test/0.tar\n"
+     ]
+    }
+   ],
+   "source": [
+    "if subj==1:\n",
+    "    num_train = 24958\n",
+    "    num_test = 2770\n",
+    "test_batch_size = num_test\n",
+    "\n",
+    "def my_split_by_node(urls): return urls\n",
+    "    \n",
+    "train_url = f\"{data_path}/wds/subj0{subj}/train/\" + \"{0..36}.tar\"\n",
+    "print(train_url)\n",
+    "\n",
+    "train_data = wds.WebDataset(train_url,resampled=False,nodesplitter=my_split_by_node)\\\n",
+    "                    .shuffle(750, initial=1500, rng=random.Random(42))\\\n",
+    "                    .decode(\"torch\")\\\n",
+    "                    .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n",
+    "                    .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n",
+    "train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)\n",
+    "\n",
+    "test_url = f\"{data_path}/wds/subj0{subj}/test/\" + \"0.tar\"\n",
+    "print(test_url)\n",
+    "\n",
+    "test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\\\n",
+    "                    .shuffle(750, initial=1500, rng=random.Random(42))\\\n",
+    "                    .decode(\"torch\")\\\n",
+    "                    .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n",
+    "                    .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n",
+    "test_dl = torch.utils.data.DataLoader(test_data, batch_size=test_batch_size, shuffle=False, drop_last=False, pin_memory=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "203b060a-2dd2-4c35-929b-c576be82eb52",
+   "metadata": {},
+   "source": [
+    "### check dataloaders are working"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "e7a9c68c-c3c9-4080-bd99-067c4486dc37",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# test_indices = []\n",
+    "# test_images = []\n",
+    "# for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):\n",
+    "#     test_indices = np.append(test_indices, behav[:,0,5].numpy())\n",
+    "#     test_images = np.append(test_images, behav[:,0,0].numpy())\n",
+    "# test_indices = test_indices.astype(np.int16)\n",
+    "# print(test_i, (test_i+1) * test_batch_size, len(test_indices))\n",
+    "# print(\"---\\n\")\n",
+    "\n",
+    "# train_indices = []\n",
+    "# train_images = []\n",
+    "# for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):\n",
+    "#     train_indices = np.append(train_indices, behav[:,0,5].long().numpy())\n",
+    "#     train_images = np.append(train_images, behav[:,0,0].numpy())\n",
+    "# train_indices = train_indices.astype(np.int16)\n",
+    "# print(train_i, (train_i+1) * batch_size, len(train_indices))\n",
+    "\n",
+    "# # train_images = np.hstack((train_images, test_images))\n",
+    "# # print(\"WARNING: ADDED TEST IMAGES TO TRAIN IMAGES\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "45fad12c-f9fb-4408-8fd4-9bca324ad634",
+   "metadata": {},
+   "source": [
+    "## Load data and images"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "039dd330-7339-4f88-8f00-45f95e47baa0",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "subj01 betas loaded into memory\n",
+      "voxels torch.Size([27750, 15729])\n",
+      "images torch.Size([73000, 3, 224, 224])\n"
+     ]
+    }
+   ],
+   "source": [
+    "# load betas\n",
+    "f = h5py.File(f'{data_path}/betas_all_subj0{subj}.hdf5', 'r')\n",
+    "voxels = f['betas'][:]\n",
+    "print(f\"subj0{subj} betas loaded into memory\")\n",
+    "voxels = torch.Tensor(voxels).to(\"cpu\").half()\n",
+    "if subj==1:\n",
+    "    voxels = torch.hstack((voxels, torch.zeros((len(voxels), 5))))\n",
+    "print(\"voxels\", voxels.shape)\n",
+    "num_voxels = voxels.shape[-1]\n",
+    "\n",
+    "# load orig images\n",
+    "f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')\n",
+    "images = f['images'][:]\n",
+    "images = torch.Tensor(images).to(\"cpu\").half()\n",
+    "print(\"images\", images.shape)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "10ec4517-dbdf-4ece-98f6-4714d5de4e15",
+   "metadata": {},
+   "source": [
+    "## Load models"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "48d6160e-1ee8-4da7-a755-9dbb452a6fa5",
+   "metadata": {},
+   "source": [
+    "### CLIP image embeddings  model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "795e2885-bd07-4e27-bed7-181473c06df9",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import transformers\n",
+    "from transformers import Blip2Processor, Blip2ForConditionalGeneration\n",
+    "\n",
+    "from PIL import Image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "b0420dc0-199e-4c1a-857d-b1747058b467",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "ViT-L/14 cuda:0\n"
+     ]
+    }
+   ],
+   "source": [
+    "from models import Clipper\n",
+    "clip_model = Clipper(\"ViT-L/14\", device=torch.device(f\"cuda:{local_rank}\"), hidden_state=True, norm_embs=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "23428fb7-2955-4295-bea1-447cebf9f72e",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [01:08<00:00, 34.47s/it]\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'from lavis.models import load_model_and_preprocess\\nfrom lavis.models import model_zoo\\nblip2_model, vis_processors, _ = load_model_and_preprocess(\\n            name=\"blip2_t5\", model_type=\"pretrain_flant5xl_vitL\", is_eval=True, device=device)\\n\\nclip_seq_dim = 257\\nclip_emb_dim = 1024\\nhidden_dim = 4096'"
+      ]
+     },
+     "execution_count": 14,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "cache_blip2 = \"/fsx/proj-fmri/shared/cache/models--Salesforce--blip2-opt-2.7b/snapshots/6e723d92ee91ebcee4ba74d7017632f11ff4217b\"\n",
+    "\n",
+    "b2_processor = Blip2Processor.from_pretrained(cache_blip2)\n",
+    "b2_model = Blip2ForConditionalGeneration.from_pretrained(cache_blip2, torch_dtype=torch.float16, device_map=\"auto\")\n",
+    "\n",
+    "#Load in blip2 as well\n",
+    "\"\"\"from lavis.models import load_model_and_preprocess\n",
+    "from lavis.models import model_zoo\n",
+    "blip2_model, vis_processors, _ = load_model_and_preprocess(\n",
+    "            name=\"blip2_t5\", model_type=\"pretrain_flant5xl_vitL\", is_eval=True, device=device)\n",
+    "\n",
+    "clip_seq_dim = 257\n",
+    "clip_emb_dim = 1024\n",
+    "hidden_dim = 4096\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "b06f3de2-a8da-4ba0-94f0-99096f738d55",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "def embed_images_b2(images):\n",
+    "    images = (images * 255).type(torch.uint8)\n",
+    "    with torch.no_grad():\n",
+    "        inputs_processed = b2_processor(images, return_tensors=\"pt\").to(\"cuda\", torch.float16)\n",
+    "        enc_imgs = b2_model.vision_model.forward(inputs_processed['pixel_values'])\n",
+    "        return enc_imgs.last_hidden_state.detach(), inputs_processed\n",
+    "\n",
+    "def embeds_to_captions_b2(embeds):\n",
+    "    with torch.no_grad():\n",
+    "        input_ids = None #inputs['input_ids']\n",
+    "        attention_mask = None\n",
+    "        batch_size = embeds.shape[0]\n",
+    "        image_embeds = embeds\n",
+    "        image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device)\n",
+    "\n",
+    "        query_tokens = b2_model.query_tokens.expand(image_embeds.shape[0], -1, -1)\n",
+    "        query_outputs = b2_model.qformer(\n",
+    "            query_embeds=query_tokens,\n",
+    "            encoder_hidden_states=image_embeds,\n",
+    "            encoder_attention_mask=image_attention_mask,\n",
+    "            return_dict=True,\n",
+    "        )\n",
+    "        query_output = query_outputs.last_hidden_state\n",
+    "\n",
+    "        language_model_inputs = b2_model.language_projection(query_output)\n",
+    "        language_attention_mask = torch.ones(\n",
+    "            language_model_inputs.size()[:-1], dtype=torch.long, device=language_model_inputs.device\n",
+    "        )\n",
+    "        if input_ids is None:\n",
+    "            input_ids = (\n",
+    "                torch.LongTensor([[b2_model.config.text_config.bos_token_id]])\n",
+    "                .repeat(batch_size, 1)\n",
+    "                .to(image_embeds.device)\n",
+    "            )\n",
+    "        if attention_mask is None:\n",
+    "            attention_mask = torch.ones_like(input_ids)\n",
+    "        attention_mask = torch.cat([language_attention_mask, attention_mask.to(language_attention_mask.device)], dim=1)\n",
+    "\n",
+    "        # concatenate query embeddings with prompt embeddings\n",
+    "        inputs_embeds = b2_model.get_input_embeddings()(input_ids)\n",
+    "        inputs_embeds = torch.cat([language_model_inputs, inputs_embeds.to(language_model_inputs.device)], dim=1)\n",
+    "\n",
+    "        outputs = b2_model.language_model.generate(\n",
+    "            inputs_embeds=inputs_embeds,\n",
+    "            attention_mask=attention_mask,\n",
+    "        )\n",
+    "        text = b2_processor.batch_decode(outputs, skip_special_tokens=True)\n",
+    "        \n",
+    "        return outputs, text\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "ec0a34d3-76e0-4a47-a9ab-6131ab2ccecd",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "image_test = images[1:20].permute(0,2,3,1)\n",
+    "#raw_image = Image.open('/fsx/proj-fmri/shared/controlNetData/target/img_t1.jpg').convert('RGB')\n",
+    "# Convert the image to a NumPy array\n",
+    "#image_test = np.array(raw_image)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "e04876a4-45c7-4015-8255-8574c8f50f14",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "\"import matplotlib.pyplot as plt\\n# Plotting one of the images (taking the first image as an example)\\nimg_to_plot = inputs_rec['pixel_values'][-1]\\n\\n# Transpose the image for correct display (PyTorch: [C, H, W], Matplotlib: [H, W, C])\\nimg_to_plot = img_to_plot.permute(1, 2, 0).to(torch.float32).to('cpu')\\nprint(img_to_plot.shape)\\n\\nplt.imshow(img_to_plot)\\nplt.show()\""
+      ]
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "\"\"\"import matplotlib.pyplot as plt\n",
+    "# Plotting one of the images (taking the first image as an example)\n",
+    "img_to_plot = inputs_rec['pixel_values'][-1]\n",
+    "\n",
+    "# Transpose the image for correct display (PyTorch: [C, H, W], Matplotlib: [H, W, C])\n",
+    "img_to_plot = img_to_plot.permute(1, 2, 0).to(torch.float32).to('cpu')\n",
+    "print(img_to_plot.shape)\n",
+    "\n",
+    "plt.imshow(img_to_plot)\n",
+    "plt.show()\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "328a17d0-593b-4d1e-812a-10a3b6efea6a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "embeds_test, inputs_rec = embed_images_b2(image_test)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "abe5f8a8-fca9-4083-8596-a913bdb57de7",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "#inputs_rec['pixel_values'].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "c5f3ca7e-b880-421e-b354-7b6c3df565e9",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "#out = b2_model.generate(**inputs_rec)\n",
+    "#print(b2_processor.decode(out[0], skip_special_tokens=True).strip())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "fb462016-78d7-46ea-8058-0d608f17ea65",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/admin/home-ckadirt/miniconda3/envs/mindeye/lib/python3.10/site-packages/transformers/generation/utils.py:1260: UserWarning: Using the model-agnostic default `max_length` (=20) to control thegeneration length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
+      "  warnings.warn(\n"
+     ]
+    }
+   ],
+   "source": [
+    "outputs_test, text_test = embeds_to_captions_b2(embeds_test)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "6a95fcdf-db87-4c02-9728-09f85605fb1c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['a cat sitting on a toilet seat\\n',\n",
+       " 'a person cutting a pizza on a cutting board\\n',\n",
+       " 'a sandwich and a drink on a table\\n',\n",
+       " 'a man crossing the street in front of a truck\\n',\n",
+       " 'a giraffe standing in front of trees\\n',\n",
+       " 'three men standing together\\n',\n",
+       " 'a bird standing on a rock next to a body of water\\n',\n",
+       " 'two men sitting on a street corner in asia\\n',\n",
+       " 'a woman and two children playing tennis on a court\\n',\n",
+       " 'a tall brick building with a clock on the side\\n',\n",
+       " 'a train is on the tracks\\n',\n",
+       " 'a man and woman in the water with a surfboard\\n',\n",
+       " 'a living room with a desk and a chair\\n',\n",
+       " 'a group of men on a basketball court\\n',\n",
+       " 'a man holding an umbrella\\n',\n",
+       " 'a man in a red shirt\\n',\n",
+       " 'a group of people holding cell phones and wine glasses\\n',\n",
+       " 'a laptop computer sitting on a table in front of a television\\n',\n",
+       " 'a baseball player is swinging a bat on a field\\n']"
+      ]
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "text_test"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "9ac69fbd-55db-435b-bed6-5ae9186450e3",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "#inputss['pixel_values'].shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "id": "0524f498-c8da-4e8a-8970-d75d2d0f6b8b",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "#image_test.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "5417541b-49eb-4e43-a3e2-d937d9653e04",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "max_lr = 1e-4"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "id": "da0ce190-1b3e-4c12-9e9f-91cbc076d044",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "clip_seq_dim = 257 #blip2 image encoder shapes\n",
+    "clip_emb_dim = 1408 #blip2 image encoder shapes\n",
+    "hidden_dim = 2048"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5b79bd38-6990-4504-8d45-4a68d57d8885",
+   "metadata": {},
+   "source": [
+    "### SD VAE (blurry images)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "id": "01baff79-8114-482b-b115-6f05aa8ad691",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "83,653,863 total\n",
+      "0 trainable\n"
+     ]
+    }
+   ],
+   "source": [
+    "from diffusers import AutoencoderKL\n",
+    "autoenc = AutoencoderKL.from_pretrained(\"madebyollin/sdxl-vae-fp16-fix\", torch_dtype=torch.float16, cache_dir=\"/fsx/proj-fmri/shared/cache\")\n",
+    "# autoenc.load_state_dict(torch.load('../train_logs/sdxl_vae_normed/best.pth')[\"model_state_dict\"])\n",
+    "autoenc.eval()\n",
+    "autoenc.requires_grad_(False)\n",
+    "autoenc.to(device)\n",
+    "utils.count_params(autoenc)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "260e5e4a-f697-4b2c-88fc-01f6a54886c0",
+   "metadata": {},
+   "source": [
+    "### MindEye modules"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "id": "c44c271b-173f-472e-b059-a2eda0f4c4c5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "MindEyeModule()"
+      ]
+     },
+     "execution_count": 28,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "class MindEyeModule(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super(MindEyeModule, self).__init__()\n",
+    "    def forward(self, x):\n",
+    "        return x\n",
+    "        \n",
+    "model = MindEyeModule()\n",
+    "model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "id": "038a5d61-4769-40b9-a004-f4e7b5b38bb0",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "32,215,040 total\n",
+      "32,215,040 trainable\n",
+      "param counts:\n",
+      "32,215,040 total\n",
+      "32,215,040 trainable\n",
+      "torch.Size([2, 1, 15729]) torch.Size([2, 1, 2048])\n"
+     ]
+    }
+   ],
+   "source": [
+    "class RidgeRegression(torch.nn.Module):\n",
+    "    # make sure to add weight_decay when initializing optimizer\n",
+    "    def __init__(self, input_size, out_features): \n",
+    "        super(RidgeRegression, self).__init__()\n",
+    "        self.out_features = out_features\n",
+    "        self.linear = torch.nn.Linear(input_size, out_features)\n",
+    "    def forward(self, x):\n",
+    "        return self.linear(x)\n",
+    "        \n",
+    "model.ridge = RidgeRegression(voxels.shape[1], out_features=hidden_dim)\n",
+    "utils.count_params(model.ridge)\n",
+    "utils.count_params(model)\n",
+    "\n",
+    "b = torch.randn((2,1,voxels.shape[1]))\n",
+    "print(b.shape, model.ridge(b).shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "id": "3602c333-d029-465c-8fb4-c3ccffdba6fd",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "772,419,072 total\n",
+      "772,419,072 trainable\n",
+      "param counts:\n",
+      "804,634,112 total\n",
+      "804,634,112 trainable\n",
+      "torch.Size([4, 2048])\n",
+      "torch.Size([4, 257, 1408])\n"
+     ]
+    }
+   ],
+   "source": [
+    "from functools import partial\n",
+    "from diffusers.models.vae import Decoder\n",
+    "class BrainNetwork(nn.Module):\n",
+    "    def __init__(self, out_dim=768, in_dim=15724, clip_size=768, h=4096, n_blocks=4, norm_type='ln', act_first=False, drop=.15, blurry_dim=16):\n",
+    "        super().__init__()\n",
+    "        self.blurry_dim = blurry_dim\n",
+    "        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)\n",
+    "        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU\n",
+    "        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)\n",
+    "        self.lin0 = nn.Linear(in_dim, h)\n",
+    "        self.mlp = nn.ModuleList([\n",
+    "            nn.Sequential(\n",
+    "                nn.Linear(h, h),\n",
+    "                *[item() for item in act_and_norm],\n",
+    "                nn.Dropout(drop)\n",
+    "            ) for _ in range(n_blocks)\n",
+    "        ])\n",
+    "        self.lin1 = nn.Linear(h, out_dim, bias=True)\n",
+    "        # self.blin1 = nn.Linear(out_dim, blurry_dim, bias=True)\n",
+    "        self.n_blocks = n_blocks\n",
+    "        self.clip_size = clip_size\n",
+    "        self.clip_proj = nn.Sequential(\n",
+    "            nn.LayerNorm(clip_size),\n",
+    "            nn.GELU(),\n",
+    "            nn.Linear(clip_size, 2048),\n",
+    "            nn.LayerNorm(2048),\n",
+    "            nn.GELU(),\n",
+    "            nn.Linear(2048, 2048),\n",
+    "            nn.LayerNorm(2048),\n",
+    "            nn.GELU(),\n",
+    "            nn.Linear(2048, clip_size)\n",
+    "        )\n",
+    "        # self.upsampler = Decoder(\n",
+    "        #         in_channels=64,\n",
+    "        #         out_channels=4,\n",
+    "        #         up_block_types=[\"UpDecoderBlock2D\",\"UpDecoderBlock2D\",\"UpDecoderBlock2D\"],\n",
+    "        #         block_out_channels=[64, 128, 256],\n",
+    "        #         layers_per_block=1,\n",
+    "        #     )\n",
+    "        \n",
+    "    def forward(self, x):\n",
+    "        x = self.lin0(x)\n",
+    "        residual = x\n",
+    "        for res_block in range(self.n_blocks):\n",
+    "            x = self.mlp[res_block](x)\n",
+    "            x += residual\n",
+    "            residual = x\n",
+    "        x = x.reshape(len(x), -1)\n",
+    "        x = self.lin1(x)\n",
+    "        # b = self.blin1(x)\n",
+    "        # b = self.upsampler(b.reshape(len(b), -1, 7, 7))\n",
+    "        c = self.clip_proj(x.reshape(len(x), -1, self.clip_size))\n",
+    "        # return c, b\n",
+    "        return c\n",
+    "\n",
+    "model.backbone = BrainNetwork(h=2048, in_dim=hidden_dim, clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, blurry_dim=64*7*7) \n",
+    "utils.count_params(model.backbone)\n",
+    "utils.count_params(model)\n",
+    "\n",
+    "b = torch.randn((4,hidden_dim))\n",
+    "print(b.shape)\n",
+    "clip_ = model.backbone(b)\n",
+    "print(clip_.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "id": "e14d0482-dc42-43b9-9ce1-953c32f2c9c1",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Done with model preparations!\n",
+      "param counts:\n",
+      "804,634,112 total\n",
+      "804,634,112 trainable\n"
+     ]
+    }
+   ],
+   "source": [
+    "no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n",
+    "opt_grouped_parameters = [\n",
+    "    {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},\n",
+    "    {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},\n",
+    "    {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},\n",
+    "]\n",
+    "\n",
+    "optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr, betas=(0.9, 0.95))\n",
+    "\n",
+    "if lr_scheduler_type == 'linear':\n",
+    "    lr_scheduler = torch.optim.lr_scheduler.LinearLR(\n",
+    "        optimizer,\n",
+    "        total_iters=int(num_epochs*(num_train*num_devices//batch_size)),\n",
+    "        last_epoch=-1\n",
+    "    )\n",
+    "elif lr_scheduler_type == 'cycle':\n",
+    "    total_steps=int(num_epochs*(num_train*num_devices//batch_size))\n",
+    "    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(\n",
+    "        optimizer, \n",
+    "        max_lr=max_lr,\n",
+    "        total_steps=total_steps,\n",
+    "        final_div_factor=1000,\n",
+    "        last_epoch=-1, pct_start=2/num_epochs\n",
+    "    )\n",
+    "    \n",
+    "def save_ckpt(tag):    \n",
+    "    ckpt_path = outdir+f'/{tag}.pth'\n",
+    "    print(f'saving {ckpt_path}',flush=True)\n",
+    "    unwrapped_model = accelerator.unwrap_model(model)\n",
+    "    try:\n",
+    "        torch.save({\n",
+    "            'epoch': epoch,\n",
+    "            'model_state_dict': unwrapped_model.state_dict(),\n",
+    "            'optimizer_state_dict': optimizer.state_dict(),\n",
+    "            'lr_scheduler': lr_scheduler.state_dict(),\n",
+    "            'train_losses': losses,\n",
+    "            'test_losses': test_losses,\n",
+    "            'lrs': lrs,\n",
+    "            }, ckpt_path)\n",
+    "    except:\n",
+    "        print(\"Couldn't save... moving on to prevent crashing.\")\n",
+    "    del unwrapped_model\n",
+    "        \n",
+    "print(\"\\nDone with model preparations!\")\n",
+    "utils.count_params(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "983f458b-35b8-49f2-b6db-80296cece730",
+   "metadata": {},
+   "source": [
+    "# Weights and Biases"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "id": "0a25a662-daa8-4de9-9233-8364800fcb6b",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "wandb mindeyev2 run captions\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mckadirt\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "wandb_config:\n",
+      " {'model_name': 'captions', 'batch_size': 128, 'num_epochs': 30, 'use_image_aug': False, 'max_lr': 0.0001, 'lr_scheduler_type': 'cycle', 'mixup_pct': 0.66, 'num_train': 24958, 'num_test': 2770, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1}\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "wandb version 0.16.0 is available!  To upgrade, please run:\n",
+       " $ pip install wandb --upgrade"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Tracking run with wandb version 0.15.5"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Run data is saved locally in <code>/fsx/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20231119_163615-o1xwsqre</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Syncing run <strong><a href='https://stability.wandb.io/ckadirt/mindeyev2/runs/o1xwsqre' target=\"_blank\">captions</a></strong> to <a href='https://stability.wandb.io/ckadirt/mindeyev2' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View project at <a href='https://stability.wandb.io/ckadirt/mindeyev2' target=\"_blank\">https://stability.wandb.io/ckadirt/mindeyev2</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View run at <a href='https://stability.wandb.io/ckadirt/mindeyev2/runs/o1xwsqre' target=\"_blank\">https://stability.wandb.io/ckadirt/mindeyev2/runs/o1xwsqre</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# params for wandb\n",
+    "if local_rank==0 and True: # only use main process for wandb logging\n",
+    "    import wandb\n",
+    "    \n",
+    "    wandb_project = 'mindeyev2'\n",
+    "    wandb_run = model_name\n",
+    "    wandb_notes = ''\n",
+    "    \n",
+    "    print(f\"wandb {wandb_project} run {wandb_run}\")\n",
+    "    wandb.login(host='https://stability.wandb.io')#, relogin=True)\n",
+    "    wandb_config = {\n",
+    "      \"model_name\": model_name,\n",
+    "      \"batch_size\": batch_size,\n",
+    "      \"num_epochs\": num_epochs,\n",
+    "      \"use_image_aug\": use_image_aug,\n",
+    "      \"max_lr\": max_lr,\n",
+    "      \"lr_scheduler_type\": lr_scheduler_type,\n",
+    "      \"mixup_pct\": mixup_pct,\n",
+    "      \"num_train\": num_train,\n",
+    "      \"num_test\": num_test,\n",
+    "      \"seed\": seed,\n",
+    "      \"distributed\": distributed,\n",
+    "      \"num_devices\": num_devices,\n",
+    "      \"world_size\": world_size,\n",
+    "    }\n",
+    "    print(\"wandb_config:\\n\",wandb_config)\n",
+    "    if False: # wandb_auto_resume\n",
+    "        print(\"wandb_id:\",model_name)\n",
+    "        wandb.init(\n",
+    "            id = model_name,\n",
+    "            project=wandb_project,\n",
+    "            name=wandb_run,\n",
+    "            config=wandb_config,\n",
+    "            notes=wandb_notes,\n",
+    "            resume=\"allow\",\n",
+    "        )\n",
+    "    else:\n",
+    "        wandb.init(\n",
+    "            project=wandb_project,\n",
+    "            name=wandb_run,\n",
+    "            config=wandb_config,\n",
+    "            notes=wandb_notes,\n",
+    "        )\n",
+    "else:\n",
+    "    wandb_log = False"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "id": "4e5de216-5318-4b45-ac02-113f03105adc",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #ff0000; text-decoration-color: #ff0000\">╭──────────────────────────────────────────────────────────────────────────────────────────────────╮</span>\n",
+       "<span style=\"color: #ff0000; text-decoration-color: #ff0000\">│</span> n++                                                                                              <span style=\"color: #ff0000; text-decoration-color: #ff0000\">│</span>\n",
+       "<span style=\"color: #ff0000; text-decoration-color: #ff0000\">│</span>    <span style=\"color: #ff0000; text-decoration-color: #ff0000; font-weight: bold\">▲</span>                                                                                             <span style=\"color: #ff0000; text-decoration-color: #ff0000\">│</span>\n",
+       "<span style=\"color: #ff0000; text-decoration-color: #ff0000\">╰──────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
+       "<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-weight: bold\">SyntaxError: </span>invalid syntax\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[91m╭──────────────────────────────────────────────────────────────────────────────────────────────────╮\u001b[0m\n",
+       "\u001b[91m│\u001b[0m n++                                                                                              \u001b[91m│\u001b[0m\n",
+       "\u001b[91m���\u001b[0m    \u001b[1;91m▲\u001b[0m                                                                                             \u001b[91m│\u001b[0m\n",
+       "\u001b[91m╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n",
+       "\u001b[1;91mSyntaxError: \u001b[0minvalid syntax\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "n++"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5b0ae095-3203-4eb8-8606-acc2db6ccf20",
+   "metadata": {},
+   "source": [
+    "# More custom functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "827ead88-7eb3-47cc-82da-31565063b927",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# using the same preprocessing as was used in MindEye + BrainDiffuser\n",
+    "pixcorr_preprocess = transforms.Compose([\n",
+    "    transforms.Resize(425, interpolation=transforms.InterpolationMode.BILINEAR),\n",
+    "])\n",
+    "def pixcorr(images,brains):\n",
+    "    # Flatten images while keeping the batch dimension\n",
+    "    all_images_flattened = pixcorr_preprocess(images).reshape(len(images), -1)\n",
+    "    all_brain_recons_flattened = pixcorr_preprocess(brains).view(len(brains), -1)\n",
+    "    corrmean = torch.diag(utils.batchwise_pearson_correlation(all_images_flattened, all_brain_recons_flattened)).mean()\n",
+    "    return corrmean"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d5690151-2131-4918-b750-e869cbd1a8a8",
+   "metadata": {},
+   "source": [
+    "# Main"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "12de6387-6e18-4e4b-b5ce-a847d625330a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "epoch = 0\n",
+    "losses, test_losses, lrs = [], [], []\n",
+    "best_test_loss = 1e9\n",
+    "soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))\n",
+    "\n",
+    "# Optionally resume from checkpoint #\n",
+    "if resume_from_ckpt:\n",
+    "    print(\"\\n---resuming from last.pth ckpt---\\n\")\n",
+    "    try:\n",
+    "        checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')\n",
+    "    except:\n",
+    "        print('last.pth failed... trying last_backup.pth')\n",
+    "        checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')\n",
+    "    epoch = checkpoint['epoch']\n",
+    "    print(\"Epoch\",epoch)\n",
+    "    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
+    "    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n",
+    "    diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])\n",
+    "    del checkpoint\n",
+    "elif wandb_log:\n",
+    "    if wandb.run.resumed:\n",
+    "        print(\"\\n---resuming from last.pth ckpt---\\n\")\n",
+    "        try:\n",
+    "            checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')\n",
+    "        except:\n",
+    "            print('last.pth failed... trying last_backup.pth')\n",
+    "            checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')\n",
+    "        epoch = checkpoint['epoch']\n",
+    "        print(\"Epoch\",epoch)\n",
+    "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
+    "        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n",
+    "        diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])\n",
+    "        del checkpoint\n",
+    "torch.cuda.empty_cache()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "99f09f76-4481-4133-b09a-a22b10dbc0c4",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "model, optimizer, train_dl, test_dl, lr_scheduler = accelerator.prepare(\n",
+    "model, optimizer, train_dl, test_dl, lr_scheduler\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bfeeda32-82ca-4364-bce1-eaa41b4f3e25",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "\"\"\"transform = transforms.Compose(\n",
+    "            [\n",
+    "                transforms.Resize(\n",
+    "                    (224, 224),\n",
+    "                ),\n",
+    "                transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),\n",
+    "            ]\n",
+    "        )\n",
+    "\n",
+    "def tensor_2_embed(image):    \n",
+    "    image_for_blip2 = transform(image)\n",
+    "    \n",
+    "    #Generate embeddings\n",
+    "    with blip2_model.maybe_autocast():\n",
+    "        blip2_target = blip2_model.ln_vision(blip2_model.visual_encoder(image_for_blip2))\n",
+    "                \n",
+    "    return blip2_target\n",
+    "\n",
+    "def embed_2_caption(image_embeds, model):\n",
+    "    image_embeds = image_embeds.float()\n",
+    "    image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(\n",
+    "                image.device)\n",
+    "\n",
+    "    query_tokens = model.query_tokens.expand(image_embeds.shape[0], -1, -1)\n",
+    "    query_output = model.Qformer.bert(\n",
+    "                query_embeds=query_tokens,\n",
+    "                encoder_hidden_states=image_embeds,\n",
+    "                encoder_attention_mask=image_atts,\n",
+    "                return_dict=True)\n",
+    "\n",
+    "    inputs_t5 = model.t5_proj(query_output.last_hidden_state)\n",
+    "    atts_t5 = torch.ones(inputs_t5.size()[:-1], dtype=torch.long).to(image.device)\n",
+    "    prompt = model.prompt\n",
+    "    input_tokens = model.t5_tokenizer(\n",
+    "                prompt, padding=\"longest\", return_tensors=\"pt\"\n",
+    "            ).to(image.device)\n",
+    "    encoder_atts = torch.cat([atts_t5, input_tokens.attention_mask], dim=1)\n",
+    "    \n",
+    "    with model.maybe_autocast(dtype=torch.bfloat16):\n",
+    "            inputs_embeds = model.t5_model.encoder.embed_tokens(input_tokens.input_ids)\n",
+    "            inputs_embeds = torch.cat([inputs_t5, inputs_embeds], dim=1)\n",
+    "\n",
+    "            outputs = model.t5_model.generate(\n",
+    "                inputs_embeds=inputs_embeds,\n",
+    "                attention_mask=encoder_atts)\n",
+    "            output_text = model.t5_tokenizer.batch_decode(\n",
+    "                outputs, skip_special_tokens=True)\n",
+    "    \n",
+    "    return output_text\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "636b4684-df9a-4e29-8683-86fb035ba690",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "wandb_log = True"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "60be0d5f-3e94-4612-9373-61b53d836393",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "print(f\"{model_name} starting with epoch {epoch} / {num_epochs}\")\n",
+    "progress_bar = tqdm(range(epoch,num_epochs), ncols=1200, disable=(local_rank!=0))\n",
+    "test_image, test_voxel = None, None\n",
+    "mse = nn.MSELoss()\n",
+    "for epoch in progress_bar:\n",
+    "    model.train()\n",
+    "    \n",
+    "    fwd_percent_correct = 0.\n",
+    "    bwd_percent_correct = 0.\n",
+    "    test_fwd_percent_correct = 0.\n",
+    "    test_bwd_percent_correct = 0.\n",
+    "\n",
+    "    loss_clip_total = 0.\n",
+    "    loss_blurry_total = 0.\n",
+    "    test_loss_clip_total = 0.\n",
+    "    test_loss_blurry_total = 0.\n",
+    "\n",
+    "    blurry_pixcorr = 0.\n",
+    "    test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1\n",
+    "    \n",
+    "    for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):\n",
+    "        if epoch == 0:\n",
+    "            lrs.append(0)\n",
+    "            break\n",
+    "        with torch.cuda.amp.autocast():\n",
+    "            optimizer.zero_grad()\n",
+    "\n",
+    "            voxel = voxels[behav[:,0,5].cpu().long()].to(device)\n",
+    "            \n",
+    "            image = images[behav[:,0,0].cpu().long()].to(device).float()\n",
+    "\n",
+    "            # blurry_image_enc = autoenc.encode(image).latent_dist.mode()\n",
+    "            \n",
+    "            if use_image_aug: image = img_augment(image)\n",
+    "            # clip_target = clip_model.embed_image(image)\n",
+    "            clip_target = embed_images_b2(image)[0].to(device) #####CHANGED\n",
+    "            assert not torch.any(torch.isnan(clip_target))\n",
+    "  \n",
+    "            if epoch < int(mixup_pct * num_epochs):\n",
+    "                voxel, perm, betas, select = utils.mixco(voxel)\n",
+    "\n",
+    "            voxel_ridge = model.ridge(voxel)\n",
+    "            \n",
+    "            # clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)\n",
+    "            clip_voxels = model.backbone(voxel_ridge)\n",
+    "            \n",
+    "            clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)\n",
+    "            clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)\n",
+    "\n",
+    "            if epoch < int(mixup_pct * num_epochs):                \n",
+    "                loss_clip = utils.mixco_nce(\n",
+    "                     clip_voxels_norm,\n",
+    "                     clip_target_norm,\n",
+    "                     temp=.006, \n",
+    "                     perm=perm, betas=betas, select=select)\n",
+    "            else:\n",
+    "                epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]\n",
+    "                loss_clip = utils.soft_clip_loss(\n",
+    "                     clip_voxels_norm,\n",
+    "                     clip_target_norm,\n",
+    "                     temp=epoch_temp)\n",
+    "            \n",
+    "            loss_mse= mse(clip_voxels, clip_target)\n",
+    "\n",
+    "            # loss_blurry = mse(blurry_image_enc_, blurry_image_enc) \n",
+    "\n",
+    "            loss_clip_total += loss_clip.item()\n",
+    "            # loss_blurry_total += loss_blurry.item()\n",
+    "\n",
+    "            # loss = loss_blurry + loss_clip\n",
+    "            loss = 0.7 * loss_clip + 0.3 * loss_mse\n",
+    "            if (train_i % 10 == 0):\n",
+    "                print(train_i, loss)\n",
+    "            # print(batch_size)\n",
+    "            utils.check_loss(loss)\n",
+    "            accelerator.backward(loss)\n",
+    "            optimizer.step()\n",
+    "    \n",
+    "            losses.append(loss.item())\n",
+    "            lrs.append(optimizer.param_groups[0]['lr'])\n",
+    "    \n",
+    "            # forward and backward top 1 accuracy        \n",
+    "            labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) \n",
+    "            fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)\n",
+    "            bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)\n",
+    "\n",
+    "            # with torch.no_grad():\n",
+    "            #     # only doing pixcorr eval on a subset (8) of the samples per batch because its costly & slow to compute autoenc.decode()\n",
+    "            #     random_samps = np.random.choice(np.arange(len(voxel)), size=8, replace=False)\n",
+    "            #     blurry_recon_images = autoenc.decode(blurry_image_enc_[random_samps]).sample.clamp(0,1)\n",
+    "            #     blurry_pixcorr += pixcorr(image[random_samps], blurry_recon_images)\n",
+    "\n",
+    "            if lr_scheduler_type is not None:\n",
+    "                lr_scheduler.step()\n",
+    "    \n",
+    "    model.eval()\n",
+    "    for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):\n",
+    "        with torch.cuda.amp.autocast():\n",
+    "            with torch.no_grad():   \n",
+    "                # all test samples should be loaded per batch such that test_i should never exceed 0\n",
+    "                if len(behav) != num_test: print(\"!\",len(behav),num_test)\n",
+    "                \n",
+    "                ## Average same-image repeats ##\n",
+    "                if test_image is None:\n",
+    "                    voxel = voxels[behav[:,0,5].cpu().long()].to(device)\n",
+    "                    \n",
+    "                    image = behav[:,0,0].cpu().long()\n",
+    "                    \n",
+    "                    unique_image, sort_indices = torch.unique(image, return_inverse=True)\n",
+    "                    for im in unique_image:\n",
+    "                        locs = torch.where(im == image)[0]\n",
+    "                        if test_image is None:\n",
+    "                            test_image = images[im][None]\n",
+    "                            test_voxel = torch.mean(voxel[locs],axis=0)[None]\n",
+    "                        else:\n",
+    "                            test_image = torch.vstack((test_image, images[im][None]))\n",
+    "                            test_voxel = torch.vstack((test_voxel, torch.mean(voxel[locs],axis=0)[None]))\n",
+    "    \n",
+    "                # sample of batch_size\n",
+    "                random_indices = torch.arange(len(test_voxel))[:batch_size] #torch.randperm(len(test_voxel))[:300]\n",
+    "                voxel = test_voxel[random_indices].to(device)\n",
+    "                image = test_image[random_indices].to(device)\n",
+    "                assert len(image) == batch_size\n",
+    "    \n",
+    "                # blurry_image_enc = autoenc.encode(image).latent_dist.mode()\n",
+    "        \n",
+    "                # clip_target = clip_model.embed_image(image.float())\n",
+    "                clip_target = embed_images_b2(image)[0].to(device) #####CHANGED\n",
+    "    \n",
+    "                voxel_ridge = model.ridge(voxel)\n",
+    "                \n",
+    "                # clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)\n",
+    "                clip_voxels = model.backbone(voxel_ridge)\n",
+    "                \n",
+    "                clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)\n",
+    "                clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)\n",
+    "        \n",
+    "                # loss_clip = utils.soft_clip_loss(\n",
+    "                #     clip_voxels_norm,\n",
+    "                #     clip_target_norm,\n",
+    "                #     temp=.006)\n",
+    "                \n",
+    "                loss_clip = mse(clip_voxels, clip_target)\n",
+    "\n",
+    "                # loss_blurry = mse(blurry_image_enc_, blurry_image_enc)\n",
+    "                \n",
+    "                # loss = loss_blurry + loss_clip\n",
+    "                loss = loss_clip\n",
+    "                \n",
+    "                utils.check_loss(loss)\n",
+    "        \n",
+    "                test_losses.append(loss.item())\n",
+    "        \n",
+    "                # forward and backward top 1 accuracy        \n",
+    "                labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) \n",
+    "                test_fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)\n",
+    "                test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)\n",
+    "\n",
+    "                # # halving the batch size because the decoder is computationally heavy\n",
+    "                # blurry_recon_images = autoenc.decode(blurry_image_enc_[:len(voxel)//2]).sample.clamp(0,1)\n",
+    "                # blurry_recon_images = torch.vstack((blurry_recon_images, autoenc.decode(blurry_image_enc_[len(voxel)//2:]).sample.clamp(0,1)))\n",
+    "                # test_blurry_pixcorr += pixcorr(image, blurry_recon_images)\n",
+    "\n",
+    "                #Find captions and print next to images\n",
+    "                #caption1 = embed_2_caption(clip_voxels[[0]], blip2_model)\n",
+    "                #caption2 = embed_2_caption(clip_voxels[[1]], blip2_model)\n",
+    "\n",
+    "                #true_embed1 = tensor_2_embed(image[[0]])\n",
+    "                #true_embed2 = tensor_2_embed(image[[1]])\n",
+    "\n",
+    "                # print(clip_voxels[[0]].shape)\n",
+    "                # print(true_embed1.shape)\n",
+    "                \n",
+    "                #true_caption1 = embed_2_caption(true_embed1, blip2_model)\n",
+    "                #true_caption2 = embed_2_caption(true_embed2, blip2_model)\n",
+    "                \n",
+    "                # transform blurry recon latents to images and plot it\n",
+    "                #fig, axes = plt.subplots(2, 2, figsize=(8, 4))\n",
+    "                #axes[0,0].imshow(utils.torch_to_Image(image[[0]]))\n",
+    "                #axes[0,1].imshow(utils.torch_to_Image(image[[1]]))\n",
+    "                #axes[0,0].axis('off'); axes[0,1].axis('off'); axes[1,0].axis('off'); axes[1,1].axis('off')\n",
+    "                #axes[0,0].set_title(caption1)\n",
+    "                #axes[0,1].set_title(caption2)\n",
+    "                #axes[1,0].set_title(true_caption1)\n",
+    "                #axes[1,1].set_title(true_caption2)\n",
+    "\n",
+    "                #plt.show()\n",
+    "                \n",
+    "                # # transform blurry recon latents to images and plot it\n",
+    "                # fig, axes = plt.subplots(1, 4, figsize=(8, 4))\n",
+    "                # axes[0].imshow(utils.torch_to_Image(image[[0]]))\n",
+    "                # axes[1].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[0]]).sample.clamp(0,1)))\n",
+    "                # axes[2].imshow(utils.torch_to_Image(image[[1]]))\n",
+    "                # axes[3].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[1]]).sample.clamp(0,1)))\n",
+    "                # axes[0].axis('off'); axes[1].axis('off'); axes[2].axis('off'); axes[3].axis('off')\n",
+    "                # axes[0].set_title(caption1)\n",
+    "                # axes[3].set_title(caption2)\n",
+    "                # plt.show()\n",
+    "                \n",
+    "\n",
+    "    if local_rank==0:      \n",
+    "        # if utils.is_interactive(): clear_output(wait=True)\n",
+    "        assert (test_i+1) == 1\n",
+    "        logs = {\"train/loss\": np.mean(losses[-(train_i+1):]),\n",
+    "            \"test/loss\": np.mean(test_losses[-(test_i+1):]),\n",
+    "            \"train/lr\": lrs[-1],\n",
+    "            \"train/num_steps\": len(losses),\n",
+    "            \"test/num_steps\": len(test_losses),\n",
+    "            \"train/fwd_pct_correct\": fwd_percent_correct / (train_i + 1),\n",
+    "            \"train/bwd_pct_correct\": bwd_percent_correct / (train_i + 1),\n",
+    "            \"test/test_fwd_pct_correct\": test_fwd_percent_correct / (test_i + 1),\n",
+    "            \"test/test_bwd_pct_correct\": test_bwd_percent_correct / (test_i + 1),\n",
+    "            \"train/loss_clip_total\": loss_clip_total / (train_i + 1),\n",
+    "            \"train/loss_blurry_total\": loss_blurry_total / (train_i + 1),\n",
+    "            \"test/loss_clip_total\": test_loss_clip_total / (test_i + 1),\n",
+    "            \"test/loss_blurry_total\": test_loss_blurry_total / (test_i + 1),\n",
+    "            \"train/blurry_pixcorr\": blurry_pixcorr / (train_i + 1),\n",
+    "            \"test/blurry_pixcorr\": test_blurry_pixcorr / (test_i + 1),\n",
+    "            }\n",
+    "        progress_bar.set_postfix(**logs)\n",
+    "                          \n",
+    "        fig, axes = plt.subplots(1, 8, figsize=(10, 4))\n",
+    "        jj=-1\n",
+    "        for j in [0,1,2,3,4,5,6,7]:\n",
+    "            jj+=1\n",
+    "            axes[jj].imshow(utils.torch_to_Image(image[j]))\n",
+    "            axes[jj].axis('off')\n",
+    "\n",
+    "        if wandb_log:\n",
+    "            generated_captions = embeds_to_captions_b2(clip_voxels[0:8])\n",
+    "            print(generated_captions[1])\n",
+    "            logs[f\"test/recons\"] = wandb.Image(fig, caption=f\"epoch{epoch:03d}\" + \"\\n\".join(generated_captions[1]))\n",
+    "            plt.close()\n",
+    "        # Save model checkpoint and reconstruct\n",
+    "        if epoch % ckpt_interval == 0:\n",
+    "            if not utils.is_interactive():\n",
+    "                save_ckpt(f'last')\n",
+    "                \n",
+    "        if wandb_log: wandb.log(logs)\n",
+    "\n",
+    "    # wait for other GPUs to catch up if needed\n",
+    "    accelerator.wait_for_everyone()\n",
+    "    torch.cuda.empty_cache()\n",
+    "    gc.collect()\n",
+    "\n",
+    "print(\"\\n===Finished!===\\n\")\n",
+    "if ckpt_saving:\n",
+    "    save_ckpt(f'last')\n",
+    "if not utils.is_interactive():\n",
+    "    sys.exit(0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "93e87fde-815d-4452-9915-f5f5dacf7c2a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "plt.plot(losses)\n",
+    "plt.show()\n",
+    "plt.plot(test_losses)\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ccfccd4f-764d-4624-842c-f931676eb43b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print('test')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f1a60e19-c440-4c9c-a634-30186209012f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def tensor_2_embed_old(tensor):\n",
+    "    embed_array = torch.zeros((tensor.shape[0],257, 1024)) \n",
+    "    to_pil = ToPILImage()\n",
+    "    for sample in range(tensor.shape[0]):\n",
+    "        PIL_image = to_pil(tensor[sample])\n",
+    "        image_for_blip2 = vis_processors[\"eval\"](PIL_image).unsqueeze(0).to(device)\n",
+    "        #Generate embeddings\n",
+    "        with blip2_model.maybe_autocast():\n",
+    "                blip2_target = blip2_model.ln_vision(blip2_model.visual_encoder(image_for_blip2))\n",
+    "                embed_array[sample] = blip2_target\n",
+    "                \n",
+    "    return embed_array"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d39ddada-47f7-4111-92fa-0dd98e8a83d6",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ec8ed96a-61fa-4c20-8da2-fcd9d0a2ed38",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6228eb1a-e8e7-4500-b7bc-d0c57bcac4c6",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.8"
+  },
+  "toc": {
+   "base_numbering": 1,
+   "nav_menu": {},
+   "number_sections": true,
+   "sideBar": true,
+   "skip_h1_title": false,
+   "title_cell": "Table of Contents",
+   "title_sidebar": "Contents",
+   "toc_cell": false,
+   "toc_position": {
+    "height": "calc(100% - 180px)",
+    "left": "10px",
+    "top": "150px",
+    "width": "165px"
+   },
+   "toc_section_display": true,
+   "toc_window_display": true
+  },
+  "toc-autonumbering": true,
+  "vscode": {
+   "interpreter": {
+    "hash": "62aae01ef0cf7b6af841ab1c8ce59175c4332e693ab3d00bc32ceffb78a35376"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/train2.py

@@ -0,0 +1,1141 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+# # Code to convert this notebook to .py if you want to run it via command line or with Slurm
+# from subprocess import call
+# command = "jupyter nbconvert Train.ipynb --to python"
+# call(command,shell=True)
+
+
+# # Import packages & functions
+
+# In[2]:
+
+
+import os
+import sys
+import json
+import argparse
+import numpy as np
+import math
+from einops import rearrange
+import time
+import random
+import h5py
+from tqdm import tqdm
+
+import webdataset as wds
+import gc
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from torchvision.transforms import ToPILImage   #CHANGED (added)
+
+from accelerate import Accelerator, DeepSpeedPlugin
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+
+global_batch_size = 128 #128
+
+
+# In[3]:
+
+
+### Multi-GPU config ###
+local_rank = os.getenv('RANK')
+if local_rank is None: 
+    local_rank = 0
+else:
+    local_rank = int(local_rank)
+print("LOCAL RANK ", local_rank)  
+
+num_devices = torch.cuda.device_count()
+if num_devices==0: num_devices = 1
+
+accelerator = Accelerator(split_batches=False)
+
+### UNCOMMENT BELOW STUFF TO USE DEEPSPEED (also comment out the immediately above "accelerator = " line) ###
+
+# if num_devices <= 1 and utils.is_interactive():
+#     # can emulate a distributed environment for deepspeed to work in jupyter notebook
+#     os.environ["MASTER_ADDR"] = "localhost"
+#     os.environ["MASTER_PORT"] = str(np.random.randint(10000)+9000)
+#     os.environ["RANK"] = "0"
+#     os.environ["LOCAL_RANK"] = "0"
+#     os.environ["WORLD_SIZE"] = "1"
+#     os.environ["GLOBAL_BATCH_SIZE"] = str(global_batch_size) # set this to your batch size!
+#     global_batch_size = os.environ["GLOBAL_BATCH_SIZE"]
+
+# # alter the deepspeed config according to your global and local batch size
+# if local_rank == 0:
+#     with open('deepspeed_config_stage2.json', 'r') as file:
+#         config = json.load(file)
+#     config['train_batch_size'] = int(os.environ["GLOBAL_BATCH_SIZE"])
+#     config['train_micro_batch_size_per_gpu'] = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices
+#     with open('deepspeed_config_stage2.json', 'w') as file:
+#         json.dump(config, file)
+# else:
+#     # give some time for the local_rank=0 gpu to prep new deepspeed config file
+#     time.sleep(10)
+# deepspeed_plugin = DeepSpeedPlugin("deepspeed_config_stage2.json")
+# accelerator = Accelerator(split_batches=False, deepspeed_plugin=deepspeed_plugin)
+
+
+# In[4]:
+
+
+print("PID of this process =",os.getpid())
+device = accelerator.device
+print("device:",device)
+num_workers = num_devices
+print(accelerator.state)
+world_size = accelerator.state.num_processes
+distributed = not accelerator.state.distributed_type == 'NO'
+print("distributed =",distributed, "num_devices =", num_devices, "local rank =", local_rank, "world size =", world_size)
+print = accelerator.print # only print if local_rank=0
+
+
+# # Configurations
+
+# In[5]:
+
+
+# if running this interactively, can specify jupyter_args here for argparser to use
+if utils.is_interactive():
+    # Example use
+    jupyter_args = f"--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset \
+                    --model_name=captions \
+                    --subj=1 --batch_size={global_batch_size} --n_samples_save=0 \
+                    --max_lr=3e-1 --mixup_pct=.66 --num_epochs=30 --ckpt_interval=999 --no-use_image_aug"
+    #max_lr=3e-5 originally
+    jupyter_args = jupyter_args.split()
+    print(jupyter_args)
+    
+    from IPython.display import clear_output # function to clear print outputs in cell
+    get_ipython().run_line_magic('load_ext', 'autoreload')
+    # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions
+    get_ipython().run_line_magic('autoreload', '2')
+
+
+# In[6]:
+
+
+parser = argparse.ArgumentParser(description="Model Training Configuration")
+parser.add_argument(
+    "--model_name", type=str, default="testing",
+    help="name of model, used for ckpt saving and wandb logging (if enabled)",
+)
+parser.add_argument(
+    "--data_path", type=str, default="/fsx/proj-fmri/shared/natural-scenes-dataset",
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,5,7],
+)
+parser.add_argument(
+    "--batch_size", type=int, default=32,
+    help="Batch size can be increased by 10x if only training v2c and not diffusion diffuser",
+)
+parser.add_argument(
+    "--wandb_log",action=argparse.BooleanOptionalAction,default=False,
+    help="whether to log to wandb",
+)
+parser.add_argument(
+    "--resume_from_ckpt",action=argparse.BooleanOptionalAction,default=False,
+    help="if not using wandb and want to resume from a ckpt",
+)
+parser.add_argument(
+    "--wandb_project",type=str,default="stability",
+    help="wandb project name",
+)
+parser.add_argument(
+    "--mixup_pct",type=float,default=.33,
+    help="proportion of way through training when to switch from BiMixCo to SoftCLIP",
+)
+parser.add_argument(
+    "--use_image_aug",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to use image augmentation",
+)
+parser.add_argument(
+    "--num_epochs",type=int,default=100,
+    help="number of epochs of training",
+)
+parser.add_argument(
+    "--lr_scheduler_type",type=str,default='cycle',choices=['cycle','linear'],
+)
+parser.add_argument(
+    "--ckpt_saving",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--ckpt_interval",type=int,default=5,
+    help="save backup ckpt and reconstruct every x epochs",
+)
+parser.add_argument(
+    "--seed",type=int,default=42,
+)
+parser.add_argument(
+    "--max_lr",type=float,default=3e-4,
+)
+parser.add_argument(
+    "--n_samples_save",type=int,default=0,choices=[0,1],
+    help="Number of reconstructions for monitoring progress, 0 will speed up training",
+)
+parser.add_argument(
+    "--clip_mse_ratio",type=float,default=0.7,
+    help="Number of reconstructions for monitoring progress, 0 will speed up training",
+)
+
+if utils.is_interactive():
+    args = parser.parse_args(jupyter_args)
+else:
+    args = parser.parse_args()
+
+# create global variables without the args prefix
+for attribute_name in vars(args).keys():
+    globals()[attribute_name] = getattr(args, attribute_name)
+
+print("global batch_size", batch_size)
+batch_size = int(batch_size / num_devices)
+print("batch_size", batch_size)
+
+
+# In[7]:
+
+
+outdir = os.path.abspath(f'../train_logs/{model_name}')
+if not os.path.exists(outdir):
+    os.makedirs(outdir,exist_ok=True)
+if use_image_aug:
+    import kornia
+    from kornia.augmentation.container import AugmentationSequential
+    img_augment = AugmentationSequential(
+        kornia.augmentation.RandomResizedCrop((224,224), (0.6,1), p=0.3),
+        kornia.augmentation.Resize((224, 224)),
+        kornia.augmentation.RandomHorizontalFlip(p=0.3),
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.3),
+        kornia.augmentation.RandomGrayscale(p=0.3),
+        same_on_batch=False,
+        data_keys=["input"],
+    )
+
+
+# In[8]:
+
+
+wandb_log = True
+
+
+# # Prep data, models, and dataloaders
+
+# ## Dataloader
+
+# In[9]:
+
+
+if subj==1:
+    num_train = 24958
+    num_test = 2770
+test_batch_size = num_test
+
+def my_split_by_node(urls): return urls
+    
+train_url = f"{data_path}/wds/subj0{subj}/train/" + "{0..36}.tar"
+print(train_url)
+
+train_data = wds.WebDataset(train_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+print(test_url)
+
+test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+test_dl = torch.utils.data.DataLoader(test_data, batch_size=test_batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+
+# ### check dataloaders are working
+
+# In[10]:
+
+
+# test_indices = []
+# test_images = []
+# for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+#     test_indices = np.append(test_indices, behav[:,0,5].numpy())
+#     test_images = np.append(test_images, behav[:,0,0].numpy())
+# test_indices = test_indices.astype(np.int16)
+# print(test_i, (test_i+1) * test_batch_size, len(test_indices))
+# print("---\n")
+
+# train_indices = []
+# train_images = []
+# for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+#     train_indices = np.append(train_indices, behav[:,0,5].long().numpy())
+#     train_images = np.append(train_images, behav[:,0,0].numpy())
+# train_indices = train_indices.astype(np.int16)
+# print(train_i, (train_i+1) * batch_size, len(train_indices))
+
+# # train_images = np.hstack((train_images, test_images))
+# # print("WARNING: ADDED TEST IMAGES TO TRAIN IMAGES")
+
+
+# ## Load data and images
+
+# In[11]:
+
+
+# load betas
+f = h5py.File(f'{data_path}/betas_all_subj0{subj}.hdf5', 'r')
+voxels = f['betas'][:]
+print(f"subj0{subj} betas loaded into memory")
+voxels = torch.Tensor(voxels).to("cpu").half()
+if subj==1:
+    voxels = torch.hstack((voxels, torch.zeros((len(voxels), 5))))
+print("voxels", voxels.shape)
+num_voxels = voxels.shape[-1]
+
+# load orig images
+f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+images = f['images'][:]
+images = torch.Tensor(images).to("cpu").half()
+print("images", images.shape)
+
+
+# ## Load models
+
+# ### CLIP image embeddings  model
+
+# In[12]:
+
+
+import transformers
+from transformers import Blip2Processor, Blip2ForConditionalGeneration
+
+from PIL import Image
+
+
+# In[13]:
+
+
+from models import Clipper
+clip_model = Clipper("ViT-L/14", device=torch.device(f"cuda:{local_rank}"), hidden_state=True, norm_embs=True)
+
+
+# In[14]:
+
+
+cache_blip2 = "/fsx/proj-fmri/shared/cache/models--Salesforce--blip2-opt-2.7b/snapshots/6e723d92ee91ebcee4ba74d7017632f11ff4217b"
+
+b2_processor = Blip2Processor.from_pretrained(cache_blip2)
+b2_model = Blip2ForConditionalGeneration.from_pretrained(cache_blip2, torch_dtype=torch.float16, device_map="auto")
+
+#Load in blip2 as well
+"""from lavis.models import load_model_and_preprocess
+from lavis.models import model_zoo
+blip2_model, vis_processors, _ = load_model_and_preprocess(
+            name="blip2_t5", model_type="pretrain_flant5xl_vitL", is_eval=True, device=device)
+
+clip_seq_dim = 257
+clip_emb_dim = 1024
+hidden_dim = 4096"""
+
+
+# In[15]:
+
+
+def embed_images_b2(images):
+    images = (images * 255).type(torch.uint8)
+    with torch.no_grad():
+        inputs_processed = b2_processor(images, return_tensors="pt").to("cuda", torch.float16)
+        enc_imgs = b2_model.vision_model.forward(inputs_processed['pixel_values'])
+        return enc_imgs.last_hidden_state.detach(), inputs_processed
+
+def embeds_to_captions_b2(embeds):
+    with torch.no_grad():
+        input_ids = None #inputs['input_ids']
+        attention_mask = None
+        batch_size = embeds.shape[0]
+        image_embeds = embeds
+        image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device)
+
+        query_tokens = b2_model.query_tokens.expand(image_embeds.shape[0], -1, -1)
+        query_outputs = b2_model.qformer(
+            query_embeds=query_tokens,
+            encoder_hidden_states=image_embeds,
+            encoder_attention_mask=image_attention_mask,
+            return_dict=True,
+        )
+        query_output = query_outputs.last_hidden_state
+
+        language_model_inputs = b2_model.language_projection(query_output)
+        language_attention_mask = torch.ones(
+            language_model_inputs.size()[:-1], dtype=torch.long, device=language_model_inputs.device
+        )
+        if input_ids is None:
+            input_ids = (
+                torch.LongTensor([[b2_model.config.text_config.bos_token_id]])
+                .repeat(batch_size, 1)
+                .to(image_embeds.device)
+            )
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+        attention_mask = torch.cat([language_attention_mask, attention_mask.to(language_attention_mask.device)], dim=1)
+
+        # concatenate query embeddings with prompt embeddings
+        inputs_embeds = b2_model.get_input_embeddings()(input_ids)
+        inputs_embeds = torch.cat([language_model_inputs, inputs_embeds.to(language_model_inputs.device)], dim=1)
+
+        outputs = b2_model.language_model.generate(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+        )
+        text = b2_processor.batch_decode(outputs, skip_special_tokens=True)
+        
+        return outputs, text
+
+
+# In[16]:
+
+
+image_test = images[1:20].permute(0,2,3,1)
+#raw_image = Image.open('/fsx/proj-fmri/shared/controlNetData/target/img_t1.jpg').convert('RGB')
+# Convert the image to a NumPy array
+#image_test = np.array(raw_image)
+
+
+# In[17]:
+
+
+"""import matplotlib.pyplot as plt
+# Plotting one of the images (taking the first image as an example)
+img_to_plot = inputs_rec['pixel_values'][-1]
+
+# Transpose the image for correct display (PyTorch: [C, H, W], Matplotlib: [H, W, C])
+img_to_plot = img_to_plot.permute(1, 2, 0).to(torch.float32).to('cpu')
+print(img_to_plot.shape)
+
+plt.imshow(img_to_plot)
+plt.show()"""
+
+
+# In[18]:
+
+
+embeds_test, inputs_rec = embed_images_b2(image_test)
+
+
+# In[19]:
+
+
+#inputs_rec['pixel_values'].shape
+
+
+# In[20]:
+
+
+#out = b2_model.generate(**inputs_rec)
+#print(b2_processor.decode(out[0], skip_special_tokens=True).strip())
+
+
+# In[21]:
+
+
+outputs_test, text_test = embeds_to_captions_b2(embeds_test)
+
+
+# In[22]:
+
+
+text_test
+
+
+# In[23]:
+
+
+#inputss['pixel_values'].shape
+
+
+# In[24]:
+
+
+#image_test.shape
+
+
+# In[25]:
+
+
+
+
+# In[26]:
+
+
+clip_seq_dim = 257 #blip2 image encoder shapes
+clip_emb_dim = 1408 #blip2 image encoder shapes
+hidden_dim = 2048
+
+
+# ### SD VAE (blurry images)
+
+# In[27]:
+
+
+from diffusers import AutoencoderKL
+autoenc = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16, cache_dir="/fsx/proj-fmri/shared/cache")
+# autoenc.load_state_dict(torch.load('../train_logs/sdxl_vae_normed/best.pth')["model_state_dict"])
+autoenc.eval()
+autoenc.requires_grad_(False)
+autoenc.to(device)
+utils.count_params(autoenc)
+
+
+# ### MindEye modules
+
+# In[28]:
+
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+model
+
+
+# In[29]:
+
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer
+    def __init__(self, input_size, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linear = torch.nn.Linear(input_size, out_features)
+    def forward(self, x):
+        return self.linear(x)
+        
+model.ridge = RidgeRegression(voxels.shape[1], out_features=hidden_dim)
+utils.count_params(model.ridge)
+utils.count_params(model)
+
+b = torch.randn((2,1,voxels.shape[1]))
+print(b.shape, model.ridge(b).shape)
+
+
+# In[30]:
+
+
+from functools import partial
+from diffusers.models.vae import Decoder
+class BrainNetwork(nn.Module):
+    def __init__(self, out_dim=768, in_dim=15724, clip_size=768, h=4096, n_blocks=4, norm_type='ln', act_first=False, drop=.35, blurry_dim=16):
+        super().__init__()
+        self.blurry_dim = blurry_dim
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        self.lin0 = nn.Linear(in_dim, h)
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, out_dim, bias=True)
+        # self.blin1 = nn.Linear(out_dim, blurry_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.clip_size = clip_size
+        self.clip_proj = nn.Sequential(
+            nn.LayerNorm(clip_size),
+            nn.GELU(),
+            nn.Linear(clip_size, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, 2048),
+            nn.LayerNorm(2048),
+            nn.GELU(),
+            nn.Linear(2048, clip_size)
+        )
+        # self.upsampler = Decoder(
+        #         in_channels=64,
+        #         out_channels=4,
+        #         up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D"],
+        #         block_out_channels=[64, 128, 256],
+        #         layers_per_block=1,
+        #     )
+        
+    def forward(self, x):
+        x = self.lin0(x)
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        # b = self.blin1(x)
+        # b = self.upsampler(b.reshape(len(b), -1, 7, 7))
+        c = self.clip_proj(x.reshape(len(x), -1, self.clip_size))
+        # return c, b
+        return c
+
+model.backbone = BrainNetwork(h=2048, in_dim=hidden_dim, clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, blurry_dim=64*7*7) 
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+b = torch.randn((4,hidden_dim))
+print(b.shape)
+clip_ = model.backbone(b)
+print(clip_.shape)
+
+
+# In[31]:
+
+
+no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+opt_grouped_parameters = [
+    {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
+]
+
+optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr, betas=(0.9, 0.95))
+
+if lr_scheduler_type == 'linear':
+    lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+        optimizer,
+        total_iters=int(num_epochs*(num_train*num_devices//batch_size)),
+        last_epoch=-1
+    )
+elif lr_scheduler_type == 'cycle':
+    total_steps=int(num_epochs*(num_train*num_devices//batch_size))
+    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer, 
+        max_lr=max_lr,
+        total_steps=total_steps,
+        final_div_factor=1000,
+        last_epoch=-1, pct_start=2/num_epochs
+    )
+    
+def save_ckpt(tag):    
+    ckpt_path = outdir+f'/{tag}.pth'
+    print(f'saving {ckpt_path}',flush=True)
+    unwrapped_model = accelerator.unwrap_model(model)
+    try:
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': unwrapped_model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'lr_scheduler': lr_scheduler.state_dict(),
+            'train_losses': losses,
+            'test_losses': test_losses,
+            'lrs': lrs,
+            }, ckpt_path)
+    except:
+        print("Couldn't save... moving on to prevent crashing.")
+    del unwrapped_model
+        
+print("\nDone with model preparations!")
+utils.count_params(model)
+
+
+# # Weights and Biases
+
+# In[32]:
+
+
+# params for wandb
+if local_rank==0 and True: # only use main process for wandb logging
+    import wandb
+    
+    wandb_project = 'mindeyev2'
+    wandb_run = model_name
+    wandb_notes = ''
+    
+    print(f"wandb {wandb_project} run {wandb_run}")
+    wandb.login(host='https://stability.wandb.io')#, relogin=True)
+    wandb_config = {
+      "model_name": model_name,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "use_image_aug": use_image_aug,
+      "max_lr": max_lr,
+      "lr_scheduler_type": lr_scheduler_type,
+      "mixup_pct": mixup_pct,
+      "num_train": num_train,
+      "num_test": num_test,
+      "seed": seed,
+      "distributed": distributed,
+      "num_devices": num_devices,
+      "world_size": world_size,
+    }
+    print("wandb_config:\n",wandb_config)
+    if False: # wandb_auto_resume
+        print("wandb_id:",model_name)
+        wandb.init(
+            id = model_name,
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+            resume="allow",
+        )
+    else:
+        wandb.init(
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+        )
+else:
+    wandb_log = False
+
+
+# # More custom functions
+
+# In[33]:
+
+
+# using the same preprocessing as was used in MindEye + BrainDiffuser
+pixcorr_preprocess = transforms.Compose([
+    transforms.Resize(425, interpolation=transforms.InterpolationMode.BILINEAR),
+])
+def pixcorr(images,brains):
+    # Flatten images while keeping the batch dimension
+    all_images_flattened = pixcorr_preprocess(images).reshape(len(images), -1)
+    all_brain_recons_flattened = pixcorr_preprocess(brains).view(len(brains), -1)
+    corrmean = torch.diag(utils.batchwise_pearson_correlation(all_images_flattened, all_brain_recons_flattened)).mean()
+    return corrmean
+
+
+# # Main
+
+# In[34]:
+
+
+epoch = 0
+losses, test_losses, lrs = [], [], []
+best_test_loss = 1e9
+soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))
+
+# Optionally resume from checkpoint #
+if resume_from_ckpt:
+    print("\n---resuming from last.pth ckpt---\n")
+    try:
+        checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+    except:
+        print('last.pth failed... trying last_backup.pth')
+        checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+    epoch = checkpoint['epoch']
+    print("Epoch",epoch)
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])
+    del checkpoint
+elif wandb_log:
+    if wandb.run.resumed:
+        print("\n---resuming from last.pth ckpt---\n")
+        try:
+            checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+        except:
+            print('last.pth failed... trying last_backup.pth')
+            checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+        epoch = checkpoint['epoch']
+        print("Epoch",epoch)
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+        diffusion_diffuser.load_state_dict(checkpoint['model_state_dict'])
+        del checkpoint
+torch.cuda.empty_cache()
+
+
+# In[35]:
+
+
+model, optimizer, train_dl, test_dl, lr_scheduler = accelerator.prepare(
+model, optimizer, train_dl, test_dl, lr_scheduler
+)
+
+
+# In[36]:
+
+
+"""transform = transforms.Compose(
+            [
+                transforms.Resize(
+                    (224, 224),
+                ),
+                transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+            ]
+        )
+
+def tensor_2_embed(image):    
+    image_for_blip2 = transform(image)
+    
+    #Generate embeddings
+    with blip2_model.maybe_autocast():
+        blip2_target = blip2_model.ln_vision(blip2_model.visual_encoder(image_for_blip2))
+                
+    return blip2_target
+
+def embed_2_caption(image_embeds, model):
+    image_embeds = image_embeds.float()
+    image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(
+                image.device)
+
+    query_tokens = model.query_tokens.expand(image_embeds.shape[0], -1, -1)
+    query_output = model.Qformer.bert(
+                query_embeds=query_tokens,
+                encoder_hidden_states=image_embeds,
+                encoder_attention_mask=image_atts,
+                return_dict=True)
+
+    inputs_t5 = model.t5_proj(query_output.last_hidden_state)
+    atts_t5 = torch.ones(inputs_t5.size()[:-1], dtype=torch.long).to(image.device)
+    prompt = model.prompt
+    input_tokens = model.t5_tokenizer(
+                prompt, padding="longest", return_tensors="pt"
+            ).to(image.device)
+    encoder_atts = torch.cat([atts_t5, input_tokens.attention_mask], dim=1)
+    
+    with model.maybe_autocast(dtype=torch.bfloat16):
+            inputs_embeds = model.t5_model.encoder.embed_tokens(input_tokens.input_ids)
+            inputs_embeds = torch.cat([inputs_t5, inputs_embeds], dim=1)
+
+            outputs = model.t5_model.generate(
+                inputs_embeds=inputs_embeds,
+                attention_mask=encoder_atts)
+            output_text = model.t5_tokenizer.batch_decode(
+                outputs, skip_special_tokens=True)
+    
+    return output_text"""
+
+
+# In[37]:
+
+
+wandb_log = True
+
+
+# In[ ]:
+
+
+print(f"{model_name} starting with epoch {epoch} / {num_epochs}")
+progress_bar = tqdm(range(epoch,num_epochs), ncols=1200, disable=(local_rank!=0))
+test_image, test_voxel = None, None
+mse = nn.MSELoss()
+for epoch in progress_bar:
+    model.train()
+    
+    fwd_percent_correct = 0.
+    bwd_percent_correct = 0.
+    test_fwd_percent_correct = 0.
+    test_bwd_percent_correct = 0.
+
+    loss_clip_total = 0.
+    loss_blurry_total = 0.
+    test_loss_clip_total = 0.
+    test_loss_blurry_total = 0.
+
+    blurry_pixcorr = 0.
+    test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1
+    
+    for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+        if epoch == 0:
+            lrs.append(0)
+            break
+        with torch.cuda.amp.autocast():
+            optimizer.zero_grad()
+
+            voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+            
+            image = images[behav[:,0,0].cpu().long()].to(device).float()
+
+            # blurry_image_enc = autoenc.encode(image).latent_dist.mode()
+            
+            if use_image_aug: image = img_augment(image)
+            # clip_target = clip_model.embed_image(image)
+            clip_target = embed_images_b2(image)[0].to(device) #####CHANGED
+            assert not torch.any(torch.isnan(clip_target))
+  
+            if epoch < int(mixup_pct * num_epochs):
+                voxel, perm, betas, select = utils.mixco(voxel)
+
+            voxel_ridge = model.ridge(voxel)
+            
+            # clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+            clip_voxels = model.backbone(voxel_ridge)
+            
+            clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+            clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+            if epoch < int(mixup_pct * num_epochs):                
+                loss_clip = utils.mixco_nce(
+                     clip_voxels_norm,
+                     clip_target_norm,
+                     temp=.006, 
+                     perm=perm, betas=betas, select=select)
+            else:
+                epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]
+                loss_clip = utils.soft_clip_loss(
+                     clip_voxels_norm,
+                     clip_target_norm,
+                     temp=epoch_temp)
+            
+            loss_mse= mse(clip_voxels, clip_target)
+
+            # loss_blurry = mse(blurry_image_enc_, blurry_image_enc) 
+
+            loss_clip_total += loss_clip.item()
+            # loss_blurry_total += loss_blurry.item()
+
+            # loss = loss_blurry + loss_clip
+            loss = (clip_mse_ratio * loss_clip) + ((1 - clip_mse_ratio) * loss_mse)
+            if (train_i % 10 == 0):
+                print(train_i, loss)
+            # print(batch_size)
+            utils.check_loss(loss)
+            accelerator.backward(loss)
+            optimizer.step()
+    
+            losses.append(loss.item())
+            lrs.append(optimizer.param_groups[0]['lr'])
+    
+            # forward and backward top 1 accuracy        
+            labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+            fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)
+            bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)
+
+            # with torch.no_grad():
+            #     # only doing pixcorr eval on a subset (8) of the samples per batch because its costly & slow to compute autoenc.decode()
+            #     random_samps = np.random.choice(np.arange(len(voxel)), size=8, replace=False)
+            #     blurry_recon_images = autoenc.decode(blurry_image_enc_[random_samps]).sample.clamp(0,1)
+            #     blurry_pixcorr += pixcorr(image[random_samps], blurry_recon_images)
+
+            if lr_scheduler_type is not None:
+                lr_scheduler.step()
+    
+    model.eval()
+    for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+        with torch.cuda.amp.autocast():
+            with torch.no_grad():   
+                # all test samples should be loaded per batch such that test_i should never exceed 0
+                if len(behav) != num_test: print("!",len(behav),num_test)
+                
+                ## Average same-image repeats ##
+                if test_image is None:
+                    voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+                    
+                    image = behav[:,0,0].cpu().long()
+                    
+                    unique_image, sort_indices = torch.unique(image, return_inverse=True)
+                    for im in unique_image:
+                        locs = torch.where(im == image)[0]
+                        if test_image is None:
+                            test_image = images[im][None]
+                            test_voxel = torch.mean(voxel[locs],axis=0)[None]
+                        else:
+                            test_image = torch.vstack((test_image, images[im][None]))
+                            test_voxel = torch.vstack((test_voxel, torch.mean(voxel[locs],axis=0)[None]))
+    
+                # sample of batch_size
+                random_indices = torch.arange(len(test_voxel))[:batch_size] #torch.randperm(len(test_voxel))[:300]
+                voxel = test_voxel[random_indices].to(device)
+                image = test_image[random_indices].to(device)
+                assert len(image) == batch_size
+    
+                # blurry_image_enc = autoenc.encode(image).latent_dist.mode()
+        
+                # clip_target = clip_model.embed_image(image.float())
+                clip_target = embed_images_b2(image)[0].to(device) #####CHANGED
+    
+                voxel_ridge = model.ridge(voxel)
+                
+                # clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+                clip_voxels = model.backbone(voxel_ridge)
+                
+                clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+                clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+        
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006)
+                
+                loss_mse = mse(clip_voxels, clip_target)
+
+                # loss_blurry = mse(blurry_image_enc_, blurry_image_enc)
+                
+                # loss = loss_blurry + loss_clip
+                loss = (clip_mse_ratio * loss_clip) + ((1 - clip_mse_ratio) * loss_mse)
+                
+                utils.check_loss(loss)
+        
+                test_losses.append(loss.item())
+        
+                # forward and backward top 1 accuracy        
+                labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+                test_fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)
+                test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)
+
+                # # halving the batch size because the decoder is computationally heavy
+                # blurry_recon_images = autoenc.decode(blurry_image_enc_[:len(voxel)//2]).sample.clamp(0,1)
+                # blurry_recon_images = torch.vstack((blurry_recon_images, autoenc.decode(blurry_image_enc_[len(voxel)//2:]).sample.clamp(0,1)))
+                # test_blurry_pixcorr += pixcorr(image, blurry_recon_images)
+
+                #Find captions and print next to images
+                #caption1 = embed_2_caption(clip_voxels[[0]], blip2_model)
+                #caption2 = embed_2_caption(clip_voxels[[1]], blip2_model)
+
+                #true_embed1 = tensor_2_embed(image[[0]])
+                #true_embed2 = tensor_2_embed(image[[1]])
+
+                # print(clip_voxels[[0]].shape)
+                # print(true_embed1.shape)
+                
+                #true_caption1 = embed_2_caption(true_embed1, blip2_model)
+                #true_caption2 = embed_2_caption(true_embed2, blip2_model)
+                
+                # transform blurry recon latents to images and plot it
+                #fig, axes = plt.subplots(2, 2, figsize=(8, 4))
+                #axes[0,0].imshow(utils.torch_to_Image(image[[0]]))
+                #axes[0,1].imshow(utils.torch_to_Image(image[[1]]))
+                #axes[0,0].axis('off'); axes[0,1].axis('off'); axes[1,0].axis('off'); axes[1,1].axis('off')
+                #axes[0,0].set_title(caption1)
+                #axes[0,1].set_title(caption2)
+                #axes[1,0].set_title(true_caption1)
+                #axes[1,1].set_title(true_caption2)
+
+                #plt.show()
+                
+                # # transform blurry recon latents to images and plot it
+                # fig, axes = plt.subplots(1, 4, figsize=(8, 4))
+                # axes[0].imshow(utils.torch_to_Image(image[[0]]))
+                # axes[1].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[0]]).sample.clamp(0,1)))
+                # axes[2].imshow(utils.torch_to_Image(image[[1]]))
+                # axes[3].imshow(utils.torch_to_Image(autoenc.decode(blurry_image_enc_[[1]]).sample.clamp(0,1)))
+                # axes[0].axis('off'); axes[1].axis('off'); axes[2].axis('off'); axes[3].axis('off')
+                # axes[0].set_title(caption1)
+                # axes[3].set_title(caption2)
+                # plt.show()
+                
+
+    if local_rank==0:      
+        # if utils.is_interactive(): clear_output(wait=True)
+        assert (test_i+1) == 1
+        logs = {"train/loss": np.mean(losses[-(train_i+1):]),
+            "test/loss": np.mean(test_losses[-(test_i+1):]),
+            "train/lr": lrs[-1],
+            "train/num_steps": len(losses),
+            "test/num_steps": len(test_losses),
+            "train/fwd_pct_correct": fwd_percent_correct / (train_i + 1),
+            "train/bwd_pct_correct": bwd_percent_correct / (train_i + 1),
+            "test/test_fwd_pct_correct": test_fwd_percent_correct / (test_i + 1),
+            "test/test_bwd_pct_correct": test_bwd_percent_correct / (test_i + 1),
+            "train/loss_clip_total": loss_clip_total / (train_i + 1),
+            "train/loss_blurry_total": loss_blurry_total / (train_i + 1),
+            "test/loss_clip_total": test_loss_clip_total / (test_i + 1),
+            "test/loss_blurry_total": test_loss_blurry_total / (test_i + 1),
+            "train/blurry_pixcorr": blurry_pixcorr / (train_i + 1),
+            "test/blurry_pixcorr": test_blurry_pixcorr / (test_i + 1),
+            }
+        progress_bar.set_postfix(**logs)
+                          
+        fig, axes = plt.subplots(1, 8, figsize=(10, 4))
+        jj=-1
+        for j in [0,1,2,3,4,5,6,7]:
+            jj+=1
+            axes[jj].imshow(utils.torch_to_Image(image[j]))
+            axes[jj].axis('off')
+
+        if wandb_log:
+            generated_captions = embeds_to_captions_b2(clip_voxels[0:8])
+            print(generated_captions[1])
+            logs[f"test/recons"] = wandb.Image(fig, caption=f"epoch{epoch:03d}" + "\n".join(generated_captions[1]))
+            plt.close()
+        # Save model checkpoint and reconstruct
+        if epoch % ckpt_interval == 0:
+            if not utils.is_interactive():
+                save_ckpt(f'last')
+                
+        if wandb_log: wandb.log(logs)
+
+    # wait for other GPUs to catch up if needed
+    accelerator.wait_for_everyone()
+    torch.cuda.empty_cache()
+    gc.collect()
+
+print("\n===Finished!===\n")
+if ckpt_saving:
+    save_ckpt(f'last')
+if not utils.is_interactive():
+    sys.exit(0)
+
+
+# In[ ]:
+
+
+plt.plot(losses)
+plt.show()
+plt.plot(test_losses)
+plt.show()
+
+
+# In[ ]:
+
+
+print('test')
+
+
+# In[ ]:
+
+
+def tensor_2_embed_old(tensor):
+    embed_array = torch.zeros((tensor.shape[0],257, 1024)) 
+    to_pil = ToPILImage()
+    for sample in range(tensor.shape[0]):
+        PIL_image = to_pil(tensor[sample])
+        image_for_blip2 = vis_processors["eval"](PIL_image).unsqueeze(0).to(device)
+        #Generate embeddings
+        with blip2_model.maybe_autocast():
+                blip2_target = blip2_model.ln_vision(blip2_model.visual_encoder(image_for_blip2))
+                embed_array[sample] = blip2_target
+                
+    return embed_array
+
+
+# In[ ]:
+
+
+
+
+
+# In[ ]:
+
+
+
+
+
+# In[ ]:
+
+
+
+

src/utils.py

@@ -0,0 +1,368 @@
+import numpy as np
+from torchvision import transforms
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import PIL
+import random
+import os
+import matplotlib.pyplot as plt
+import pandas as pd
+import math
+import webdataset as wds
+import tempfile
+from torchvision.utils import make_grid
+# from diffusers.utils import randn_tensor
+
+import json
+from torchmetrics.image.fid import FrechetInceptionDistance
+from PIL import Image
+import requests
+import io
+import time 
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+def is_interactive():
+    import __main__ as main
+    return not hasattr(main, '__file__')
+
+def seed_everything(seed=0, cudnn_deterministic=True):
+    random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    if cudnn_deterministic:
+        torch.backends.cudnn.deterministic = True
+    else:
+        ## needs to be False to use conv3D
+        print('Note: not using cudnn.deterministic')
+
+def np_to_Image(x):
+    if x.ndim==4:
+        x=x[0]
+    return PIL.Image.fromarray((x.transpose(1, 2, 0)*127.5+128).clip(0,255).astype('uint8'))
+
+def torch_to_Image(x):
+    if x.ndim==4:
+        x=x[0]
+    return transforms.ToPILImage()(x)
+
+def Image_to_torch(x):
+    try:
+        x = (transforms.ToTensor()(x)[:3].unsqueeze(0)-.5)/.5
+    except:
+        x = (transforms.ToTensor()(x[0])[:3].unsqueeze(0)-.5)/.5
+    return x
+
+def torch_to_matplotlib(x,device=device):
+    if torch.mean(x)>10:
+        x = (x.permute(0, 2, 3, 1)).clamp(0, 255).to(torch.uint8)
+    else:
+        x = (x.permute(0, 2, 3, 1) * 255).clamp(0, 255).to(torch.uint8)
+    if device=='cpu':
+        return x[0]
+    else:
+        return x.cpu().numpy()[0]
+
+def pairwise_cosine_similarity(A, B, dim=1, eps=1e-8):
+    #https://stackoverflow.com/questions/67199317/pytorch-cosine-similarity-nxn-elements
+    numerator = A @ B.T
+    A_l2 = torch.mul(A, A).sum(axis=dim)
+    B_l2 = torch.mul(B, B).sum(axis=dim)
+    denominator = torch.max(torch.sqrt(torch.outer(A_l2, B_l2)), torch.tensor(eps))
+    return torch.div(numerator, denominator)
+
+def batchwise_pearson_correlation(Z, B):
+    # Calculate means
+    Z_mean = torch.mean(Z, dim=1, keepdim=True)
+    B_mean = torch.mean(B, dim=1, keepdim=True)
+
+    # Subtract means
+    Z_centered = Z - Z_mean
+    B_centered = B - B_mean
+
+    # Calculate Pearson correlation coefficient
+    numerator = Z_centered @ B_centered.T
+    Z_centered_norm = torch.linalg.norm(Z_centered, dim=1, keepdim=True)
+    B_centered_norm = torch.linalg.norm(B_centered, dim=1, keepdim=True)
+    denominator = Z_centered_norm @ B_centered_norm.T
+
+    pearson_correlation = (numerator / denominator)
+    return pearson_correlation
+
+def batchwise_cosine_similarity(Z,B):
+    # https://www.h4pz.co/blog/2021/4/2/batch-cosine-similarity-in-pytorch-or-numpy-jax-cupy-etc
+    B = B.T
+    Z_norm = torch.linalg.norm(Z, dim=1, keepdim=True)  # Size (n, 1).
+    B_norm = torch.linalg.norm(B, dim=0, keepdim=True)  # Size (1, b).
+    cosine_similarity = ((Z @ B) / (Z_norm @ B_norm)).T
+    return cosine_similarity
+
+def topk(similarities,labels,k=5):
+    if k > similarities.shape[0]:
+        k = similarities.shape[0]
+    topsum=0
+    for i in range(k):
+        topsum += torch.sum(torch.argsort(similarities,axis=1)[:,-(i+1)] == labels)/len(labels)
+    return topsum
+
+def get_non_diagonals(a):
+    a = torch.triu(a,diagonal=1)+torch.tril(a,diagonal=-1)
+    # make diagonals -1
+    a=a.fill_diagonal_(-1)
+    return a
+
+def gather_features(image_features, voxel_features, accelerator):  
+    all_image_features = accelerator.gather(image_features.contiguous())
+    if voxel_features is not None:
+        all_voxel_features = accelerator.gather(voxel_features.contiguous())
+        return all_image_features, all_voxel_features
+    return all_image_features
+
+def soft_clip_loss(preds, targs, temp=0.125): #, distributed=False, accelerator=None):
+    # if not distributed:
+    clip_clip = (targs @ targs.T)/temp
+    brain_clip = (preds @ targs.T)/temp
+    # else:
+    #     all_targs = gather_features(targs, None, accelerator)
+    #     clip_clip = (targs @ all_targs.T)/temp
+    #     brain_clip = (preds @ all_targs.T)/temp
+    
+    loss1 = -(brain_clip.log_softmax(-1) * clip_clip.softmax(-1)).sum(-1).mean()
+    loss2 = -(brain_clip.T.log_softmax(-1) * clip_clip.softmax(-1)).sum(-1).mean()
+    
+    loss = (loss1 + loss2)/2
+    return loss
+
+def soft_siglip_loss(preds, targs, temp, bias):
+    temp = torch.exp(temp)
+    
+    logits = (preds @ targs.T) * temp + bias
+    # diagonals (aka paired samples) should be >0 and off-diagonals <0
+    labels = (targs @ targs.T) - 1 + (torch.eye(len(targs)).to(targs.dtype).to(targs.device))
+
+    loss1 = -torch.sum(nn.functional.logsigmoid(logits * labels[:len(preds)])) / len(preds)
+    loss2 = -torch.sum(nn.functional.logsigmoid(logits.T * labels[:,:len(preds)])) / len(preds)
+    loss = (loss1 + loss2)/2
+    return loss
+
+def mixco_hard_siglip_loss(preds, targs, temp, bias, perm, betas):
+    temp = torch.exp(temp)
+    
+    probs = torch.diag(betas)
+    probs[torch.arange(preds.shape[0]).to(preds.device), perm] = 1 - betas
+
+    logits = (preds @ targs.T) * temp + bias
+    labels = probs * 2 - 1
+    #labels = torch.eye(len(targs)).to(targs.dtype).to(targs.device) * 2 - 1
+    
+    loss1 = -torch.sum(nn.functional.logsigmoid(logits * labels)) / len(preds)
+    loss2 = -torch.sum(nn.functional.logsigmoid(logits.T * labels)) / len(preds)
+    loss = (loss1 + loss2)/2
+    return loss
+
+def mixco(voxels, beta=0.15, s_thresh=0.5, perm=None, betas=None, select=None):
+    if perm is None:
+        perm = torch.randperm(voxels.shape[0])
+    voxels_shuffle = voxels[perm].to(voxels.device,dtype=voxels.dtype)
+    if betas is None:
+        betas = torch.distributions.Beta(beta, beta).sample([voxels.shape[0]]).to(voxels.device,dtype=voxels.dtype)
+    if select is None:
+        select = (torch.rand(voxels.shape[0]) <= s_thresh).to(voxels.device)
+    betas_shape = [-1] + [1]*(len(voxels.shape)-1)
+    voxels[select] = voxels[select] * betas[select].reshape(*betas_shape) + \
+        voxels_shuffle[select] * (1 - betas[select]).reshape(*betas_shape)
+    betas[~select] = 1
+    return voxels, perm, betas, select
+
+def mixco_clip_target(clip_target, perm, select, betas):
+    clip_target_shuffle = clip_target[perm]
+    clip_target[select] = clip_target[select] * betas[select].reshape(-1, 1) + \
+        clip_target_shuffle[select] * (1 - betas[select]).reshape(-1, 1)
+    return clip_target
+
+def mixco_nce(preds, targs, temp=0.1, perm=None, betas=None, select=None, distributed=False, 
+              accelerator=None, local_rank=None, bidirectional=True):
+    brain_clip = (preds @ targs.T)/temp
+    
+    if perm is not None and betas is not None and select is not None:
+        probs = torch.diag(betas)
+        probs[torch.arange(preds.shape[0]).to(preds.device), perm] = 1 - betas
+
+        loss = -(brain_clip.log_softmax(-1) * probs).sum(-1).mean()
+        if bidirectional:
+            loss2 = -(brain_clip.T.log_softmax(-1) * probs.T).sum(-1).mean()
+            loss = (loss + loss2)/2
+        return loss
+    else:
+        loss =  F.cross_entropy(brain_clip, torch.arange(brain_clip.shape[0]).to(brain_clip.device))
+        if bidirectional:
+            loss2 = F.cross_entropy(brain_clip.T, torch.arange(brain_clip.shape[0]).to(brain_clip.device))
+            loss = (loss + loss2)/2
+        return loss
+    
+def count_params(model):
+    total = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print('param counts:\n{:,} total\n{:,} trainable'.format(total, trainable))
+
+def image_grid(imgs, rows, cols):
+    w, h = imgs[0].size
+    grid = PIL.Image.new('RGB', size=(cols*w, rows*h))
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i%cols*w, i//cols*h))
+    return grid
+    
+def check_loss(loss):
+    if loss.isnan().any():
+        raise ValueError('NaN loss')
+
+def cosine_anneal(start, end, steps):
+    return end + (start - end)/2 * (1 + torch.cos(torch.pi*torch.arange(steps)/(steps-1)))
+
+def resize(img, img_size=128):
+    if img.ndim == 3: img = img[None]
+    return nn.functional.interpolate(img, size=(img_size, img_size), mode='nearest')
+
+def patchify(img, patch_size=16):
+    B, C, H, W = img.size()
+    patches = img.unfold(2, patch_size, patch_size).unfold(3, patch_size, patch_size)
+    patches = patches.contiguous().view(B, C, -1, patch_size, patch_size)
+    return patches.permute(0, 2, 1, 3, 4)
+
+def unpatchify(patches):
+    B, N, C, H, W = patches.shape  # B=Batch size, N=Number of patches, C=Channels, H=Height, W=Width
+    patches = patches.view(B, int(N**0.5), int(N**0.5), C, H, W)
+    patches = patches.permute(0, 3, 1, 4, 2, 5).contiguous()
+    return patches.view(B, C, H*int(N**0.5), W*int(N**0.5))
+
+import braceexpand
+def get_dataloaders(
+    batch_size,
+    image_var='images',
+    num_devices=None,
+    num_workers=None,
+    train_url=None,
+    val_url=None,
+    meta_url=None,
+    num_train=None,
+    num_val=None,
+    cache_dir="/scratch/tmp/wds-cache",
+    seed=0,
+    voxels_key="nsdgeneral.npy",
+    val_batch_size=None,
+    to_tuple=["voxels", "images", "trial"],
+    local_rank=0,
+    world_size=1,
+):
+    print("Getting dataloaders...")
+    assert image_var == 'images'
+    
+    def my_split_by_node(urls):
+        return urls
+    
+    train_url = list(braceexpand.braceexpand(train_url))
+    val_url = list(braceexpand.braceexpand(val_url))
+
+    if num_devices is None:
+        num_devices = torch.cuda.device_count()
+    
+    if num_workers is None:
+        num_workers = num_devices
+    
+    if num_train is None:
+        metadata = json.load(open(meta_url))
+        num_train = metadata['totals']['train']
+    if num_val is None:
+        metadata = json.load(open(meta_url))
+        num_val = metadata['totals']['val']
+
+    if val_batch_size is None:
+        val_batch_size = batch_size
+        
+    global_batch_size = batch_size * num_devices
+    num_batches = math.floor(num_train / global_batch_size)
+    num_worker_batches = math.floor(num_batches / num_workers)
+    if num_worker_batches == 0: num_worker_batches = 1
+    
+    print("\nnum_train",num_train)
+    print("global_batch_size",global_batch_size)
+    print("batch_size",batch_size)
+    print("num_workers",num_workers)
+    print("num_batches",num_batches)
+    print("num_worker_batches", num_worker_batches)
+    
+    # train_url = train_url[local_rank:world_size]
+    train_data = wds.WebDataset(train_url, resampled=False, cache_dir=cache_dir, nodesplitter=my_split_by_node)\
+        .shuffle(500, initial=500, rng=random.Random(42))\
+        .decode("torch")\
+        .rename(images="jpg;png", voxels=voxels_key, trial="trial.npy", coco="coco73k.npy", reps="num_uniques.npy")\
+        .to_tuple(*to_tuple)#\
+        # .batched(batch_size, partial=True)#\
+        # .with_epoch(num_worker_batches)
+    
+    # BATCH SIZE SHOULD BE NONE!!! FOR TRAIN AND VAL | resampled=True for train | .batched(val_batch_size, partial=False)
+    train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, num_workers=1, shuffle=False)
+
+    # Validation 
+    print("val_batch_size",val_batch_size)
+    val_data = wds.WebDataset(val_url, resampled=False, cache_dir=cache_dir, nodesplitter=my_split_by_node)\
+        .shuffle(500, initial=500, rng=random.Random(42))\
+        .decode("torch")\
+        .rename(images="jpg;png", voxels=voxels_key, trial="trial.npy", coco="coco73k.npy", reps="num_uniques.npy")\
+        .to_tuple(*to_tuple)#\
+        # .batched(val_batch_size, partial=True)
+    val_dl = torch.utils.data.DataLoader(val_data, batch_size=val_batch_size, num_workers=1, shuffle=False, drop_last=True)
+
+    return train_dl, val_dl, num_train, num_val
+
+pixcorr_preprocess = transforms.Compose([
+    transforms.Resize(425, interpolation=transforms.InterpolationMode.BILINEAR),
+])
+def pixcorr(images,brains):
+    all_images_flattened = pixcorr_preprocess(images).reshape(len(images), -1)
+    all_brain_recons_flattened = pixcorr_preprocess(brains).view(len(brains), -1)
+    corrmean = torch.diag(batchwise_pearson_correlation(all_images_flattened, all_brain_recons_flattened)).mean()
+    return corrmean
+    
+pixcorr_origsize_nanmean_preprocess = transforms.Compose([
+    transforms.Resize(128, interpolation=transforms.InterpolationMode.BILINEAR),
+])
+def pixcorr_origsize_nanmean(images,brains):
+    all_images_flattened = pixcorr_origsize_nanmean_preprocess(images).reshape(len(images), -1)
+    all_brain_recons_flattened = brains.view(len(brains), -1) # assuming it's already 128 size
+    corrmean = torch.nanmean(torch.diag(batchwise_pearson_correlation(all_images_flattened, all_brain_recons_flattened)))
+    return corrmean
+
+def select_annotations(annots, random=False):
+    """
+    There are 5 annotations per image. Select one of them for each image.
+    """
+    for i, b in enumerate(annots):
+        t = ''
+        if random:
+            # select random non-empty annotation
+            while t == '':
+                rand = torch.randint(5, (1,1))[0][0]
+                t = b[rand]
+        else:
+            # select first non-empty annotation
+            for j in range(5):
+                if b[j] != '':
+                    t = b[j]
+                    break
+        if i == 0:
+            txt = np.array(t)
+        else:
+            txt = np.vstack((txt, t))
+    txt = txt.flatten()
+    return txt
+
+def add_saturation(image, alpha=2):
+    gray_image = 0.2989 * image[:, 0, :, :] + 0.5870 * image[:, 1, :, :] + 0.1140 * image[:, 2, :, :]
+    gray_image = gray_image.unsqueeze(1).expand_as(image)
+    saturated_image = alpha * image + (1 - alpha) * gray_image
+    return torch.clamp(saturated_image, 0, 1)

train_mem_logs/error.pth

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

train_mem_logs/error_tensors.pth

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

train_mem_logs/test/last.pth

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

train_mem_logs/test_mem/last.pth

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