Add files using upload-large-folder tool

MindEyeV2/src/TrainB5k.py

@@ -0,0 +1,1438 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# # Import packages & functions
+
+# In[1]:
+
+
+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 matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from accelerate import Accelerator
+
+# SDXL unCLIP requires code from https://github.com/Stability-AI/generative-models/tree/main
+sys.path.append('generative_models/')
+import sgm
+from generative_models.sgm.modules.encoders.modules import FrozenOpenCLIPImageEmbedder # bigG embedder
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+
+
+# In[2]:
+
+
+### 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)  
+
+data_type = torch.float16 # change depending on your mixed_precision
+num_devices = torch.cuda.device_count()
+if num_devices==0: num_devices = 1
+
+# First use "accelerate config" in terminal and setup using deepspeed stage 2 with CPU offloading!
+accelerator = Accelerator(split_batches=False, mixed_precision="fp16")
+if utils.is_interactive(): # set batch size here if using interactive notebook instead of submitting job
+    global_batch_size = batch_size = 16
+    
+else:
+    global_batch_size = os.environ["GLOBAL_BATCH_SIZE"]
+    batch_size = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices
+
+
+# In[3]:
+
+
+print("PID of this process =",os.getpid())
+device = accelerator.device
+print("device:",device)
+world_size = accelerator.state.num_processes
+distributed = not accelerator.state.distributed_type == 'NO'
+num_devices = torch.cuda.device_count()
+if num_devices==0 or not distributed: num_devices = 1
+num_workers = num_devices
+print(accelerator.state)
+
+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[4]:
+
+
+# if running this interactively, can specify jupyter_args here for argparser to use
+if utils.is_interactive():
+    model_name = "testing"
+    print("model_name:", model_name)
+    
+    # global_batch_size and batch_size should already be defined in the 2nd cell block
+    jupyter_args = f"--data_path=/weka/proj-medarc/shared/bold5000/nsd_general \
+                    --cache_dir=/weka/proj-medarc/shared/cache \
+                    --model_name={model_name} \
+                    --no-multi_subject --subj=1 --batch_size={batch_size} --num_sessions=15 \
+                    --hidden_dim=1024 --clip_scale=1. \
+                    --no-blurry_recon --blur_scale=.5  \
+                    --use_prior --prior_scale=30 \
+                    --n_blocks=4 --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --no-use_image_aug \
+                    --ckpt_interval=999 --no-ckpt_saving --no-wandb_log"
+    # --multisubject_ckpt=../train_logs/multisubject_subj01_1024_24bs_nolow
+
+    print(jupyter_args)
+    jupyter_args = jupyter_args.split()
+    
+    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[5]:
+
+
+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=os.getcwd(),
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--cache_dir", type=str, default=os.getcwd(),
+    help="Path to where misc. files downloaded from huggingface are stored. Defaults to current src directory.",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,3,4,5,6,7,8],
+    help="Validate on which subject?",
+)
+parser.add_argument(
+    "--multisubject_ckpt", type=str, default=None,
+    help="Path to pre-trained multisubject model to finetune a single subject from. multisubject must be False.",
+)
+parser.add_argument(
+    "--num_sessions", type=int, default=1,
+    help="Number of training sessions to include",
+)
+parser.add_argument(
+    "--use_prior",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to train diffusion prior (True) or just rely on retrieval part of the pipeline (False)",
+)
+parser.add_argument(
+    "--batch_size", type=int, default=16,
+    help="Batch size can be increased by 10x if only training retreival submodule and not diffusion prior",
+)
+parser.add_argument(
+    "--wandb_log",action=argparse.BooleanOptionalAction,default=False,
+    help="whether to log to wandb",
+)
+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(
+    "--blur_scale",type=float,default=.5,
+    help="multiply loss from blurry recons by this number",
+)
+parser.add_argument(
+    "--clip_scale",type=float,default=1.,
+    help="multiply contrastive loss by this number",
+)
+parser.add_argument(
+    "--prior_scale",type=float,default=30,
+    help="multiply diffusion prior loss by this",
+)
+parser.add_argument(
+    "--use_image_aug",action=argparse.BooleanOptionalAction,default=False,
+    help="whether to use image augmentation",
+)
+parser.add_argument(
+    "--num_epochs",type=int,default=150,
+    help="number of epochs of training",
+)
+parser.add_argument(
+    "--multi_subject",action=argparse.BooleanOptionalAction,default=False,
+)
+parser.add_argument(
+    "--new_test",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--n_blocks",type=int,default=4,
+)
+parser.add_argument(
+    "--hidden_dim",type=int,default=1024,
+)
+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)
+    
+# seed all random functions
+utils.seed_everything(seed)
+
+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 or blurry_recon:
+    import kornia
+    from kornia.augmentation.container import AugmentationSequential
+if use_image_aug:
+    img_augment = AugmentationSequential(
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1, p=0.3),
+        same_on_batch=False,
+        data_keys=["input"],
+    )
+    
+if multi_subject:
+    subj_list = np.arange(1,9)
+    subj_list = subj_list[subj_list != subj]
+else:
+    subj_list = [subj]
+
+print("subj_list", subj_list, "num_sessions", num_sessions)
+
+
+# # Prep data, models, and dataloaders
+
+# ### Creating wds dataloader, preload betas and all 73k possible images
+
+# In[6]:
+
+
+def my_split_by_node(urls): return urls
+num_voxels_list = []
+
+if multi_subject:
+    nsessions_allsubj=np.array([15, 15, 15, 9])
+    num_samples_per_epoch = (370*15) // num_devices 
+else:
+    num_samples_per_epoch = (370*num_sessions) // num_devices 
+
+print("dividing batch size by subj_list, which will then be concatenated across subj during training...") 
+batch_size = batch_size // len(subj_list)
+samples_per_session = 370
+
+
+# In[9]:
+
+
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.data.distributed import DistributedSampler
+from sklearn.model_selection import train_test_split
+
+class BoldDataset(Dataset):
+    def __init__(self, voxels, images, sessions_mask, indexes):
+        self.voxels = voxels
+        self.images = images
+        self.sessions_mask = sessions_mask
+        self.indexes = indexes
+        self.num_samples = len(voxels)
+    def __len__(self):
+        return self.num_samples
+    def __getitem__(self, idx):
+        return self.voxels[idx], self.images[idx], self.sessions_mask[idx], self.indexes[idx]
+  
+def get_voxels_and_images(subj, num_sessions, data_path, mode, random_state = 42):
+    """
+        subj: int, subject number
+        num_sessions: int, number of sessions to include, can be -1 for all
+        data_path: str, path to the data
+        mode: str, mode can be 'train', 'val', 'shared' and 'all' and "t&v"
+        random_state: int, random state for train_test_split
+    """
+
+    # load the voxels
+    all_voxels = torch.load(os.path.join(data_path, f"subj{subj}/wholebrain.pt"))
+    # all_voxels = torch.load(os.path.join(data_path, f"subj{subj}/voxels.pt"))
+    all_voxels = torch.where(torch.isnan(all_voxels), torch.zeros_like(all_voxels), all_voxels)
+    # load the other files
+    all_images = torch.load(os.path.join(data_path, f"subj{subj}/images_stimuli.pt"))
+    shared_images_mask = torch.load(os.path.join(data_path, f"subj{subj}/shared_images_mask.pt"))
+    sessions_mask = torch.load(os.path.join(data_path, f"subj{subj}/sessions_mask.pt"))
+    indexes = np.linspace(0, all_voxels.shape[0],num=all_voxels.shape[0], dtype=int)
+    # split the data into train and val with sklearn train_test_split
+    voxels_train, voxels_val, images_train, images_val, shared_images_mask_train, shared_images_mask_val, sessions_mask_train, sessions_mask_val, indexes_train, indexes_val = train_test_split(all_voxels, all_images, shared_images_mask, sessions_mask, indexes, test_size=0.1, random_state=random_state)
+
+    if mode == 'train':
+        # apply the shared images mask removing the ones in which is 1
+        voxels_train = voxels_train[shared_images_mask_train==0]
+        images_train = images_train[shared_images_mask_train==0]
+        sessions_mask_train = sessions_mask_train[shared_images_mask_train==0]
+        indexes_train = indexes_train[shared_images_mask_train==0]
+
+        # apply the sessions mask, the mask has values from 1 to 15, if num_sessions is 4, we will take the elements in which the value is 1 to 4
+        if num_sessions != -1:
+            voxels_train = voxels_train[sessions_mask_train <= num_sessions]
+            images_train = images_train[sessions_mask_train <= num_sessions]
+            indexes_train = indexes_train[sessions_mask_train <= num_sessions]
+            sessions_mask_train = sessions_mask_train[sessions_mask_train <= num_sessions]
+        
+        return voxels_train, images_train, sessions_mask_train, indexes_train
+
+    elif mode == 'val':
+        # apply the shared images mask removing the ones in which is 1
+        voxels_val = voxels_val[shared_images_mask_val==0]
+        images_val = images_val[shared_images_mask_val==0]
+        sessions_mask_val = sessions_mask_val[shared_images_mask_val==0]
+        indexes_val = indexes_val[shared_images_mask_val==0]
+
+        return voxels_val, images_val, sessions_mask_val, indexes_val
+    
+    elif mode == 'shared':
+        # apply the shared images mask taking only the ones in which is 1
+        voxels_shared = all_voxels[shared_images_mask==1]
+        images_shared = all_images[shared_images_mask==1]
+        indexes_shared = indexes[shared_images_mask==1]
+        sessions_mask_shared = sessions_mask[shared_images_mask==1]
+
+        return voxels_shared, images_shared, sessions_mask_shared, indexes_shared
+    
+    elif mode == 't&v':
+        voxels_tv = all_voxels[shared_images_mask==0]
+        images_tv = all_images[shared_images_mask==0]
+        indexes_tv = indexes[shared_images_mask==0]
+        sessions_mask_tv = sessions_mask[shared_images_mask==0]
+
+        return voxels_tv, images_tv, sessions_mask_tv, indexes_tv
+    
+    elif mode == 'all':
+        return all_voxels, all_images, sessions_mask, indexes
+
+# voxels, images, sessions_idx, sample_idx = get_voxels_and_images(subj, num_sessions, data_path, 'train')
+# print("voxels.shape", voxels.shape, "images.shape", images.shape, "sessions_idx.shape", sessions_idx.shape, "sample_idx.shape", sample_idx.shape)
+
+
+# In[11]:
+
+
+train_data = {}
+train_dl = {}
+num_voxels = {}
+voxels = {}
+for s in subj_list:
+    # TODO: Check num_sessions is working, replace -1
+    voxels_train, images_train, sessions_idx_train, sample_idx_train = get_voxels_and_images(s, -1, data_path, 'train')
+    train_data[f'subj0{s}'] = BoldDataset(voxels_train, images_train, sessions_idx_train, sample_idx_train)
+    train_dl[f'subj0{s}'] = DataLoader(train_data[f'subj0{s}'], batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True, drop_last=True)
+    print(f"Training with {num_sessions} sessions")
+    num_voxels[f'subj0{s}'] = voxels_train.shape[-1]
+    num_voxels_list.append(num_voxels[f'subj0{s}'])
+    # TODO: Delete this part if it's not needed and delete the variable declaration above as well
+    voxels[f'subj0{s}'] = voxels_train
+
+print("Loaded all subj train dls and betas!\n")
+
+    # if multi_subject:
+    #     train_url = f"{data_path}/wds/subj0{s}/train/" + "{0.." + f"{nsessions_allsubj[s-1]-1}" + "}.tar"
+    # else:
+    #     train_url = f"{data_path}/wds/subj0{s}/train/" + "{0.." + f"{num_sessions-1}" + "}.tar"
+    # print(train_url)
+    
+    # train_data[f'subj0{s}'] = wds.WebDataset(train_url,resampled=True,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[f'subj0{s}'] = torch.utils.data.DataLoader(train_data[f'subj0{s}'], batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+    # f = h5py.File(f'{data_path}/betas_all_subj0{s}_fp32_renorm.hdf5', 'r')
+    # betas = f['betas'][:]
+    # betas = torch.Tensor(betas).to("cpu").to(data_type)
+    # num_voxels_list.append(betas[0].shape[-1])
+    # num_voxels[f'subj0{s}'] = betas[0].shape[-1]
+    # voxels[f'subj0{s}'] = betas
+    # print(f"num_voxels for subj0{s}: {num_voxels[f'subj0{s}']}")
+
+print("Loaded all subj train dls and betas!\n")
+
+# Validate only on one subject
+if multi_subject: 
+    subj = subj_list[0] # cant validate on the actual held out person so picking first in subj_list
+# if not new_test: # using old test set from before full dataset released (used in original MindEye paper)
+#     if subj==3:
+#         num_test=2113
+#     elif subj==4:
+#         num_test=1985
+#     elif subj==6:
+#         num_test=2113
+#     elif subj==8:
+#         num_test=1985
+#     else:
+#         num_test=2770
+#     test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+# elif new_test: # using larger test set from after full dataset released
+#     if subj==3:
+#         num_test=2371
+#     elif subj==4:
+#         num_test=2188
+#     elif subj==6:
+#         num_test=2371
+#     elif subj==8:
+#         num_test=2188
+#     else:
+#         num_test=3000
+#     test_url = f"{data_path}/wds/subj0{subj}/new_test/" + "0.tar"
+# print(test_url)
+
+voxels_test, images_test, sessions_idx_test, sample_idx_test = get_voxels_and_images(subj, num_sessions, data_path, 'val')
+test_data = BoldDataset(voxels_test, images_test, sessions_idx_test, sample_idx_test)
+test_dl = DataLoader(test_data, batch_size=voxels_test.shape[0], shuffle=False, num_workers=num_workers, pin_memory=True)
+num_test = voxels_test.shape[0]
+# 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=num_test, shuffle=False, drop_last=True, pin_memory=True)
+print(f"Loaded test dl for subj{subj}!\n")
+
+
+# In[13]:
+
+
+num_samples_per_epoch = 0
+for dt_c in train_data:
+    num_samples_per_epoch = num_samples_per_epoch + train_data[dt_c].voxels.shape[0]
+# num_iterations_per_epoch = num_samples_per_epoch // (batch_size*len(subj_list))
+num_iterations_per_epoch = num_samples_per_epoch // (batch_size*len(subj_list))
+
+print("batch_size =", batch_size, "num_iterations_per_epoch =",num_iterations_per_epoch, "num_samples_per_epoch =",num_samples_per_epoch)
+
+
+# In[14]:
+
+
+# # now plot inside the traning loop
+# for batch in train_dl[f'subj0{subj}']:
+#     voxels, images, sessions_idx, sample_idx = batch
+#     print("voxels.shape", voxels.shape, "images.shape", images.shape, "sessions_idx.shape", sessions_idx.shape, "sample_idx.shape", sample_idx.shape)
+#     indexes_to_plot = [1,3,4,5,6]
+#     fig, axs = plt.subplots(1, len(indexes_to_plot), figsize=(20, 5))
+#     for i, idx in enumerate(indexes_to_plot):
+#         axs[i].imshow(images[idx].permute(1, 2, 0))
+#         axs[i].axis('off')
+#     plt.show()
+#     break
+
+# # now plot inside the traning loop
+# for batch in test_dl:
+#     voxels, images, sessions_idx, sample_idx = batch
+#     print("voxels.shape", voxels.shape, "images.shape", images.shape, "sessions_idx.shape", sessions_idx.shape, "sample_idx.shape", sample_idx.shape)
+#     indexes_to_plot = [1,3,4,5,6]
+#     fig, axs = plt.subplots(1, len(indexes_to_plot), figsize=(20, 5))
+#     for i, idx in enumerate(indexes_to_plot):
+#         axs[i].imshow(images[idx].permute(1, 2, 0))
+#         axs[i].axis('off')
+#     plt.show()
+#     break
+
+
+# In[15]:
+
+
+# train_data = {}
+# train_dl = {}
+# num_voxels = {}
+# voxels = {}
+# for s in subj_list:
+#     print(f"Training with {num_sessions} sessions")
+#     if multi_subject:
+#         train_url = f"{data_path}/wds/subj0{s}/train/" + "{0.." + f"{nsessions_allsubj[s-1]-1}" + "}.tar"
+#     else:
+#         train_url = f"{data_path}/wds/subj0{s}/train/" + "{0.." + f"{num_sessions-1}" + "}.tar"
+#     print(train_url)
+    
+#     train_data[f'subj0{s}'] = wds.WebDataset(train_url,resampled=True,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[f'subj0{s}'] = torch.utils.data.DataLoader(train_data[f'subj0{s}'], batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+#     f = h5py.File(f'{data_path}/betas_all_subj0{s}_fp32_renorm.hdf5', 'r')
+#     betas = f['betas'][:]
+#     betas = torch.Tensor(betas).to("cpu").to(data_type)
+#     num_voxels_list.append(betas[0].shape[-1])
+#     num_voxels[f'subj0{s}'] = betas[0].shape[-1]
+#     voxels[f'subj0{s}'] = betas
+#     print(f"num_voxels for subj0{s}: {num_voxels[f'subj0{s}']}")
+
+# print("Loaded all subj train dls and betas!\n")
+
+# # Validate only on one subject
+# if multi_subject: 
+#     subj = subj_list[0] # cant validate on the actual held out person so picking first in subj_list
+# if not new_test: # using old test set from before full dataset released (used in original MindEye paper)
+#     if subj==3:
+#         num_test=2113
+#     elif subj==4:
+#         num_test=1985
+#     elif subj==6:
+#         num_test=2113
+#     elif subj==8:
+#         num_test=1985
+#     else:
+#         num_test=2770
+#     test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+# elif new_test: # using larger test set from after full dataset released
+#     if subj==3:
+#         num_test=2371
+#     elif subj==4:
+#         num_test=2188
+#     elif subj==6:
+#         num_test=2371
+#     elif subj==8:
+#         num_test=2188
+#     else:
+#         num_test=3000
+#     test_url = f"{data_path}/wds/subj0{subj}/new_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=num_test, shuffle=False, drop_last=True, pin_memory=True)
+# print(f"Loaded test dl for subj{subj}!\n")
+
+
+# In[16]:
+
+
+# # Load 73k NSD images
+# f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+# images = f['images']
+# print("Loaded all 73k possible NSD images to cpu!", images.shape)
+
+
+# ## Load models
+
+# ### CLIP image embeddings  model
+
+# In[17]:
+
+
+clip_img_embedder = FrozenOpenCLIPImageEmbedder(
+    arch="ViT-bigG-14",
+    version="laion2b_s39b_b160k",
+    output_tokens=True,
+    only_tokens=True,
+)
+clip_img_embedder.to(device)
+
+clip_seq_dim = 256
+clip_emb_dim = 1664
+
+
+# ### SD VAE
+
+# In[18]:
+
+
+if blurry_recon:
+    from diffusers import AutoencoderKL    
+    autoenc = AutoencoderKL(
+        down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D'],
+        up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D'],
+        block_out_channels=[128, 256, 512, 512],
+        layers_per_block=2,
+        sample_size=256,
+    )
+    ckpt = torch.load(f'{cache_dir}/sd_image_var_autoenc.pth')
+    autoenc.load_state_dict(ckpt)
+    
+    autoenc.eval()
+    autoenc.requires_grad_(False)
+    autoenc.to(device)
+    utils.count_params(autoenc)
+    
+    from autoencoder.convnext import ConvnextXL
+    cnx = ConvnextXL(f'{cache_dir}/convnext_xlarge_alpha0.75_fullckpt.pth')
+    cnx.requires_grad_(False)
+    cnx.eval()
+    cnx.to(device)
+    
+    mean = torch.tensor([0.485, 0.456, 0.406]).to(device).reshape(1,3,1,1)
+    std = torch.tensor([0.228, 0.224, 0.225]).to(device).reshape(1,3,1,1)
+    
+    blur_augs = AugmentationSequential(
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.8),
+        kornia.augmentation.RandomGrayscale(p=0.1),
+        kornia.augmentation.RandomSolarize(p=0.1),
+        kornia.augmentation.RandomResizedCrop((224,224), scale=(.9,.9), ratio=(1,1), p=1.0),
+        data_keys=["input"],
+    )
+
+
+# ### MindEye modules
+
+# In[19]:
+
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+model
+
+
+# In[20]:
+
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer to enable regularization
+    def __init__(self, input_sizes, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linears = torch.nn.ModuleList([
+                torch.nn.Linear(input_size, out_features) for input_size in input_sizes
+            ])
+    def forward(self, x, subj_idx):
+        out = self.linears[subj_idx](x[:,0]).unsqueeze(1)
+        return out
+        
+model.ridge = RidgeRegression(num_voxels_list, out_features=hidden_dim)
+utils.count_params(model.ridge)
+utils.count_params(model)
+
+# test on subject 1 with fake data
+b = torch.randn((2,1,num_voxels_list[0]))
+print(b.shape, model.ridge(b,0).shape)
+
+
+# In[21]:
+
+
+from models import BrainNetwork
+model.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, seq_len=1, n_blocks=n_blocks,
+                          clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, 
+                          blurry_recon=blurry_recon, clip_scale=clip_scale)
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+# test that the model works on some fake data
+b = torch.randn((2,1,hidden_dim))
+print("b.shape",b.shape)
+
+backbone_, clip_, blur_ = model.backbone(b)
+print(backbone_.shape, clip_.shape, blur_[0].shape, blur_[1].shape)
+
+
+# ### Adding diffusion prior + unCLIP if use_prior=True
+
+# In[22]:
+
+
+if use_prior:
+    from models import *
+
+    # setup diffusion prior network
+    out_dim = clip_emb_dim
+    depth = 6
+    dim_head = 52
+    heads = clip_emb_dim//52 # heads * dim_head = clip_emb_dim
+    timesteps = 100
+
+    prior_network = PriorNetwork(
+            dim=out_dim,
+            depth=depth,
+            dim_head=dim_head,
+            heads=heads,
+            causal=False,
+            num_tokens = clip_seq_dim,
+            learned_query_mode="pos_emb"
+        )
+
+    model.diffusion_prior = BrainDiffusionPrior(
+        net=prior_network,
+        image_embed_dim=out_dim,
+        condition_on_text_encodings=False,
+        timesteps=timesteps,
+        cond_drop_prob=0.2,
+        image_embed_scale=None,
+    )
+    
+    utils.count_params(model.diffusion_prior)
+    utils.count_params(model)
+
+
+# ### Setup optimizer / lr / ckpt saving
+
+# In[23]:
+
+
+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},
+]
+if use_prior:
+    opt_grouped_parameters.extend([
+        {'params': [p for n, p in model.diffusion_prior.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},
+        {'params': [p for n, p in model.diffusion_prior.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_iterations_per_epoch)),
+        last_epoch=-1
+    )
+elif lr_scheduler_type == 'cycle':
+    total_steps=int(np.floor(num_epochs*num_iterations_per_epoch))
+    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'
+    if accelerator.is_main_process:
+        unwrapped_model = accelerator.unwrap_model(model)
+        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)
+    print(f"\n---saved {outdir}/{tag} ckpt!---\n")
+
+def load_ckpt(tag,load_lr=True,load_optimizer=True,load_epoch=True,strict=True,outdir=outdir,multisubj_loading=False): 
+    print(f"\n---loading {outdir}/{tag}.pth ckpt---\n")
+    checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+    state_dict = checkpoint['model_state_dict']
+    if multisubj_loading: # remove incompatible ridge layer that will otherwise error
+        state_dict.pop('ridge.linears.0.weight',None)
+    model.load_state_dict(state_dict, strict=strict)
+    if load_epoch:
+        globals()["epoch"] = checkpoint['epoch']
+        print("Epoch",epoch)
+    if load_optimizer:
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    if load_lr:
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    del checkpoint
+
+print("\nDone with model preparations!")
+num_params = utils.count_params(model)
+
+
+# # Weights and Biases
+
+# In[24]:
+
+
+if local_rank==0 and wandb_log: # only use main process for wandb logging
+    import wandb
+    wandb_project = 'mindeye'
+    print(f"wandb {wandb_project} run {model_name}")
+    # need to configure wandb beforehand in terminal with "wandb init"!
+    wandb_config = {
+      "model_name": model_name,
+      "global_batch_size": global_batch_size,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "num_sessions": num_sessions,
+      "num_params": num_params,
+      "clip_scale": clip_scale,
+      "prior_scale": prior_scale,
+      "blur_scale": blur_scale,
+      "use_image_aug": use_image_aug,
+      "max_lr": max_lr,
+      "mixup_pct": mixup_pct,
+      "num_samples_per_epoch": num_samples_per_epoch,
+      "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)
+    print("wandb_id:",model_name)
+    wandb.init(
+        id=model_name,
+        project=wandb_project,
+        name=model_name,
+        config=wandb_config,
+        resume="allow",
+    )
+else:
+    wandb_log = False
+
+
+# # Main
+
+# In[25]:
+
+
+epoch = 0
+losses, test_losses, lrs = [], [], []
+best_test_loss = 1e9
+torch.cuda.empty_cache()
+
+
+# In[26]:
+
+
+# load multisubject stage1 ckpt if set
+if multisubject_ckpt is not None:
+    load_ckpt("last",outdir=multisubject_ckpt,load_lr=False,load_optimizer=False,load_epoch=False,strict=False,multisubj_loading=True)
+
+
+# In[27]:
+
+
+train_dls = [train_dl[f'subj0{s}'] for s in subj_list]
+
+model, optimizer, *train_dls, lr_scheduler = accelerator.prepare(model, optimizer, *train_dls, lr_scheduler)
+# leaving out test_dl since we will only have local_rank 0 device do evals
+
+
+# In[28]:
+
+
+for voxels, images, sessions_mask, sample_idx in train_dls[0]:
+    print(voxels.shape, images.shape, sessions_mask.shape, sample_idx.shape)
+    break
+
+
+# In[29]:
+
+
+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()
+soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))
+
+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.
+    
+    recon_cossim = 0.
+    test_recon_cossim = 0.
+    recon_mse = 0.
+    test_recon_mse = 0.
+
+    loss_clip_total = 0.
+    loss_blurry_total = 0.
+    loss_blurry_cont_total = 0.
+    test_loss_clip_total = 0.
+    
+    loss_prior_total = 0.
+    test_loss_prior_total = 0.
+
+    blurry_pixcorr = 0.
+    test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1
+
+    # pre-load all batches for this epoch (it's MUCH faster to pre-load in bulk than to separate loading per batch)
+    voxel_iters = {} # empty dict because diff subjects have differing # of voxels
+    image_iters = torch.zeros(num_iterations_per_epoch, batch_size*len(subj_list), 3, 224, 224).float()
+    annot_iters = {}
+    perm_iters, betas_iters, select_iters = {}, {}, {}
+    for s, train_dl in enumerate(train_dls):
+        with torch.cuda.amp.autocast(dtype=data_type):
+            iter = -1
+            for voxels, images, sessions_mask, sample_idx in train_dl:
+            # for behav0, past_behav0, future_behav0, old_behav0 in train_dl: 
+                # Load images to cpu from hdf5 (requires sorted indexing)
+                # image_idx = behav0[:,0,0].cpu().long().numpy()
+                # image0, image_sorted_idx = np.unique(image_idx, return_index=True)                
+                # if len(image0) != len(image_idx): # hdf5 cant handle duplicate indexing
+                #     continue
+                iter += 1
+                # image0 = torch.tensor(images[image0], dtype=data_type)
+                # image_iters[iter,s*batch_size:s*batch_size+batch_size] = image0
+                image_iters[iter,s*batch_size:s*batch_size+batch_size] = images
+
+                # Load voxels for current batch, matching above indexing
+                # voxel_idx = behav0[:,0,5].cpu().long().numpy()
+                # voxel_sorted_idx = voxel_idx[image_sorted_idx]
+                # voxel0 = voxels[f'subj0{subj_list[s]}'][voxel_sorted_idx]
+                # voxel0 = torch.Tensor(voxel0).unsqueeze(1)
+                # print(voxels.shape)
+                voxel0 = voxels.unsqueeze(1)
+                # print(voxel0.shape)
+                if epoch < int(mixup_pct * num_epochs):
+                    voxel0, perm, betas, select = utils.mixco(voxel0)
+                    # voxel0, perm, betas, select = utils.mixco(voxel0)
+                    perm_iters[f"subj0{subj_list[s]}_iter{iter}"] = perm
+                    betas_iters[f"subj0{subj_list[s]}_iter{iter}"] = betas
+                    select_iters[f"subj0{subj_list[s]}_iter{iter}"] = select
+                    
+                voxel_iters[f"subj0{subj_list[s]}_iter{iter}"] = voxel0
+                # print(voxel0.shape)
+
+                if iter >= num_iterations_per_epoch-1:
+                    break
+
+    # you now have voxel_iters and image_iters with num_iterations_per_epoch batches each
+    for train_i in range(num_iterations_per_epoch):
+        with torch.cuda.amp.autocast(dtype=data_type):
+            optimizer.zero_grad()
+            loss=0.
+
+            voxel_list = [voxel_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list]
+            image = image_iters[train_i].detach()
+            image = image.to(device)
+
+            if use_image_aug: 
+                image = img_augment(image)
+
+            clip_target = clip_img_embedder(image)
+            assert not torch.any(torch.isnan(clip_target))
+
+            if epoch < int(mixup_pct * num_epochs):
+                perm_list = [perm_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list]
+                perm = torch.cat(perm_list, dim=0)
+                betas_list = [betas_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list]
+                betas = torch.cat(betas_list, dim=0)
+                select_list = [select_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list]
+                select = torch.cat(select_list, dim=0)
+
+            voxel_ridge_list = [model.ridge(voxel_list[si],si) for si,s in enumerate(subj_list)]
+            voxel_ridge = torch.cat(voxel_ridge_list, dim=0)
+
+            backbone, clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+
+            if clip_scale>0:
+                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 use_prior:
+                loss_prior, prior_out = model.diffusion_prior(text_embed=backbone, image_embed=clip_target)
+                loss_prior_total += loss_prior.item()
+                loss_prior *= prior_scale
+                loss += loss_prior
+
+                recon_cossim += nn.functional.cosine_similarity(prior_out, clip_target).mean().item()
+                recon_mse += mse(prior_out, clip_target).item()
+
+            if clip_scale>0:
+                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:     
+                image_enc_pred, transformer_feats = blurry_image_enc_
+
+                image_enc = autoenc.encode(2*image-1).latent_dist.mode() * 0.18215
+                loss_blurry = l1(image_enc_pred, image_enc)
+                loss_blurry_total += loss_blurry.item()
+
+                if epoch < int(mixup_pct * num_epochs):
+                    image_enc_shuf = image_enc[perm]
+                    betas_shape = [-1] + [1]*(len(image_enc.shape)-1)
+                    image_enc[select] = image_enc[select] * betas[select].reshape(*betas_shape) + \
+                        image_enc_shuf[select] * (1 - betas[select]).reshape(*betas_shape)
+
+                image_norm = (image - mean)/std
+                image_aug = (blur_augs(image) - mean)/std
+                _, cnx_embeds = cnx(image_norm)
+                _, cnx_aug_embeds = cnx(image_aug)
+
+                cont_loss = utils.soft_cont_loss(
+                    nn.functional.normalize(transformer_feats.reshape(-1, transformer_feats.shape[-1]), dim=-1),
+                    nn.functional.normalize(cnx_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),
+                    nn.functional.normalize(cnx_aug_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),
+                    temp=0.2)
+                loss_blurry_cont_total += cont_loss.item()
+
+                loss += (loss_blurry + 0.1*cont_loss) * blur_scale #/.18215
+
+            if clip_scale>0:
+                # forward and backward top 1 accuracy        
+                labels = torch.arange(len(clip_voxels_norm)).to(clip_voxels_norm.device) 
+                fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item()
+                bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()
+                # print(utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item(), utils.topk(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(image)), size=len(image)//5, replace=False)
+                    blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample/ 2 + 0.5).clamp(0,1)
+                    pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)
+                    blurry_pixcorr += pixcorr.item()
+
+            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):  
+            for test_i, (voxels, images, sessions_mask, sample_idx) 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[f'subj0{subj}'][behav[:,0,5].cpu().long()].unsqueeze(1)
+                #     voxel = voxels.unsqueeze(1)
+                #     # image = behav[:,0,0].cpu().long()
+                #     image = images
+
+                #     unique_image, sort_indices = torch.unique(image, return_inverse=True)
+                #     for im in unique_image:
+                #         locs = torch.where(im == image)[0]
+                #         if len(locs)==1:
+                #             locs = locs.repeat(3)
+                #         elif len(locs)==2:
+                #             locs = locs.repeat(2)[:3]
+                #         assert len(locs)==3
+                #         if test_image is None:
+                #             test_image = torch.Tensor(images[im][None])
+                #             test_voxel = voxel[locs][None]
+                #         else:
+                #             test_image = torch.vstack((test_image, torch.Tensor(images[im][None])))
+                #             test_voxel = torch.vstack((test_voxel, voxel[locs][None]))
+                test_image = images
+                test_voxel = voxels.unsqueeze(1)
+
+                loss=0.
+                            
+                # test_indices = torch.arange(len(test_voxel))[:] 
+                # voxel = test_voxel[test_indices].to(device)
+                # image = test_image[test_indices].to(device)
+                # assert len(image) == 100
+                image = test_image.to(device)
+                voxel = test_voxel.to(device)
+
+                clip_target = clip_img_embedder(image.float())
+
+                # REPETITION AVERAGING INFERENCE DOESN'T APPLY SINCE WE'RE USING VALIDATION SET
+                # for rep in range(3):
+                #     voxel_ridge = model.ridge(voxel[:,rep],0) # 0th index of subj_list
+                #     backbone0, clip_voxels0, blurry_image_enc_ = model.backbone(voxel_ridge)
+                #     if rep==0:
+                #         clip_voxels = clip_voxels0
+                #         backbone = backbone0
+                #     else:
+                #         clip_voxels += clip_voxels0
+                #         backbone += backbone0
+                # clip_voxels /= 3
+                # backbone /= 3
+
+                voxel_ridge = model.ridge(voxel,0) # 0th index of subj_list
+
+                backbone, clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+
+                if clip_scale>0:
+                    clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+                    clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+                
+                # for some evals, only doing a subset of the samples per batch because of computational cost
+                random_samps = np.random.choice(np.arange(len(image)), size=len(image)//5, replace=False)
+                
+                if use_prior:
+                    loss_prior, contaminated_prior_out = model.diffusion_prior(text_embed=backbone[random_samps], image_embed=clip_target[random_samps])
+                    test_loss_prior_total += loss_prior.item()
+                    loss_prior *= prior_scale
+                    loss += loss_prior
+                    
+                if clip_scale>0:
+                    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:
+                    image_enc_pred, _ = blurry_image_enc_
+                    blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample / 2 + 0.5).clamp(0,1)
+                    pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)
+                    test_blurry_pixcorr += pixcorr.item()
+
+                if clip_scale>0:
+                    # forward and backward top 1 accuracy        
+                    labels = torch.arange(len(clip_voxels_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())
+
+            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),
+                "train/loss_blurry_cont_total": loss_blurry_cont_total / (train_i + 1),
+                "test/loss_clip_total": test_loss_clip_total / (test_i + 1),
+                "train/blurry_pixcorr": blurry_pixcorr / (train_i + 1),
+                "test/blurry_pixcorr": test_blurry_pixcorr / (test_i + 1),
+                "train/recon_cossim": recon_cossim / (train_i + 1),
+                "test/recon_cossim": test_recon_cossim / (test_i + 1),
+                "train/recon_mse": recon_mse / (train_i + 1),
+                "test/recon_mse": test_recon_mse / (test_i + 1),
+                "train/loss_prior": loss_prior_total / (train_i + 1),
+                "test/loss_prior": test_loss_prior_total / (test_i + 1),
+                }
+
+            # if finished training, save jpg recons if they exist
+            if (epoch == num_epochs-1) or (epoch % ckpt_interval == 0):
+                if blurry_recon:    
+                    image_enc = autoenc.encode(2*image[:4]-1).latent_dist.mode() * 0.18215
+                    # 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(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(image_enc_pred[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                        axes[jj].axis('off')
+
+                    if wandb_log:
+                        logs[f"test/blur_recons"] = wandb.Image(fig, caption=f"epoch{epoch:03d}")
+                        plt.close()
+                    else:
+                        plt.show()
+
+            progress_bar.set_postfix(**logs)
+
+            if wandb_log: wandb.log(logs)
+            
+    # Save model checkpoint and reconstruct
+    if (ckpt_saving) and (epoch % ckpt_interval == 0):
+        save_ckpt(f'last')
+
+    # wait for other GPUs to catch up if needed
+    accelerator.wait_for_everyone()
+    torch.cuda.empty_cache()
+
+print("\n===Finished!===\n")
+if ckpt_saving:
+    save_ckpt(f'last')
+
+
+# In[ ]:
+
+
+# voxel_iters.keys()
+
+
+# In[ ]:
+
+
+# # create a new dl
+# test_voxels, test_images, test_sessions_idx, test_sample_idx = get_voxels_and_images(subj, num_sessions, data_path, 'shared')
+# test_data = BoldDataset(test_voxels, test_images, test_sessions_idx, test_sample_idx)
+# test_dl = DataLoader(test_data, batch_size=test_voxels.shape[0], shuffle=False, num_workers=num_workers, pin_memory=True)
+
+
+# In[ ]:
+
+
+aaa
+
+
+# In[ ]:
+
+
+# # find the unique images in the test set
+# unique_images = torch.unique(test_images, dim=0)
+# print("unique_images.shape", unique_images.shape)
+
+
+# In[ ]:
+
+
+v1_images = torch.load('/weka/proj-medarc/shared/bold5000/v1/subj1/images_stimuli.pt')
+wb_images = torch.load('/weka/proj-medarc/shared/bold5000/formated/subj1/images_stimuli.pt')
+
+
+# In[ ]:
+
+
+wb_images.device
+
+
+# In[ ]:
+
+
+# wb_images = torch.load('/weka/proj-medarc/shared/bold5000/formated/subj1/images_stimuli.pt') # (5254, 3, 224, 224)
+
+# iterate one in train_dl
+for b1_voxels, b1_images, b1_sessions_idx, b1_sample_idx in test_dl:
+    print(b1_voxels.shape, b1_images.shape, b1_sessions_idx.shape, b1_sample_idx.shape)
+    break
+
+b1_sample = b1_images[1].to('cpu') # (3, 224, 224)
+# check the index of the first image b1_images[1] in wb_images
+
+mask = (wb_images == b1_sample.unsqueeze(0))
+# Reduce over the last three dimensions one by one:
+mask = mask.all(dim=3).all(dim=2).all(dim=1)
+
+indices = torch.where(mask)[0]
+if len(indices) > 0:
+    print("Found b1_images[1] at index(es):", indices.tolist())
+else:
+    print("No exact match found.")
+
+
+# In[ ]:
+
+
+v1_images[4095,0,0,10:100], b1_images[1,0,0,10:100], wb_voxels[4095,1000:1020], b1_voxels[1,1000:1020]
+
+
+# In[ ]:
+
+
+b1_voxels.min()
+
+
+# In[ ]:
+
+
+wb_images.shape
+
+
+# In[ ]:
+
+
+v1_voxels = torch.load(f'/weka/proj-medarc/shared/bold5000/v1/subj1/voxels.pt')
+wb_voxels = torch.load(f'/weka/proj-medarc/shared/bold5000/formated/subj1/wholebrain.pt')
+
+
+# In[ ]:
+
+
+v1_voxels.max(), v1_voxels.min(), wb_voxels.max(), wb_voxels.min()
+
+
+# In[ ]:
+
+
+# for i in train_dl
+
+
+# In[ ]:
+
+
+from tqdm import tqdm
+import torch
+
+# Extract the first example from each tensor
+v1_example = v1_voxels[0]  # shape: [1685]
+wb_example = wb_voxels[0]  # shape: [174019]
+
+# Round both sets of voxels to 4 decimal places
+v1_rounded = torch.round(v1_example * 1000) / 1000
+wb_rounded = torch.round(wb_example * 1000) / 1000
+
+# Convert wb_rounded to a set for O(1) membership checks
+wb_set = set(wb_rounded.tolist())
+
+count = 0
+# Iterate over each voxel in v1_rounded and check if it exists in wb_rounded
+for voxel in tqdm(v1_rounded, desc="Checking V1 voxels"):
+    if voxel.item() in wb_set:
+        count += 1
+
+print(f"Number of V1 voxels present in the WB example (up to 4 decimals): {count}")
+
+
+# In[ ]:
+
+
+
+
+
+# In[ ]:
+
+
+# Extract the first example from each tensor
+v1_example = v1_voxels[0]        # shape: [1685]
+wb_example = wb_voxels[0]        # shape: [174019]
+
+# Convert wb_example to a set for O(1) membership checks
+wb_set = set(wb_example.tolist())
+
+count = 0
+# Iterate over each voxel in v1_example and check if it exists in wb_example
+for voxel in tqdm(v1_example, desc="Checking V1 voxels"):
+    value = voxel.item()
+    if count == 0:
+        value = 1.0938
+    
+    if value in wb_set:
+        count += 1
+
+print(f"Number of V1 voxels present in the WB example: {count}")
+
+
+# In[ ]:
+
+
+v1_voxels = torch.load(f'/weka/proj-medarc/shared/bold5000/v1/subj1/voxels.pt')
+wb_voxels = torch.load(f'/weka/proj-medarc/shared/bold5000/formated/subj1/wholebrain.pt')
+
+
+# In[ ]:
+
+
+v1_voxels.shape, wb_voxels[0,1006:1060]
+
+
+# In[ ]:
+
+
+v1_voxels.shape, wb_voxels[0,1006:1060]
+
+
+# In[ ]:
+
+
+# # plot train images
+indexes_to_plot = [1,3,4,5,6]
+fig, axs = plt.subplots(1, len(indexes_to_plot), figsize=(20, 5))
+for i, idx in enumerate(indexes_to_plot):
+    axs[i].imshow(v1_images[idx].permute(1, 2, 0))
+    axs[i].axis('off')
+plt.show()
+
+
+# In[ ]:
+
+
+# # plot train images
+indexes_to_plot = [1,3,4,5,6]
+fig, axs = plt.subplots(1, len(indexes_to_plot), figsize=(20, 5))
+for i, idx in enumerate(indexes_to_plot):
+    axs[i].imshow(wb_images[idx].permute(1, 2, 0))
+    axs[i].axis('off')
+plt.show()
+
+
+# In[ ]:
+
+
+# train_i
+
+
+# In[ ]:
+
+
+# fwd_percent_correct/
+
+
+# In[ ]:
+
+
+plt.plot(losses)
+plt.show()
+plt.plot(test_losses)
+plt.show()
+

MindEyeV2/src/TrainBold5000-Copy1.ipynb

@@ -0,0 +1,1669 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b0f0f4f3",
+   "metadata": {},
+   "source": [
+    "# Import packages & functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "5bad764b-45c1-45ce-a716-8d055e09821a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "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 string\n",
+    "import h5py\n",
+    "from tqdm import tqdm\n",
+    "import webdataset as wds\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torchvision import transforms\n",
+    "from accelerate import Accelerator\n",
+    "\n",
+    "# SDXL unCLIP requires code from https://github.com/Stability-AI/generative-models/tree/main\n",
+    "sys.path.append('generative_models/')\n",
+    "import sgm\n",
+    "from generative_models.sgm.modules.encoders.modules import FrozenOpenCLIPImageEmbedder # bigG embedder\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"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "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",
+    "data_type = torch.float16 # change depending on your mixed_precision\n",
+    "num_devices = torch.cuda.device_count()\n",
+    "if num_devices==0: num_devices = 1\n",
+    "\n",
+    "# First use \"accelerate config\" in terminal and setup using deepspeed stage 2 with CPU offloading!\n",
+    "accelerator = Accelerator(split_batches=False, mixed_precision=\"fp16\")\n",
+    "if utils.is_interactive(): # set batch size here if using interactive notebook instead of submitting job\n",
+    "    global_batch_size = batch_size = 8\n",
+    "else:\n",
+    "    global_batch_size = os.environ[\"GLOBAL_BATCH_SIZE\"]\n",
+    "    batch_size = int(os.environ[\"GLOBAL_BATCH_SIZE\"]) // num_devices"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "b767ab6f-d4a9-47a5-b3bf-f56bf6760c0c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "PID of this process = 2098275\n",
+      "device: cuda\n",
+      "Distributed environment: DistributedType.NO\n",
+      "Num processes: 1\n",
+      "Process index: 0\n",
+      "Local process index: 0\n",
+      "Device: cuda\n",
+      "\n",
+      "Mixed precision type: fp16\n",
+      "\n",
+      "distributed = False num_devices = 1 local rank = 0 world size = 1 data_type = torch.float16\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(\"PID of this process =\",os.getpid())\n",
+    "device = accelerator.device\n",
+    "print(\"device:\",device)\n",
+    "world_size = accelerator.state.num_processes\n",
+    "distributed = not accelerator.state.distributed_type == 'NO'\n",
+    "num_devices = torch.cuda.device_count()\n",
+    "if num_devices==0 or not distributed: num_devices = 1\n",
+    "num_workers = num_devices\n",
+    "print(accelerator.state)\n",
+    "\n",
+    "print(\"distributed =\",distributed, \"num_devices =\", num_devices, \"local rank =\", local_rank, \"world size =\", world_size, \"data_type =\", data_type)\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": 4,
+   "id": "2b61fec7-72a0-4b67-86da-1375f1d9fbd3",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "model_name: testingbold5000\n",
+      "--data_path=/weka/proj-medarc/shared/bold5000/v1                     --cache_dir=/weka/proj-medarc/shared/cache                     --model_name=testingbold5000                     --no-multi_subject --subj=1 --batch_size=8 --num_sessions=14                     --hidden_dim=1024 --clip_scale=1.                     --no-blurry_recon --blur_scale=.5                      --use_prior --prior_scale=30                     --n_blocks=4 --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --no-use_image_aug                     --ckpt_interval=999 --no-ckpt_saving --wandb_log\n"
+     ]
+    }
+   ],
+   "source": [
+    "# if running this interactively, can specify jupyter_args here for argparser to use\n",
+    "if utils.is_interactive():\n",
+    "    model_name = \"testingbold5000\"\n",
+    "    print(\"model_name:\", model_name)\n",
+    "    \n",
+    "    # global_batch_size and batch_size should already be defined in the 2nd cell block\n",
+    "    jupyter_args = f\"--data_path=/weka/proj-medarc/shared/bold5000/v1 \\\n",
+    "                    --cache_dir=/weka/proj-medarc/shared/cache \\\n",
+    "                    --model_name={model_name} \\\n",
+    "                    --no-multi_subject --subj=1 --batch_size={batch_size} --num_sessions=14 \\\n",
+    "                    --hidden_dim=1024 --clip_scale=1. \\\n",
+    "                    --no-blurry_recon --blur_scale=.5  \\\n",
+    "                    --use_prior --prior_scale=30 \\\n",
+    "                    --n_blocks=4 --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --no-use_image_aug \\\n",
+    "                    --ckpt_interval=999 --no-ckpt_saving --wandb_log\"\n",
+    "    # --multisubject_ckpt=../train_logs/multisubject_subj01_1024_24bs_nolow\n",
+    "\n",
+    "    print(jupyter_args)\n",
+    "    jupyter_args = jupyter_args.split()\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": 5,
+   "id": "2028bdf0-2f41-46d9-b6e7-86b870dbf16c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "subj_list [1] num_sessions 14\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=os.getcwd(),\n",
+    "    help=\"Path to where NSD data is stored / where to download it to\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--cache_dir\", type=str, default=os.getcwd(),\n",
+    "    help=\"Path to where misc. files downloaded from huggingface are stored. Defaults to current src directory.\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--subj\",type=int, default=1, choices=[1,2,3,4,5,6,7,8],\n",
+    "    help=\"Validate on which subject?\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--multisubject_ckpt\", type=str, default=None,\n",
+    "    help=\"Path to pre-trained multisubject model to finetune a single subject from. multisubject must be False.\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--num_sessions\", type=int, default=1,\n",
+    "    help=\"Number of training sessions to include\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--use_prior\",action=argparse.BooleanOptionalAction,default=True,\n",
+    "    help=\"whether to train diffusion prior (True) or just rely on retrieval part of the pipeline (False)\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--batch_size\", type=int, default=16,\n",
+    "    help=\"Batch size can be increased by 10x if only training retreival submodule and not diffusion prior\",\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",
+    "    \"--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",
+    "    \"--blurry_recon\",action=argparse.BooleanOptionalAction,default=True,\n",
+    "    help=\"whether to output blurry reconstructions\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--blur_scale\",type=float,default=.5,\n",
+    "    help=\"multiply loss from blurry recons by this number\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--clip_scale\",type=float,default=1.,\n",
+    "    help=\"multiply contrastive loss by this number\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--prior_scale\",type=float,default=30,\n",
+    "    help=\"multiply diffusion prior loss by this\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--use_image_aug\",action=argparse.BooleanOptionalAction,default=False,\n",
+    "    help=\"whether to use image augmentation\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--num_epochs\",type=int,default=150,\n",
+    "    help=\"number of epochs of training\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--multi_subject\",action=argparse.BooleanOptionalAction,default=False,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--new_test\",action=argparse.BooleanOptionalAction,default=True,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--n_blocks\",type=int,default=4,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--hidden_dim\",type=int,default=1024,\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",
+    "\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",
+    "# seed all random functions\n",
+    "utils.seed_everything(seed)\n",
+    "\n",
+    "outdir = os.path.abspath(f'../train_logs/{model_name}')\n",
+    "if not os.path.exists(outdir) and ckpt_saving:\n",
+    "    os.makedirs(outdir,exist_ok=True)\n",
+    "    \n",
+    "if use_image_aug or blurry_recon:\n",
+    "    import kornia\n",
+    "    from kornia.augmentation.container import AugmentationSequential\n",
+    "if use_image_aug:\n",
+    "    img_augment = AugmentationSequential(\n",
+    "        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1, p=0.3),\n",
+    "        same_on_batch=False,\n",
+    "        data_keys=[\"input\"],\n",
+    "    )\n",
+    "    \n",
+    "if multi_subject:\n",
+    "    subj_list = np.arange(1,9)\n",
+    "    subj_list = subj_list[subj_list != subj]\n",
+    "else:\n",
+    "    subj_list = [subj]\n",
+    "\n",
+    "print(\"subj_list\", subj_list, \"num_sessions\", num_sessions)"
+   ]
+  },
+  {
+   "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": [
+    "### Creating wds dataloader, preload betas and all 73k possible images"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "aefe7c27-ab39-4b2c-90f4-480f4087b7ab",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "dividing batch size by subj_list, which will then be concatenated across subj during training...\n",
+      "batch_size = 8 num_iterations_per_epoch = 656 num_samples_per_epoch = 5254\n"
+     ]
+    }
+   ],
+   "source": [
+    "def my_split_by_node(urls): return urls\n",
+    "num_voxels_list = []\n",
+    "\n",
+    "n_trials_per_subject = [5254, 5254, 5254, 3108] \n",
+    "\n",
+    "if multi_subject:\n",
+    "    nsessions_allsubj=np.array([15,15,15,9])\n",
+    "    num_samples_per_epoch = (750*40) // num_devices \n",
+    "else:\n",
+    "    num_samples_per_epoch = n_trials_per_subject[subj_list[0]] // num_devices \n",
+    "\n",
+    "print(\"dividing batch size by subj_list, which will then be concatenated across subj during training...\") \n",
+    "batch_size = batch_size // len(subj_list)\n",
+    "\n",
+    "num_iterations_per_epoch = num_samples_per_epoch // (batch_size*len(subj_list))\n",
+    "\n",
+    "print(\"batch_size =\", batch_size, \"num_iterations_per_epoch =\",num_iterations_per_epoch, \"num_samples_per_epoch =\",num_samples_per_epoch)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "8e8a3cc6-0627-4b61-b826-a6d2df93f122",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# load the data\n",
+    "images1 = torch.load('/weka/proj-medarc/shared/bold5000/v1/subj1/images_stimuli.pt')\n",
+    "voxelsv1 = torch.load('/weka/proj-medarc/shared/bold5000/v1/subj1/voxels.pt')\n",
+    "\n",
+    "voxels1 = torch.load('/weka/proj-medarc/shared/bold5000/formated/subj1/wholebrain.pt')\n",
+    "sessions_mask = torch.load(os.path.join('/weka/proj-medarc/shared/bold5000/formated/', f\"subj{subj}/sessions_mask.pt\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "9f412984-30f6-46d5-9d00-f8277aa0e30c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([5254, 174019]) torch.Size([5254])\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(voxels1.shape, sessions_mask.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "74070791-7980-49aa-a8d3-0c238432c62f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "unique_sessions = sessions_mask.unique()\n",
+    "\n",
+    "# Prepare a container for the output (copy first, or you can create a new one)\n",
+    "voxels1_zscore = torch.empty_like(voxels1)\n",
+    "\n",
+    "# Loop over sessions\n",
+    "for sess in unique_sessions:\n",
+    "    # Create a boolean mask for trials belonging to the current session\n",
+    "    sess_mask = (sessions_mask == sess)\n",
+    "\n",
+    "    # Extract the trials corresponding to this session\n",
+    "    sess_data = voxels1[sess_mask, :]  # shape: [num_session_trials, num_voxels]\n",
+    "\n",
+    "    # Compute mean and std along trials (dim=0 is trials, dim=1 is voxels)\n",
+    "    mean_sess = sess_data.mean(dim=0)      # shape: [num_voxels]\n",
+    "    std_sess = sess_data.std(dim=0)        # shape: [num_voxels]\n",
+    "\n",
+    "    # Avoid division by zero: if std is zero, set it to 1 (or handle as desired)\n",
+    "    std_sess[std_sess == 0] = 1.0\n",
+    "\n",
+    "    # Z-score the session data\n",
+    "    sess_data_zscore = (sess_data - mean_sess) / std_sess\n",
+    "\n",
+    "    # Place the z-scored data back into the output tensor\n",
+    "    voxels1_zscore[sess_mask, :] = sess_data_zscore"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "8a402d18-641f-45b4-aecd-cf6c7ec38447",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "--- Summary for voxelsv1 ---\n",
+      "Min: -10.2188\n",
+      "Max: 8.6066\n",
+      "Mean: -0.0000\n",
+      "Std: 1.0000\n",
+      "Number of Outliers (|value - mean| > 3*std): 30371\n",
+      "--- Summary for voxels1 ---\n",
+      "Min: -5801.5703\n",
+      "Max: 4001.2046\n",
+      "Mean: 0.0192\n",
+      "Std: 1.7900\n",
+      "Number of Outliers (|value - mean| > 3*std): 8926010\n",
+      "--- Summary for voxels1_zscore ---\n",
+      "Min: -14.8227\n",
+      "Max: 11.8500\n",
+      "Mean: 0.0000\n",
+      "Std: 0.9876\n",
+      "Number of Outliers (|value - mean| > 3*std): 3272907\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "def print_summary(tensor, name=\"Tensor\"):\n",
+    "    # Convert tensor to CPU if needed, and flatten for convenience\n",
+    "    data = tensor.detach().cpu().numpy().flatten() if torch.is_tensor(tensor) else tensor\n",
+    "\n",
+    "    # Compute basic statistics\n",
+    "    t_min = data.min()\n",
+    "    t_max = data.max()\n",
+    "    t_mean = data.mean()\n",
+    "    t_std = data.std()\n",
+    "\n",
+    "    # Simple outlier detection using a threshold, e.g., values beyond mean ± 3*std\n",
+    "    lower_bound = t_mean - 3 * t_std\n",
+    "    upper_bound = t_mean + 3 * t_std\n",
+    "    outliers = data[(data < lower_bound) | (data > upper_bound)]\n",
+    "\n",
+    "    # Print the summary\n",
+    "    print(f\"--- Summary for {name} ---\")\n",
+    "    print(f\"Min: {t_min:.4f}\")\n",
+    "    print(f\"Max: {t_max:.4f}\")\n",
+    "    print(f\"Mean: {t_mean:.4f}\")\n",
+    "    print(f\"Std: {t_std:.4f}\")\n",
+    "    print(f\"Number of Outliers (|value - mean| > 3*std): {len(outliers)}\")\n",
+    "\n",
+    "# Example usage:\n",
+    "print_summary(voxelsv1, \"voxelsv1\")\n",
+    "print_summary(voxels1, \"voxels1\")\n",
+    "print_summary(voxels1_zscore, \"voxels1_zscore\")\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "7eaebdf4-22b7-4c9a-9572-9e69870735df",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "voxels1 = voxels1_zscore"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "81084834-035f-4465-ad59-59e6b806a2f5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'train_data = {}\\ntrain_dl = {}\\nnum_voxels = {}\\nvoxels = {}\\nfor s in subj_list:\\n    print(f\"Training with {num_sessions} sessions\")\\n    if multi_subject:\\n        train_url = f\"{data_path}/wds/subj0{s}/train/\" + \"{0..\" + f\"{nsessions_allsubj[s-1]-1}\" + \"}.tar\"\\n    else:\\n        train_url = f\"{data_path}/wds/subj0{s}/train/\" + \"{0..\" + f\"{num_sessions-1}\" + \"}.tar\"\\n    print(train_url)\\n    \\n    train_data[f\\'subj0{s}\\'] = wds.WebDataset(train_url,resampled=True,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\"])\\n    train_dl[f\\'subj0{s}\\'] = torch.utils.data.DataLoader(train_data[f\\'subj0{s}\\'], batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)\\n\\n    f = h5py.File(f\\'{data_path}/betas_all_subj0{s}_fp32_renorm.hdf5\\', \\'r\\')\\n    betas = f[\\'betas\\'][:]\\n    betas = torch.Tensor(betas).to(\"cpu\").to(data_type)\\n    num_voxels_list.append(betas[0].shape[-1])\\n    num_voxels[f\\'subj0{s}\\'] = betas[0].shape[-1]\\n    voxels[f\\'subj0{s}\\'] = betas\\n    print(f\"num_voxels for subj0{s}: {num_voxels[f\\'subj0{s}\\']}\")\\n\\nprint(\"Loaded all subj train dls and betas!\\n\")\\n\\n# Validate only on one subject\\nif multi_subject: \\n    subj = subj_list[0] # cant validate on the actual held out person so picking first in subj_list\\nif not new_test: # using old test set from before full dataset released (used in original MindEye paper)\\n    if subj==3:\\n        num_test=2113\\n    elif subj==4:\\n        num_test=1985\\n    elif subj==6:\\n        num_test=2113\\n    elif subj==8:\\n        num_test=1985\\n    else:\\n        num_test=2770\\n    test_url = f\"{data_path}/wds/subj0{subj}/test/\" + \"0.tar\"\\nelif new_test: # using larger test set from after full dataset released\\n    if subj==3:\\n        num_test=2371\\n    elif subj==4:\\n        num_test=2188\\n    elif subj==6:\\n        num_test=2371\\n    elif subj==8:\\n        num_test=2188\\n    else:\\n        num_test=3000\\n    test_url = f\"{data_path}/wds/subj0{subj}/new_test/\" + \"0.tar\"\\nprint(test_url)\\ntest_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\"])\\ntest_dl = torch.utils.data.DataLoader(test_data, batch_size=num_test, shuffle=False, drop_last=True, pin_memory=True)\\nprint(f\"Loaded test dl for subj{subj}!\\n\")'"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "\"\"\"train_data = {}\n",
+    "train_dl = {}\n",
+    "num_voxels = {}\n",
+    "voxels = {}\n",
+    "for s in subj_list:\n",
+    "    print(f\"Training with {num_sessions} sessions\")\n",
+    "    if multi_subject:\n",
+    "        train_url = f\"{data_path}/wds/subj0{s}/train/\" + \"{0..\" + f\"{nsessions_allsubj[s-1]-1}\" + \"}.tar\"\n",
+    "    else:\n",
+    "        train_url = f\"{data_path}/wds/subj0{s}/train/\" + \"{0..\" + f\"{num_sessions-1}\" + \"}.tar\"\n",
+    "    print(train_url)\n",
+    "    \n",
+    "    train_data[f'subj0{s}'] = wds.WebDataset(train_url,resampled=True,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[f'subj0{s}'] = torch.utils.data.DataLoader(train_data[f'subj0{s}'], batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)\n",
+    "\n",
+    "    f = h5py.File(f'{data_path}/betas_all_subj0{s}_fp32_renorm.hdf5', 'r')\n",
+    "    betas = f['betas'][:]\n",
+    "    betas = torch.Tensor(betas).to(\"cpu\").to(data_type)\n",
+    "    num_voxels_list.append(betas[0].shape[-1])\n",
+    "    num_voxels[f'subj0{s}'] = betas[0].shape[-1]\n",
+    "    voxels[f'subj0{s}'] = betas\n",
+    "    print(f\"num_voxels for subj0{s}: {num_voxels[f'subj0{s}']}\")\n",
+    "\n",
+    "print(\"Loaded all subj train dls and betas!\\n\")\n",
+    "\n",
+    "# Validate only on one subject\n",
+    "if multi_subject: \n",
+    "    subj = subj_list[0] # cant validate on the actual held out person so picking first in subj_list\n",
+    "if not new_test: # using old test set from before full dataset released (used in original MindEye paper)\n",
+    "    if subj==3:\n",
+    "        num_test=2113\n",
+    "    elif subj==4:\n",
+    "        num_test=1985\n",
+    "    elif subj==6:\n",
+    "        num_test=2113\n",
+    "    elif subj==8:\n",
+    "        num_test=1985\n",
+    "    else:\n",
+    "        num_test=2770\n",
+    "    test_url = f\"{data_path}/wds/subj0{subj}/test/\" + \"0.tar\"\n",
+    "elif new_test: # using larger test set from after full dataset released\n",
+    "    if subj==3:\n",
+    "        num_test=2371\n",
+    "    elif subj==4:\n",
+    "        num_test=2188\n",
+    "    elif subj==6:\n",
+    "        num_test=2371\n",
+    "    elif subj==8:\n",
+    "        num_test=2188\n",
+    "    else:\n",
+    "        num_test=3000\n",
+    "    test_url = f\"{data_path}/wds/subj0{subj}/new_test/\" + \"0.tar\"\n",
+    "print(test_url)\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=num_test, shuffle=False, drop_last=True, pin_memory=True)\n",
+    "print(f\"Loaded test dl for subj{subj}!\\n\")\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "c13b4b84-094c-4b5b-bace-26c155aa6181",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'# Load 73k NSD images\\nf = h5py.File(f\\'{data_path}/coco_images_224_float16.hdf5\\', \\'r\\')\\nimages = f[\\'images\\']\\nprint(\"Loaded all 73k possible NSD images to cpu!\", images.shape)'"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "\"\"\"# Load 73k NSD images\n",
+    "f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')\n",
+    "images = f['images']\n",
+    "print(\"Loaded all 73k possible NSD images to cpu!\", 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": 14,
+   "id": "b0420dc0-199e-4c1a-857d-b1747058b467",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "clip_img_embedder = FrozenOpenCLIPImageEmbedder(\n",
+    "    arch=\"ViT-bigG-14\",\n",
+    "    version=\"laion2b_s39b_b160k\",\n",
+    "    output_tokens=True,\n",
+    "    only_tokens=True,\n",
+    ")\n",
+    "clip_img_embedder.to(device)\n",
+    "\n",
+    "clip_seq_dim = 256\n",
+    "clip_emb_dim = 1664"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5b79bd38-6990-4504-8d45-4a68d57d8885",
+   "metadata": {},
+   "source": [
+    "### SD VAE"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "01baff79-8114-482b-b115-6f05aa8ad691",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "if blurry_recon:\n",
+    "    from diffusers import AutoencoderKL    \n",
+    "    autoenc = AutoencoderKL(\n",
+    "        down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D'],\n",
+    "        up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D'],\n",
+    "        block_out_channels=[128, 256, 512, 512],\n",
+    "        layers_per_block=2,\n",
+    "        sample_size=256,\n",
+    "    )\n",
+    "    ckpt = torch.load(f'{cache_dir}/sd_image_var_autoenc.pth')\n",
+    "    autoenc.load_state_dict(ckpt)\n",
+    "    \n",
+    "    autoenc.eval()\n",
+    "    autoenc.requires_grad_(False)\n",
+    "    autoenc.to(device)\n",
+    "    utils.count_params(autoenc)\n",
+    "    \n",
+    "    from autoencoder.convnext import ConvnextXL\n",
+    "    cnx = ConvnextXL(f'{cache_dir}/convnext_xlarge_alpha0.75_fullckpt.pth')\n",
+    "    cnx.requires_grad_(False)\n",
+    "    cnx.eval()\n",
+    "    cnx.to(device)\n",
+    "    \n",
+    "    mean = torch.tensor([0.485, 0.456, 0.406]).to(device).reshape(1,3,1,1)\n",
+    "    std = torch.tensor([0.228, 0.224, 0.225]).to(device).reshape(1,3,1,1)\n",
+    "    \n",
+    "    blur_augs = AugmentationSequential(\n",
+    "        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.8),\n",
+    "        kornia.augmentation.RandomGrayscale(p=0.1),\n",
+    "        kornia.augmentation.RandomSolarize(p=0.1),\n",
+    "        kornia.augmentation.RandomResizedCrop((224,224), scale=(.9,.9), ratio=(1,1), p=1.0),\n",
+    "        data_keys=[\"input\"],\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "260e5e4a-f697-4b2c-88fc-01f6a54886c0",
+   "metadata": {},
+   "source": [
+    "### MindEye modules"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "c44c271b-173f-472e-b059-a2eda0f4c4c5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "MindEyeModule()"
+      ]
+     },
+     "execution_count": 16,
+     "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": 17,
+   "id": "038a5d61-4769-40b9-a004-f4e7b5b38bb0",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "178,196,480 total\n",
+      "178,196,480 trainable\n",
+      "param counts:\n",
+      "178,196,480 total\n",
+      "178,196,480 trainable\n",
+      "torch.Size([2, 1, 174019]) torch.Size([2, 1, 1024])\n"
+     ]
+    }
+   ],
+   "source": [
+    "num_voxels_list = [voxels1.shape[-1]]\n",
+    "class RidgeRegression(torch.nn.Module):\n",
+    "    # make sure to add weight_decay when initializing optimizer to enable regularization\n",
+    "    def __init__(self, input_sizes, out_features): \n",
+    "        super(RidgeRegression, self).__init__()\n",
+    "        self.out_features = out_features\n",
+    "        self.linears = torch.nn.ModuleList([\n",
+    "                torch.nn.Linear(input_size, out_features) for input_size in input_sizes\n",
+    "            ])\n",
+    "    def forward(self, x, subj_idx):\n",
+    "        out = self.linears[subj_idx](x[:,0]).unsqueeze(1)\n",
+    "        return out\n",
+    "        \n",
+    "model.ridge = RidgeRegression(num_voxels_list, out_features=hidden_dim)\n",
+    "utils.count_params(model.ridge)\n",
+    "utils.count_params(model)\n",
+    "\n",
+    "# test on subject 1 with fake data\n",
+    "b = torch.randn((2,1,num_voxels_list[0]))\n",
+    "print(b.shape, model.ridge(b,0).shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "7b8de65a-6d3b-4248-bea9-9b6f4d562321",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "453,360,280 total\n",
+      "453,360,280 trainable\n",
+      "param counts:\n",
+      "631,556,760 total\n",
+      "631,556,760 trainable\n",
+      "b.shape torch.Size([2, 1, 1024])\n",
+      "torch.Size([2, 256, 1664]) torch.Size([2, 256, 1664]) torch.Size([1]) torch.Size([1])\n"
+     ]
+    }
+   ],
+   "source": [
+    "from models import BrainNetwork\n",
+    "model.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, seq_len=1, n_blocks=n_blocks,\n",
+    "                          clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, \n",
+    "                          blurry_recon=blurry_recon, clip_scale=clip_scale)\n",
+    "utils.count_params(model.backbone)\n",
+    "utils.count_params(model)\n",
+    "\n",
+    "# test that the model works on some fake data\n",
+    "b = torch.randn((2,1,hidden_dim))\n",
+    "print(\"b.shape\",b.shape)\n",
+    "\n",
+    "backbone_, clip_, blur_ = model.backbone(b)\n",
+    "print(backbone_.shape, clip_.shape, blur_[0].shape, blur_[1].shape)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b397c0d7-52a3-4153-823b-c27d2eb3eeba",
+   "metadata": {},
+   "source": [
+    "### Adding diffusion prior + unCLIP if use_prior=True"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "69965344-9346-4592-9cc5-e537e31d5fce",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "259,865,216 total\n",
+      "259,865,200 trainable\n",
+      "param counts:\n",
+      "891,421,976 total\n",
+      "891,421,960 trainable\n"
+     ]
+    }
+   ],
+   "source": [
+    "if use_prior:\n",
+    "    from models import *\n",
+    "\n",
+    "    # setup diffusion prior network\n",
+    "    out_dim = clip_emb_dim\n",
+    "    depth = 6\n",
+    "    dim_head = 52\n",
+    "    heads = clip_emb_dim//52 # heads * dim_head = clip_emb_dim\n",
+    "    timesteps = 100\n",
+    "\n",
+    "    prior_network = PriorNetwork(\n",
+    "            dim=out_dim,\n",
+    "            depth=depth,\n",
+    "            dim_head=dim_head,\n",
+    "            heads=heads,\n",
+    "            causal=False,\n",
+    "            num_tokens = clip_seq_dim,\n",
+    "            learned_query_mode=\"pos_emb\"\n",
+    "        )\n",
+    "\n",
+    "    model.diffusion_prior = BrainDiffusionPrior(\n",
+    "        net=prior_network,\n",
+    "        image_embed_dim=out_dim,\n",
+    "        condition_on_text_encodings=False,\n",
+    "        timesteps=timesteps,\n",
+    "        cond_drop_prob=0.2,\n",
+    "        image_embed_scale=None,\n",
+    "    )\n",
+    "    \n",
+    "    utils.count_params(model.diffusion_prior)\n",
+    "    utils.count_params(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ec25271a-2209-400c-8026-df3b8ddc1eef",
+   "metadata": {},
+   "source": [
+    "### Setup optimizer / lr / ckpt saving"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "e14d0482-dc42-43b9-9ce1-953c32f2c9c1",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "total_steps 98400\n",
+      "\n",
+      "Done with model preparations!\n",
+      "param counts:\n",
+      "891,421,976 total\n",
+      "891,421,960 trainable\n"
+     ]
+    }
+   ],
+   "source": [
+    "no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n",
+    "\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",
+    "if use_prior:\n",
+    "    opt_grouped_parameters.extend([\n",
+    "        {'params': [p for n, p in model.diffusion_prior.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.diffusion_prior.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)\n",
+    "\n",
+    "if lr_scheduler_type == 'linear':\n",
+    "    lr_scheduler = torch.optim.lr_scheduler.LinearLR(\n",
+    "        optimizer,\n",
+    "        total_iters=int(np.floor(num_epochs*num_iterations_per_epoch)),\n",
+    "        last_epoch=-1\n",
+    "    )\n",
+    "elif lr_scheduler_type == 'cycle':\n",
+    "    total_steps=int(np.floor(num_epochs*num_iterations_per_epoch))\n",
+    "    print(\"total_steps\", total_steps)\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",
+    "    if accelerator.is_main_process:\n",
+    "        unwrapped_model = accelerator.unwrap_model(model)\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",
+    "    print(f\"\\n---saved {outdir}/{tag} ckpt!---\\n\")\n",
+    "\n",
+    "def load_ckpt(tag,load_lr=True,load_optimizer=True,load_epoch=True,strict=True,outdir=outdir,multisubj_loading=False): \n",
+    "    print(f\"\\n---loading {outdir}/{tag}.pth ckpt---\\n\")\n",
+    "    checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')\n",
+    "    state_dict = checkpoint['model_state_dict']\n",
+    "    if multisubj_loading: # remove incompatible ridge layer that will otherwise error\n",
+    "        state_dict.pop('ridge.linears.0.weight',None)\n",
+    "    model.load_state_dict(state_dict, strict=strict)\n",
+    "    if load_epoch:\n",
+    "        globals()[\"epoch\"] = checkpoint['epoch']\n",
+    "        print(\"Epoch\",epoch)\n",
+    "    if load_optimizer:\n",
+    "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
+    "    if load_lr:\n",
+    "        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n",
+    "    del checkpoint\n",
+    "\n",
+    "print(\"\\nDone with model preparations!\")\n",
+    "num_params = utils.count_params(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "983f458b-35b8-49f2-b6db-80296cece730",
+   "metadata": {},
+   "source": [
+    "# Weights and Biases"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "e671cfe9-3a30-461f-bfea-511ff34839b3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "num_test = 370"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "0a25a662-daa8-4de9-9233-8364800fcb6b",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "wandb bold5000 run testingbold5000\n",
+      "wandb_config:\n",
+      " {'model_name': 'testingbold5000', 'global_batch_size': 8, 'batch_size': 8, 'num_epochs': 150, 'num_sessions': 14, 'num_params': 891421960, 'clip_scale': 1.0, 'prior_scale': 30.0, 'blur_scale': 0.5, 'use_image_aug': False, 'max_lr': 0.0003, 'mixup_pct': 0.33, 'num_samples_per_epoch': 5254, 'num_test': 370, 'ckpt_interval': 999, 'ckpt_saving': False, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1}\n",
+      "wandb_id: testingbold5000\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"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "wandb version 0.19.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.17.1"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Run data is saved locally in <code>/weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20241207_162441-testingbold5000</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Resuming run <strong><a href='https://stability.wandb.io/ckadirt/bold5000/runs/testingbold5000' target=\"_blank\">testingbold5000</a></strong> to <a href='https://stability.wandb.io/ckadirt/bold5000' 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/bold5000' target=\"_blank\">https://stability.wandb.io/ckadirt/bold5000</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/bold5000/runs/testingbold5000' target=\"_blank\">https://stability.wandb.io/ckadirt/bold5000/runs/testingbold5000</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "if local_rank==0 and wandb_log: # only use main process for wandb logging\n",
+    "    import wandb\n",
+    "    wandb_project = 'bold5000'\n",
+    "    print(f\"wandb {wandb_project} run {model_name}\")\n",
+    "    # need to configure wandb beforehand in terminal with \"wandb init\"!\n",
+    "    wandb_config = {\n",
+    "      \"model_name\": model_name,\n",
+    "      \"global_batch_size\": global_batch_size,\n",
+    "      \"batch_size\": batch_size,\n",
+    "      \"num_epochs\": num_epochs,\n",
+    "      \"num_sessions\": num_sessions,\n",
+    "      \"num_params\": num_params,\n",
+    "      \"clip_scale\": clip_scale,\n",
+    "      \"prior_scale\": prior_scale,\n",
+    "      \"blur_scale\": blur_scale,\n",
+    "      \"use_image_aug\": use_image_aug,\n",
+    "      \"max_lr\": max_lr,\n",
+    "      \"mixup_pct\": mixup_pct,\n",
+    "      \"num_samples_per_epoch\": num_samples_per_epoch,\n",
+    "      \"num_test\": num_test,\n",
+    "      \"ckpt_interval\": ckpt_interval,\n",
+    "      \"ckpt_saving\": ckpt_saving,\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",
+    "    print(\"wandb_id:\",model_name)\n",
+    "    wandb.init(\n",
+    "        id=model_name,\n",
+    "        project=wandb_project,\n",
+    "        name=model_name,\n",
+    "        config=wandb_config,\n",
+    "        resume=\"allow\",\n",
+    "    )\n",
+    "else:\n",
+    "    wandb_log = False"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d5690151-2131-4918-b750-e869cbd1a8a8",
+   "metadata": {},
+   "source": [
+    "# Main"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "12de6387-6e18-4e4b-b5ce-a847d625330a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "epoch = 0\n",
+    "losses, test_losses, lrs = [], [], []\n",
+    "best_test_loss = 1e9\n",
+    "torch.cuda.empty_cache()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "id": "607a7c7b-fe5e-41a4-80bf-d2814b3a57cc",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# load multisubject stage1 ckpt if set\n",
+    "if multisubject_ckpt is not None:\n",
+    "    load_ckpt(\"last\",outdir=multisubject_ckpt,load_lr=False,load_optimizer=False,load_epoch=False,strict=False,multisubj_loading=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "99f09f76-4481-4133-b09a-a22b10dbc0c4",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "#train_dls = [train_dl[f'subj0{s}'] for s in subj_list]\n",
+    "images_train = images1[0:-370]\n",
+    "images_test = images1[-370::]\n",
+    "voxels_train = voxels1[0:-370]\n",
+    "voxels_test = voxels1[-370::]\n",
+    "\n",
+    "model, optimizer, images_train, images_test, voxels_train, voxels_test, lr_scheduler = accelerator.prepare(model, optimizer, images_train, images_test, voxels_train, voxels_test, lr_scheduler)\n",
+    "# leaving out test_dl since we will only have local_rank 0 device do evals"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "id": "60be0d5f-3e94-4612-9373-61b53d836393",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "testingbold5000 starting with epoch 0 / 150\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      " 55%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▌                                                                                                                                                                                                                                                                                             | 83/150 [1:58:25<1:35:35, 85.61s/it, test/blurry_pixcorr=0, test/loss=19.9, test/loss_clip_total=6.55, test/loss_prior=0.443, test/num_steps=83, test/recon_cossim=0, test/recon_mse=0, test/test_bwd_pct_correct=0.0162, test/test_fwd_pct_correct=0.0108, train/blurry_pixcorr=0, train/bwd_pct_correct=0.999, train/fwd_pct_correct=0.999, train/loss=7.47, train/loss_blurry_cont_total=0, train/loss_blurry_total=0, train/loss_clip_total=0.00171, train/loss_prior=0.249, train/lr=0.000146, train/num_steps=50630, train/recon_cossim=0.813, train/recon_mse=0.249]\n"
+     ]
+    },
+    {
+     "ename": "ValueError",
+     "evalue": "NaN loss",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[26], line 174\u001b[0m\n\u001b[1;32m    171\u001b[0m         pixcorr \u001b[38;5;241m=\u001b[39m utils\u001b[38;5;241m.\u001b[39mpixcorr(image[random_samps], blurry_recon_images)\n\u001b[1;32m    172\u001b[0m         blurry_pixcorr \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m pixcorr\u001b[38;5;241m.\u001b[39mitem()\n\u001b[0;32m--> 174\u001b[0m \u001b[43mutils\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcheck_loss\u001b[49m\u001b[43m(\u001b[49m\u001b[43mloss\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    175\u001b[0m accelerator\u001b[38;5;241m.\u001b[39mbackward(loss)\n\u001b[1;32m    176\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mstep()\n",
+      "File \u001b[0;32m/weka/proj-fmri/ckadirt/MindEyeV2/src/utils.py:209\u001b[0m, in \u001b[0;36mcheck_loss\u001b[0;34m(loss)\u001b[0m\n\u001b[1;32m    207\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcheck_loss\u001b[39m(loss):\n\u001b[1;32m    208\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m loss\u001b[38;5;241m.\u001b[39misnan()\u001b[38;5;241m.\u001b[39many():\n\u001b[0;32m--> 209\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;124mNaN loss\u001b[39m\u001b[38;5;124m'\u001b[39m)\n",
+      "\u001b[0;31mValueError\u001b[0m: NaN loss"
+     ]
+    }
+   ],
+   "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",
+    "l1 = nn.L1Loss()\n",
+    "soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))\n",
+    "\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",
+    "    recon_cossim = 0.\n",
+    "    test_recon_cossim = 0.\n",
+    "    recon_mse = 0.\n",
+    "    test_recon_mse = 0.\n",
+    "\n",
+    "    loss_clip_total = 0.\n",
+    "    loss_blurry_total = 0.\n",
+    "    loss_blurry_cont_total = 0.\n",
+    "    test_loss_clip_total = 0.\n",
+    "    \n",
+    "    loss_prior_total = 0.\n",
+    "    test_loss_prior_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",
+    "    # pre-load all batches for this epoch (it's MUCH faster to pre-load in bulk than to separate loading per batch)\n",
+    "    voxel_iters = {} # empty dict because diff subjects have differing # of voxels\n",
+    "    image_iters = torch.zeros(num_iterations_per_epoch, batch_size*len(subj_list), 3, 224, 224).float()\n",
+    "    annot_iters = {}\n",
+    "    perm_iters, betas_iters, select_iters = {}, {}, {}\n",
+    "    \"\"\"for s, train_dl in enumerate(train_dls):\n",
+    "        with torch.cuda.amp.autocast(dtype=data_type):\n",
+    "            iter = -1\n",
+    "            for behav0, past_behav0, future_behav0, old_behav0 in train_dl: \n",
+    "                # Load images to cpu from hdf5 (requires sorted indexing)\n",
+    "                image_idx = behav0[:,0,0].cpu().long().numpy()\n",
+    "                image0, image_sorted_idx = np.unique(image_idx, return_index=True)                \n",
+    "                if len(image0) != len(image_idx): # hdf5 cant handle duplicate indexing\n",
+    "                    continue\n",
+    "                iter += 1\n",
+    "                image0 = torch.tensor(images[image0], dtype=data_type)\n",
+    "                image_iters[iter,s*batch_size:s*batch_size+batch_size] = image0\n",
+    "                \n",
+    "                # Load voxels for current batch, matching above indexing\n",
+    "                voxel_idx = behav0[:,0,5].cpu().long().numpy()\n",
+    "                voxel_sorted_idx = voxel_idx[image_sorted_idx]\n",
+    "                voxel0 = voxels[f'subj0{subj_list[s]}'][voxel_sorted_idx]\n",
+    "                voxel0 = torch.Tensor(voxel0).unsqueeze(1)\n",
+    "\n",
+    "                if epoch < int(mixup_pct * num_epochs):\n",
+    "                    voxel0, perm, betas, select = utils.mixco(voxel0)\n",
+    "                    perm_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = perm\n",
+    "                    betas_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = betas\n",
+    "                    select_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = select\n",
+    "\n",
+    "                voxel_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = voxel0\n",
+    "\n",
+    "                if iter >= num_iterations_per_epoch-1:\n",
+    "                    break\"\"\"\n",
+    "\n",
+    "    # you now have voxel_iters and image_iters with num_iterations_per_epoch batches each\n",
+    "    for train_i in range(int((len(voxels1)-370)/batch_size)):\n",
+    "        with torch.cuda.amp.autocast(dtype=data_type):\n",
+    "            random_numbers = [random.randint(0, len(voxels_train)-1) for _ in range(batch_size)]\n",
+    "            optimizer.zero_grad()\n",
+    "            loss=0.\n",
+    "\n",
+    "            #voxel_list = [voxel_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "            #image = image_iters[train_i].detach()\n",
+    "\n",
+    "            image = images_train[random_numbers]#.unsqueeze(0)\n",
+    "            image = image.to(device)\n",
+    "            #print(image.shape)\n",
+    "            if use_image_aug: \n",
+    "                image = img_augment(image)\n",
+    "\n",
+    "            clip_target = clip_img_embedder(image)\n",
+    "            assert not torch.any(torch.isnan(clip_target))\n",
+    "\n",
+    "            if False: #epoch < int(mixup_pct * num_epochs):\n",
+    "                perm_list = [perm_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "                perm = torch.cat(perm_list, dim=0)\n",
+    "                betas_list = [betas_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "                betas = torch.cat(betas_list, dim=0)\n",
+    "                select_list = [select_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "                select = torch.cat(select_list, dim=0)\n",
+    "\n",
+    "            voxel = voxels_train[random_numbers].unsqueeze(1).to(device)\n",
+    "            #print(voxel.shape)\n",
+    "            #voxel_ridge_list = [model.ridge(voxel_list[si],si) for si,s in enumerate(subj_list)]\n",
+    "            #voxel_ridge = torch.cat(voxel_ridge_list, dim=0)\n",
+    "            voxel_ridge = model.ridge(voxel,0)\n",
+    "\n",
+    "            backbone, clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)\n",
+    "\n",
+    "            if clip_scale>0:\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 use_prior:\n",
+    "                loss_prior, prior_out = model.diffusion_prior(text_embed=backbone, image_embed=clip_target)\n",
+    "                loss_prior_total += loss_prior.item()\n",
+    "                loss_prior *= prior_scale\n",
+    "                loss += loss_prior\n",
+    "\n",
+    "                recon_cossim += nn.functional.cosine_similarity(prior_out, clip_target).mean().item()\n",
+    "                recon_mse += mse(prior_out, clip_target).item()\n",
+    "\n",
+    "            if clip_scale>0:\n",
+    "                if False: #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_clip_total += loss_clip.item()\n",
+    "                loss_clip *= clip_scale\n",
+    "                loss += loss_clip\n",
+    "\n",
+    "            if blurry_recon:     \n",
+    "                image_enc_pred, transformer_feats = blurry_image_enc_\n",
+    "\n",
+    "                image_enc = autoenc.encode(2*image-1).latent_dist.mode() * 0.18215\n",
+    "                loss_blurry = l1(image_enc_pred, image_enc)\n",
+    "                loss_blurry_total += loss_blurry.item()\n",
+    "\n",
+    "                if False: #epoch < int(mixup_pct * num_epochs):\n",
+    "                    image_enc_shuf = image_enc[perm]\n",
+    "                    betas_shape = [-1] + [1]*(len(image_enc.shape)-1)\n",
+    "                    image_enc[select] = image_enc[select] * betas[select].reshape(*betas_shape) + \\\n",
+    "                        image_enc_shuf[select] * (1 - betas[select]).reshape(*betas_shape)\n",
+    "\n",
+    "                image_norm = (image - mean)/std\n",
+    "                image_aug = (blur_augs(image) - mean)/std\n",
+    "                _, cnx_embeds = cnx(image_norm)\n",
+    "                _, cnx_aug_embeds = cnx(image_aug)\n",
+    "\n",
+    "                cont_loss = utils.soft_cont_loss(\n",
+    "                    nn.functional.normalize(transformer_feats.reshape(-1, transformer_feats.shape[-1]), dim=-1),\n",
+    "                    nn.functional.normalize(cnx_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),\n",
+    "                    nn.functional.normalize(cnx_aug_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),\n",
+    "                    temp=0.2)\n",
+    "                loss_blurry_cont_total += cont_loss.item()\n",
+    "\n",
+    "                loss += (loss_blurry + 0.1*cont_loss) * blur_scale #/.18215\n",
+    "\n",
+    "            if clip_scale>0:\n",
+    "                # forward and backward top 1 accuracy        \n",
+    "                labels = torch.arange(len(clip_voxels_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).item()\n",
+    "                bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()\n",
+    "\n",
+    "            if blurry_recon:\n",
+    "                with torch.no_grad():\n",
+    "                    # only doing pixcorr eval on a subset of the samples per batch because its costly & slow to compute autoenc.decode()\n",
+    "                    random_samps = np.random.choice(np.arange(len(image)), size=len(image)//5, replace=False)\n",
+    "                    blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample/ 2 + 0.5).clamp(0,1)\n",
+    "                    pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)\n",
+    "                    blurry_pixcorr += pixcorr.item()\n",
+    "\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",
+    "            if lr_scheduler_type is not None:\n",
+    "                lr_scheduler.step()\n",
+    "\n",
+    "    model.eval()\n",
+    "    if local_rank==0:\n",
+    "        with torch.no_grad(), torch.cuda.amp.autocast(dtype=data_type): \n",
+    "            for test_i in range(1):\n",
+    "                # all test samples should be loaded per batch such that test_i should never exceed 0\n",
+    "                # assert len(behav) == num_test\n",
+    "\n",
+    "                ## Average same-image repeats ##\n",
+    "                if False: #test_image is None:\n",
+    "                    voxel = voxels[f'subj0{subj}'][behav[:,0,5].cpu().long()].unsqueeze(1)\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 len(locs)==1:\n",
+    "                            locs = locs.repeat(3)\n",
+    "                        elif len(locs)==2:\n",
+    "                            locs = locs.repeat(2)[:3]\n",
+    "                        assert len(locs)==3\n",
+    "                        if test_image is None:\n",
+    "                            test_image = torch.Tensor(images[im][None])\n",
+    "                            test_voxel = voxel[locs][None]\n",
+    "                        else:\n",
+    "                            test_image = torch.vstack((test_image, torch.Tensor(images[im][None])))\n",
+    "                            test_voxel = torch.vstack((test_voxel, voxel[locs][None]))\n",
+    "\n",
+    "                loss=0.\n",
+    "                            \n",
+    "                #test_indices = torch.arange(len(test_voxel))[:300]\n",
+    "                voxel = voxels_test.to(device).unsqueeze(1)\n",
+    "                image = images_test.to(device)\n",
+    "                #assert len(image) == 300\n",
+    "\n",
+    "                clip_target = clip_img_embedder(image.float())\n",
+    "\n",
+    "                voxel_ridge = model.ridge(voxel,0) # 0th index of subj_list\n",
+    "                backbone, clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)\n",
+    "                \n",
+    "                #for rep in range(3):\n",
+    "                #    voxel_ridge = model.ridge(voxel[:,rep],0) # 0th index of subj_list\n",
+    "                #    backbone0, clip_voxels0, blurry_image_enc_ = model.backbone(voxel_ridge)\n",
+    "                #    if rep==0:\n",
+    "                #        clip_voxels = clip_voxels0\n",
+    "                #        backbone = backbone0\n",
+    "                #    else:\n",
+    "                #        clip_voxels += clip_voxels0\n",
+    "                #        backbone += backbone0\n",
+    "                #clip_voxels /= 3\n",
+    "                #backbone /= 3\n",
+    "\n",
+    "                if clip_scale>0:\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",
+    "                # for some evals, only doing a subset of the samples per batch because of computational cost\n",
+    "                random_samps = np.random.choice(np.arange(len(image)), size=len(image)//5, replace=False)\n",
+    "                \n",
+    "                if use_prior:\n",
+    "                    loss_prior, contaminated_prior_out = model.diffusion_prior(text_embed=backbone[random_samps], image_embed=clip_target[random_samps])\n",
+    "                    test_loss_prior_total += loss_prior.item()\n",
+    "                    loss_prior *= prior_scale\n",
+    "                    loss += loss_prior\n",
+    "                    \n",
+    "                if clip_scale>0:\n",
+    "                    loss_clip = utils.soft_clip_loss(\n",
+    "                        clip_voxels_norm,\n",
+    "                        clip_target_norm,\n",
+    "                        temp=.006)\n",
+    "\n",
+    "                    test_loss_clip_total += loss_clip.item()\n",
+    "                    loss_clip = loss_clip * clip_scale\n",
+    "                    loss += loss_clip\n",
+    "\n",
+    "                if blurry_recon:\n",
+    "                    image_enc_pred, _ = blurry_image_enc_\n",
+    "                    blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample / 2 + 0.5).clamp(0,1)\n",
+    "                    pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)\n",
+    "                    test_blurry_pixcorr += pixcorr.item()\n",
+    "\n",
+    "                if clip_scale>0:\n",
+    "                    # forward and backward top 1 accuracy        \n",
+    "                    labels = torch.arange(len(clip_voxels_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).item()\n",
+    "                    test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()\n",
+    "                \n",
+    "                utils.check_loss(loss)                \n",
+    "                test_losses.append(loss.item())\n",
+    "\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",
+    "                \"train/loss_blurry_cont_total\": loss_blurry_cont_total / (train_i + 1),\n",
+    "                \"test/loss_clip_total\": test_loss_clip_total / (test_i + 1),\n",
+    "                \"train/blurry_pixcorr\": blurry_pixcorr / (train_i + 1),\n",
+    "                \"test/blurry_pixcorr\": test_blurry_pixcorr / (test_i + 1),\n",
+    "                \"train/recon_cossim\": recon_cossim / (train_i + 1),\n",
+    "                \"test/recon_cossim\": test_recon_cossim / (test_i + 1),\n",
+    "                \"train/recon_mse\": recon_mse / (train_i + 1),\n",
+    "                \"test/recon_mse\": test_recon_mse / (test_i + 1),\n",
+    "                \"train/loss_prior\": loss_prior_total / (train_i + 1),\n",
+    "                \"test/loss_prior\": test_loss_prior_total / (test_i + 1),\n",
+    "                }\n",
+    "\n",
+    "            # if finished training, save jpg recons if they exist\n",
+    "            if (epoch == num_epochs-1) or (epoch % ckpt_interval == 0):\n",
+    "                if blurry_recon:    \n",
+    "                    image_enc = autoenc.encode(2*image[:4]-1).latent_dist.mode() * 0.18215\n",
+    "                    # transform blurry recon latents to images and plot it\n",
+    "                    fig, axes = plt.subplots(1, 8, figsize=(10, 4))\n",
+    "                    jj=-1\n",
+    "                    for j in [0,1,2,3]:\n",
+    "                        jj+=1\n",
+    "                        axes[jj].imshow(utils.torch_to_Image((autoenc.decode(image_enc[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))\n",
+    "                        axes[jj].axis('off')\n",
+    "                        jj+=1\n",
+    "                        axes[jj].imshow(utils.torch_to_Image((autoenc.decode(image_enc_pred[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))\n",
+    "                        axes[jj].axis('off')\n",
+    "\n",
+    "                    if wandb_log:\n",
+    "                        logs[f\"test/blur_recons\"] = wandb.Image(fig, caption=f\"epoch{epoch:03d}\")\n",
+    "                        plt.close()\n",
+    "                    else:\n",
+    "                        plt.show()\n",
+    "\n",
+    "            progress_bar.set_postfix(**logs)\n",
+    "\n",
+    "            if wandb_log: wandb.log(logs)\n",
+    "            \n",
+    "    # Save model checkpoint and reconstruct\n",
+    "    if (ckpt_saving) and (epoch % ckpt_interval == 0):\n",
+    "        save_ckpt(f'last')\n",
+    "\n",
+    "    # wait for other GPUs to catch up if needed\n",
+    "    accelerator.wait_for_everyone()\n",
+    "    torch.cuda.empty_cache()\n",
+    "\n",
+    "print(\"\\n===Finished!===\\n\")\n",
+    "if ckpt_saving:\n",
+    "    save_ckpt(f'last')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2e2fce4d-56aa-46dc-92e9-a15f80d0646b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "losses"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a7e81ae3-171f-40ad-a3e8-24bee4472325",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "plt.plot(losses)\n",
+    "plt.show()\n",
+    "plt.plot(test_losses)\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "fmri",
+   "language": "python",
+   "name": "fmri"
+  },
+  "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.11.9"
+  },
+  "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
+}

MindEyeV2/src/TrainCluster-Copy2.ipynb

@@ -0,0 +1,1910 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b0f0f4f3",
+   "metadata": {},
+   "source": [
+    "# Import packages & functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "5bad764b-45c1-45ce-a716-8d055e09821a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "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 string\n",
+    "import h5py\n",
+    "from tqdm import tqdm\n",
+    "import webdataset as wds\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torchvision import transforms\n",
+    "from accelerate import Accelerator\n",
+    "\n",
+    "# SDXL unCLIP requires code from https://github.com/Stability-AI/generative-models/tree/main\n",
+    "sys.path.append('generative_models/')\n",
+    "import sgm\n",
+    "from generative_models.sgm.modules.encoders.modules import FrozenOpenCLIPImageEmbedder # bigG embedder\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"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "id": "f9cdfdb2-a81e-495f-a777-31acd22d9746",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Top-1 Precision: 0.00%\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "import torch.nn.functional as F\n",
+    "\n",
+    "def classPrecision(logits, y_true, top=1):\n",
+    "    \"\"\"\n",
+    "    Calculate the precision of the top-n predictions.\n",
+    "    \n",
+    "    Parameters:\n",
+    "    logits (torch.Tensor): The output logits from the model (shape: [batch_size, num_classes]).\n",
+    "    y_true (torch.Tensor): The ground truth labels (shape: [batch_size]).\n",
+    "    top (int): The number of top predictions to consider.\n",
+    "    \n",
+    "    Returns:\n",
+    "    float: The precision percentage of the top-n predictions.\n",
+    "    \"\"\"\n",
+    "    # Apply softmax to get probabilities\n",
+    "    probs = F.softmax(logits, dim=1).detach().cpu()\n",
+    "    \n",
+    "    # Get the top-n predictions\n",
+    "    top_n_preds = torch.topk(probs, top, dim=1).indices.detach().cpu()\n",
+    "\n",
+    "    # Move y_true to CPU and detach\n",
+    "    y_true = y_true.detach().cpu()\n",
+    "\n",
+    "    # Check if y_true is in top-n predictions\n",
+    "    correct = top_n_preds.eq(y_true.view(-1, 1).expand_as(top_n_preds))\n",
+    "\n",
+    "    # Calculate precision\n",
+    "    precision = correct.sum().item() / y_true.size(0)\n",
+    "    \n",
+    "    return precision * 100\n",
+    "\n",
+    "# Example usage:\n",
+    "logits = torch.randn(8, 41)  # Example logits tensor\n",
+    "y_true = torch.randint(0, 41, (8,))  # Example ground truth labels\n",
+    "\n",
+    "top_n_precision = classPrecision(logits, y_true, top=1)\n",
+    "print(f\"Top-1 Precision: {top_n_precision:.2f}%\")\n"
+   ]
+  },
+  {
+   "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",
+    "data_type = torch.float16 # change depending on your mixed_precision\n",
+    "num_devices = torch.cuda.device_count()\n",
+    "if num_devices==0: num_devices = 1\n",
+    "\n",
+    "# First use \"accelerate config\" in terminal and setup using deepspeed stage 2 with CPU offloading!\n",
+    "accelerator = Accelerator(split_batches=False, mixed_precision=\"fp16\")\n",
+    "if utils.is_interactive(): # set batch size here if using interactive notebook instead of submitting job\n",
+    "    global_batch_size = batch_size = 16\n",
+    "else:\n",
+    "    global_batch_size = os.environ[\"GLOBAL_BATCH_SIZE\"]\n",
+    "    batch_size = int(os.environ[\"GLOBAL_BATCH_SIZE\"]) // num_devices"
+   ]
+  },
+  {
+   "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 = 5887\n",
+      "device: cuda\n",
+      "Distributed environment: DistributedType.NO\n",
+      "Num processes: 1\n",
+      "Process index: 0\n",
+      "Local process index: 0\n",
+      "Device: cuda\n",
+      "\n",
+      "Mixed precision type: fp16\n",
+      "\n",
+      "distributed = False num_devices = 1 local rank = 0 world size = 1 data_type = torch.float16\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(\"PID of this process =\",os.getpid())\n",
+    "device = accelerator.device\n",
+    "print(\"device:\",device)\n",
+    "world_size = accelerator.state.num_processes\n",
+    "distributed = not accelerator.state.distributed_type == 'NO'\n",
+    "num_devices = torch.cuda.device_count()\n",
+    "if num_devices==0 or not distributed: num_devices = 1\n",
+    "num_workers = num_devices\n",
+    "print(accelerator.state)\n",
+    "\n",
+    "print(\"distributed =\",distributed, \"num_devices =\", num_devices, \"local rank =\", local_rank, \"world size =\", world_size, \"data_type =\", data_type)\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": [
+      "model_name: testing-tt3\n",
+      "--data_path=/weka/proj-medarc/shared/mindeyev2_dataset                     --cache_dir=/weka/proj-medarc/shared/cache                     --model_name=testing-tt3                     --no-multi_subject --subj=1 --batch_size=16 --num_sessions=40                     --hidden_dim=1024 --clip_scale=1.                     --no-blurry_recon --blur_scale=.5                      --use_prior --prior_scale=30                     --n_blocks=4 --max_lr=1e-5 --mixup_pct=.33 --num_epochs=150 --no-use_image_aug                     --ckpt_interval=999 --no-ckpt_saving --wandb_log\n"
+     ]
+    }
+   ],
+   "source": [
+    "# if running this interactively, can specify jupyter_args here for argparser to use\n",
+    "if utils.is_interactive():\n",
+    "    model_name = \"testing-tt3\"\n",
+    "    print(\"model_name:\", model_name)\n",
+    "    \n",
+    "    # global_batch_size and batch_size should already be defined in the 2nd cell block\n",
+    "    jupyter_args = f\"--data_path=/weka/proj-medarc/shared/mindeyev2_dataset \\\n",
+    "                    --cache_dir=/weka/proj-medarc/shared/cache \\\n",
+    "                    --model_name={model_name} \\\n",
+    "                    --no-multi_subject --subj=1 --batch_size={batch_size} --num_sessions=40 \\\n",
+    "                    --hidden_dim=1024 --clip_scale=1. \\\n",
+    "                    --no-blurry_recon --blur_scale=.5  \\\n",
+    "                    --use_prior --prior_scale=30 \\\n",
+    "                    --n_blocks=4 --max_lr=1e-5 --mixup_pct=.33 --num_epochs=150 --no-use_image_aug \\\n",
+    "                    --ckpt_interval=999 --no-ckpt_saving --wandb_log\"\n",
+    "    # --multisubject_ckpt=../train_logs/multisubject_subj01_1024_24bs_nolow\n",
+    "\n",
+    "    print(jupyter_args)\n",
+    "    jupyter_args = jupyter_args.split()\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": [
+      "subj_list [1] num_sessions 40\n"
+     ]
+    }
+   ],
+   "source": [
+    "parser = argparse.ArgumentParser(description=\"Model Training Configuration\")\n",
+    "parser.add_argument(\n",
+    "    \"--model_name\", type=str, default=\"testing2\",\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=os.getcwd(),\n",
+    "    help=\"Path to where NSD data is stored / where to download it to\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--cache_dir\", type=str, default=os.getcwd(),\n",
+    "    help=\"Path to where misc. files downloaded from huggingface are stored. Defaults to current src directory.\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--subj\",type=int, default=1, choices=[1,2,3,4,5,6,7,8],\n",
+    "    help=\"Validate on which subject?\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--multisubject_ckpt\", type=str, default=None,\n",
+    "    help=\"Path to pre-trained multisubject model to finetune a single subject from. multisubject must be False.\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--num_sessions\", type=int, default=1,\n",
+    "    help=\"Number of training sessions to include\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--use_prior\",action=argparse.BooleanOptionalAction,default=True,\n",
+    "    help=\"whether to train diffusion prior (True) or just rely on retrieval part of the pipeline (False)\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--batch_size\", type=int, default=16,\n",
+    "    help=\"Batch size can be increased by 10x if only training retreival submodule and not diffusion prior\",\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",
+    "    \"--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",
+    "    \"--blurry_recon\",action=argparse.BooleanOptionalAction,default=True,\n",
+    "    help=\"whether to output blurry reconstructions\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--blur_scale\",type=float,default=.5,\n",
+    "    help=\"multiply loss from blurry recons by this number\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--clip_scale\",type=float,default=1.,\n",
+    "    help=\"multiply contrastive loss by this number\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--prior_scale\",type=float,default=30,\n",
+    "    help=\"multiply diffusion prior loss by this\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--use_image_aug\",action=argparse.BooleanOptionalAction,default=False,\n",
+    "    help=\"whether to use image augmentation\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--num_epochs\",type=int,default=150,\n",
+    "    help=\"number of epochs of training\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--multi_subject\",action=argparse.BooleanOptionalAction,default=False,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--new_test\",action=argparse.BooleanOptionalAction,default=True,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--n_blocks\",type=int,default=4,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--hidden_dim\",type=int,default=1024,\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-5,\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",
+    "# seed all random functions\n",
+    "utils.seed_everything(seed)\n",
+    "\n",
+    "outdir = os.path.abspath(f'../train_logs/{model_name}')\n",
+    "if not os.path.exists(outdir) and ckpt_saving:\n",
+    "    os.makedirs(outdir,exist_ok=True)\n",
+    "    \n",
+    "if use_image_aug or blurry_recon:\n",
+    "    import kornia\n",
+    "    from kornia.augmentation.container import AugmentationSequential\n",
+    "if use_image_aug:\n",
+    "    img_augment = AugmentationSequential(\n",
+    "        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1, p=0.3),\n",
+    "        same_on_batch=False,\n",
+    "        data_keys=[\"input\"],\n",
+    "    )\n",
+    "    \n",
+    "if multi_subject:\n",
+    "    subj_list = np.arange(1,9)\n",
+    "    subj_list = subj_list[subj_list != subj]\n",
+    "else:\n",
+    "    subj_list = [subj]\n",
+    "\n",
+    "print(\"subj_list\", subj_list, \"num_sessions\", num_sessions)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "2fb18bf3-27f5-470c-be05-002215d391b9",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "1e-05"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "max_lr"
+   ]
+  },
+  {
+   "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": [
+    "### Creating wds dataloader, preload betas and all 73k possible images"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "aefe7c27-ab39-4b2c-90f4-480f4087b7ab",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "dividing batch size by subj_list, which will then be concatenated across subj during training...\n",
+      "batch_size = 16 num_iterations_per_epoch = 1875 num_samples_per_epoch = 30000\n"
+     ]
+    }
+   ],
+   "source": [
+    "def my_split_by_node(urls): return urls\n",
+    "num_voxels_list = []\n",
+    "\n",
+    "if multi_subject:\n",
+    "    nsessions_allsubj=np.array([40, 40, 32, 30, 40, 32, 40, 30])\n",
+    "    num_samples_per_epoch = (750*40) // num_devices \n",
+    "else:\n",
+    "    num_samples_per_epoch = (750*num_sessions) // num_devices \n",
+    "\n",
+    "print(\"dividing batch size by subj_list, which will then be concatenated across subj during training...\") \n",
+    "batch_size = batch_size // len(subj_list)\n",
+    "\n",
+    "num_iterations_per_epoch = num_samples_per_epoch // (batch_size*len(subj_list))\n",
+    "\n",
+    "print(\"batch_size =\", batch_size, \"num_iterations_per_epoch =\",num_iterations_per_epoch, \"num_samples_per_epoch =\",num_samples_per_epoch)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "81084834-035f-4465-ad59-59e6b806a2f5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Training with 40 sessions\n",
+      "/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar\n",
+      "num_voxels for subj01: 15724\n",
+      "Loaded all subj train dls and betas!\n",
+      "\n",
+      "/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar\n",
+      "Loaded test dl for subj1!\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "train_data = {}\n",
+    "train_dl = {}\n",
+    "num_voxels = {}\n",
+    "voxels = {}\n",
+    "for s in subj_list:\n",
+    "    print(f\"Training with {num_sessions} sessions\")\n",
+    "    if multi_subject:\n",
+    "        train_url = f\"{data_path}/wds/subj0{s}/train/\" + \"{0..\" + f\"{nsessions_allsubj[s-1]-1}\" + \"}.tar\"\n",
+    "    else:\n",
+    "        train_url = f\"{data_path}/wds/subj0{s}/train/\" + \"{0..\" + f\"{num_sessions-1}\" + \"}.tar\"\n",
+    "    print(train_url)\n",
+    "    \n",
+    "    train_data[f'subj0{s}'] = wds.WebDataset(train_url,resampled=True,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[f'subj0{s}'] = torch.utils.data.DataLoader(train_data[f'subj0{s}'], batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)\n",
+    "\n",
+    "    f = h5py.File(f'{data_path}/betas_all_subj0{s}_fp32_renorm.hdf5', 'r')\n",
+    "    betas = f['betas'][:]\n",
+    "    betas = torch.Tensor(betas).to(\"cpu\").to(data_type)\n",
+    "    num_voxels_list.append(betas[0].shape[-1])\n",
+    "    num_voxels[f'subj0{s}'] = betas[0].shape[-1]\n",
+    "    voxels[f'subj0{s}'] = betas\n",
+    "    print(f\"num_voxels for subj0{s}: {num_voxels[f'subj0{s}']}\")\n",
+    "\n",
+    "print(\"Loaded all subj train dls and betas!\\n\")\n",
+    "\n",
+    "# Validate only on one subject\n",
+    "if multi_subject: \n",
+    "    subj = subj_list[0] # cant validate on the actual held out person so picking first in subj_list\n",
+    "if not new_test: # using old test set from before full dataset released (used in original MindEye paper)\n",
+    "    if subj==3:\n",
+    "        num_test=2113\n",
+    "    elif subj==4:\n",
+    "        num_test=1985\n",
+    "    elif subj==6:\n",
+    "        num_test=2113\n",
+    "    elif subj==8:\n",
+    "        num_test=1985\n",
+    "    else:\n",
+    "        num_test=2770\n",
+    "    test_url = f\"{data_path}/wds/subj0{subj}/test/\" + \"0.tar\"\n",
+    "elif new_test: # using larger test set from after full dataset released\n",
+    "    if subj==3:\n",
+    "        num_test=2371\n",
+    "    elif subj==4:\n",
+    "        num_test=2188\n",
+    "    elif subj==6:\n",
+    "        num_test=2371\n",
+    "    elif subj==8:\n",
+    "        num_test=2188\n",
+    "    else:\n",
+    "        num_test=3000\n",
+    "    test_url = f\"{data_path}/wds/subj0{subj}/new_test/\" + \"0.tar\"\n",
+    "print(test_url)\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=num_test, shuffle=False, drop_last=True, pin_memory=True)\n",
+    "print(f\"Loaded test dl for subj{subj}!\\n\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "c13b4b84-094c-4b5b-bace-26c155aa6181",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Loaded all 73k possible NSD images to cpu! (73000, 3, 224, 224)\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Load 73k NSD images\n",
+    "f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')\n",
+    "images = f['images']\n",
+    "print(\"Loaded all 73k possible NSD images to cpu!\", 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": 11,
+   "id": "b0420dc0-199e-4c1a-857d-b1747058b467",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "clip_img_embedder = FrozenOpenCLIPImageEmbedder(\n",
+    "    arch=\"ViT-bigG-14\",\n",
+    "    version=\"laion2b_s39b_b160k\",\n",
+    "    output_tokens=True,\n",
+    "    only_tokens=True,\n",
+    ")\n",
+    "clip_img_embedder.to(device)\n",
+    "\n",
+    "clip_seq_dim = 256\n",
+    "clip_emb_dim = 1664"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5b79bd38-6990-4504-8d45-4a68d57d8885",
+   "metadata": {},
+   "source": [
+    "### SD VAE"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "01baff79-8114-482b-b115-6f05aa8ad691",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "if blurry_recon:\n",
+    "    from diffusers import AutoencoderKL    \n",
+    "    autoenc = AutoencoderKL(\n",
+    "        down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D'],\n",
+    "        up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D'],\n",
+    "        block_out_channels=[128, 256, 512, 512],\n",
+    "        layers_per_block=2,\n",
+    "        sample_size=256,\n",
+    "    )\n",
+    "    ckpt = torch.load(f'{cache_dir}/sd_image_var_autoenc.pth')\n",
+    "    autoenc.load_state_dict(ckpt)\n",
+    "    \n",
+    "    autoenc.eval()\n",
+    "    autoenc.requires_grad_(False)\n",
+    "    autoenc.to(device)\n",
+    "    utils.count_params(autoenc)\n",
+    "    \n",
+    "    from autoencoder.convnext import ConvnextXL\n",
+    "    cnx = ConvnextXL(f'{cache_dir}/convnext_xlarge_alpha0.75_fullckpt.pth')\n",
+    "    cnx.requires_grad_(False)\n",
+    "    cnx.eval()\n",
+    "    cnx.to(device)\n",
+    "    \n",
+    "    mean = torch.tensor([0.485, 0.456, 0.406]).to(device).reshape(1,3,1,1)\n",
+    "    std = torch.tensor([0.228, 0.224, 0.225]).to(device).reshape(1,3,1,1)\n",
+    "    \n",
+    "    blur_augs = AugmentationSequential(\n",
+    "        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.8),\n",
+    "        kornia.augmentation.RandomGrayscale(p=0.1),\n",
+    "        kornia.augmentation.RandomSolarize(p=0.1),\n",
+    "        kornia.augmentation.RandomResizedCrop((224,224), scale=(.9,.9), ratio=(1,1), p=1.0),\n",
+    "        data_keys=[\"input\"],\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "260e5e4a-f697-4b2c-88fc-01f6a54886c0",
+   "metadata": {},
+   "source": [
+    "### MindEye modules"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "c44c271b-173f-472e-b059-a2eda0f4c4c5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "MindEyeModule()"
+      ]
+     },
+     "execution_count": 13,
+     "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": 14,
+   "id": "038a5d61-4769-40b9-a004-f4e7b5b38bb0",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "16,102,400 total\n",
+      "16,102,400 trainable\n",
+      "param counts:\n",
+      "16,102,400 total\n",
+      "16,102,400 trainable\n",
+      "torch.Size([2, 1, 15724]) torch.Size([2, 1, 1024])\n"
+     ]
+    }
+   ],
+   "source": [
+    "class RidgeRegression(torch.nn.Module):\n",
+    "    # make sure to add weight_decay when initializing optimizer to enable regularization\n",
+    "    def __init__(self, input_sizes, out_features): \n",
+    "        super(RidgeRegression, self).__init__()\n",
+    "        self.out_features = out_features\n",
+    "        self.linears = torch.nn.ModuleList([\n",
+    "                torch.nn.Linear(input_size, out_features) for input_size in input_sizes\n",
+    "            ])\n",
+    "    def forward(self, x, subj_idx):\n",
+    "        out = self.linears[subj_idx](x[:,0]).unsqueeze(1)\n",
+    "        return out\n",
+    "        \n",
+    "class IndividRidgeRegression(torch.nn.Module):\n",
+    "    def __init__(self, input_size, out_features):\n",
+    "        super(IndividRidgeRegression, self).__init__()\n",
+    "        self.out_features = out_features\n",
+    "        self.linear = torch.nn.Linear(input_size, out_features)\n",
+    "    def forward(self, x):\n",
+    "        out = self.linear(x)\n",
+    "        return out\n",
+    "    \n",
+    "model.ridge = RidgeRegression(num_voxels_list, out_features=hidden_dim)\n",
+    "utils.count_params(model.ridge)\n",
+    "utils.count_params(model)\n",
+    "\n",
+    "# test on subject 1 with fake data\n",
+    "b = torch.randn((2,1,num_voxels_list[0]))\n",
+    "print(b.shape, model.ridge(b,0).shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "7b8de65a-6d3b-4248-bea9-9b6f4d562321",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "453,360,280 total\n",
+      "453,360,280 trainable\n",
+      "param counts:\n",
+      "469,462,680 total\n",
+      "469,462,680 trainable\n",
+      "b.shape torch.Size([2, 1, 1024])\n",
+      "torch.Size([2, 256, 1664]) torch.Size([2, 256, 1664]) torch.Size([1]) torch.Size([1])\n"
+     ]
+    }
+   ],
+   "source": [
+    "from models import BrainNetwork\n",
+    "model.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, seq_len=1, n_blocks=n_blocks,\n",
+    "                          clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, \n",
+    "                          blurry_recon=blurry_recon, clip_scale=clip_scale)\n",
+    "utils.count_params(model.backbone)\n",
+    "utils.count_params(model)\n",
+    "\n",
+    "# test that the model works on some fake data\n",
+    "b = torch.randn((2,1,hidden_dim))\n",
+    "print(\"b.shape\",b.shape)\n",
+    "\n",
+    "backbone_, clip_, blur_ = model.backbone(b)\n",
+    "print(backbone_.shape, clip_.shape, blur_[0].shape, blur_[1].shape)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e63b26a9-3a36-4638-be3c-dbf28705dd76",
+   "metadata": {},
+   "source": [
+    "### Load semantic clusters"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "2950add9-3828-4d55-bbf7-0d9f3519de5c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "semantic_cluster_onehot.shape torch.Size([73000, 41])\n",
+      "num_seman_clusters 41\n"
+     ]
+    }
+   ],
+   "source": [
+    "path_semantic_names = \"/weka/proj-medarc/shared/mindeyev2_dataset/semantic_cluster_names.npy\"\n",
+    "path_semantic_cluster = \"/weka/proj-fmri/ckadirt/MindEyeV2/src/COCO_73k_semantic_cluster.npy\"\n",
+    "semantic_cluster_names = np.load(path_semantic_names)\n",
+    "semantic_cluster = np.load(path_semantic_cluster)\n",
+    "possible_semantic_clusters = np.unique(semantic_cluster)\n",
+    "\n",
+    "# one-hot encode semantic clusters\n",
+    "# move possible_semantic_clusters to numbers and create a dictionary\n",
+    "semantic_cluster_dict = {cluster: i for i, cluster in enumerate(possible_semantic_clusters)}\n",
+    "semantic_cluster_onehot = torch.zeros((len(semantic_cluster), len(possible_semantic_clusters)))\n",
+    "for i, cluster in enumerate(semantic_cluster):\n",
+    "    semantic_cluster_onehot[i, semantic_cluster_dict[cluster]] = 1\n",
+    "\n",
+    "\n",
+    "print(\"semantic_cluster_onehot.shape\", semantic_cluster_onehot.shape)\n",
+    "\n",
+    "num_seman_clusters = len(np.unique(semantic_cluster))\n",
+    "print(\"num_seman_clusters\", num_seman_clusters)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b397c0d7-52a3-4153-823b-c27d2eb3eeba",
+   "metadata": {},
+   "source": [
+    "### Adding the ridge regression to the class"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "69965344-9346-4592-9cc5-e537e31d5fce",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "17,465,385 total\n",
+      "17,465,385 trainable\n",
+      "param counts:\n",
+      "486,928,065 total\n",
+      "486,928,065 trainable\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "486928065"
+      ]
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# if use_prior:\n",
+    "#     from models import *\n",
+    "\n",
+    "#     # setup diffusion prior network\n",
+    "#     out_dim = clip_emb_dim\n",
+    "#     depth = 6\n",
+    "#     dim_head = 52\n",
+    "#     heads = clip_emb_dim//52 # heads * dim_head = clip_emb_dim\n",
+    "#     timesteps = 100\n",
+    "\n",
+    "#     prior_network = PriorNetwork(\n",
+    "#             dim=out_dim,\n",
+    "#             depth=depth,\n",
+    "#             dim_head=dim_head,\n",
+    "#             heads=heads,\n",
+    "#             causal=False,\n",
+    "#             num_tokens = clip_seq_dim,\n",
+    "#             learned_query_mode=\"pos_emb\"\n",
+    "#         )\n",
+    "\n",
+    "#     model.diffusion_prior = BrainDiffusionPrior(\n",
+    "#         net=prior_network,\n",
+    "#         image_embed_dim=out_dim,\n",
+    "#         condition_on_text_encodings=False,\n",
+    "#         timesteps=timesteps,\n",
+    "#         cond_drop_prob=0.2,\n",
+    "#         image_embed_scale=None,\n",
+    "#     )\n",
+    "    \n",
+    "#     utils.count_params(model.diffusion_prior)\n",
+    "#     utils.count_params(model)\n",
+    "\n",
+    "model.RRClassifier = IndividRidgeRegression(clip_emb_dim*clip_seq_dim, out_features=num_seman_clusters)\n",
+    "utils.count_params(model.RRClassifier)\n",
+    "utils.count_params(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ec25271a-2209-400c-8026-df3b8ddc1eef",
+   "metadata": {},
+   "source": [
+    "### Setup optimizer / lr / ckpt saving"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "e14d0482-dc42-43b9-9ce1-953c32f2c9c1",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "total_steps 281250\n",
+      "\n",
+      "Done with model preparations!\n",
+      "param counts:\n",
+      "486,928,065 total\n",
+      "486,928,065 trainable\n"
+     ]
+    }
+   ],
+   "source": [
+    "no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n",
+    "\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",
+    "    {'params': [p for n, p in model.RRClassifier.named_parameters()], 'weight_decay': 1e-2},\n",
+    "]\n",
+    "# if use_prior:\n",
+    "#     opt_grouped_parameters.extend([\n",
+    "#         {'params': [p for n, p in model.diffusion_prior.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.diffusion_prior.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}\n",
+    "#     ])\n",
+    "# opt_grouped_parameters.extend([\n",
+    "#     \n",
+    "# ])\n",
+    "\n",
+    "optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr)\n",
+    "\n",
+    "if lr_scheduler_type == 'linear':\n",
+    "    lr_scheduler = torch.optim.lr_scheduler.LinearLR(\n",
+    "        optimizer,\n",
+    "        total_iters=int(np.floor(num_epochs*num_iterations_per_epoch)),\n",
+    "        last_epoch=-1\n",
+    "    )\n",
+    "elif lr_scheduler_type == 'cycle':\n",
+    "    total_steps=int(np.floor(num_epochs*num_iterations_per_epoch))\n",
+    "    print(\"total_steps\", total_steps)\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",
+    "    if accelerator.is_main_process:\n",
+    "        unwrapped_model = accelerator.unwrap_model(model)\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",
+    "    print(f\"\\n---saved {outdir}/{tag} ckpt!---\\n\")\n",
+    "\n",
+    "def load_ckpt(tag,load_lr=True,load_optimizer=True,load_epoch=True,strict=True,outdir=outdir,multisubj_loading=False): \n",
+    "    print(f\"\\n---loading {outdir}/{tag}.pth ckpt---\\n\")\n",
+    "    checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')\n",
+    "    state_dict = checkpoint['model_state_dict']\n",
+    "    if multisubj_loading: # remove incompatible ridge layer that will otherwise error\n",
+    "        state_dict.pop('ridge.linears.0.weight',None)\n",
+    "    model.load_state_dict(state_dict, strict=strict)\n",
+    "    if load_epoch:\n",
+    "        globals()[\"epoch\"] = checkpoint['epoch']\n",
+    "        print(\"Epoch\",epoch)\n",
+    "    if load_optimizer:\n",
+    "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
+    "    if load_lr:\n",
+    "        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n",
+    "    del checkpoint\n",
+    "\n",
+    "print(\"\\nDone with model preparations!\")\n",
+    "num_params = utils.count_params(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "983f458b-35b8-49f2-b6db-80296cece730",
+   "metadata": {},
+   "source": [
+    "# Weights and Biases"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "0a25a662-daa8-4de9-9233-8364800fcb6b",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "wandb mindeye_semantic_cluster run testing-tt3\n",
+      "wandb_config:\n",
+      " {'model_name': 'testing-tt3', 'global_batch_size': 16, 'batch_size': 16, 'num_epochs': 150, 'num_sessions': 40, 'num_params': 486928065, 'clip_scale': 1.0, 'prior_scale': 30.0, 'blur_scale': 0.5, 'use_image_aug': False, 'max_lr': 1e-05, 'mixup_pct': 0.33, 'num_samples_per_epoch': 30000, 'num_test': 3000, 'ckpt_interval': 999, 'ckpt_saving': False, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1, 'train_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar', 'test_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar'}\n",
+      "wandb_id: testing-tt3\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"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "wandb version 0.17.4 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.17.1"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Run data is saved locally in <code>/weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240709_011732-testing-tt3</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/mindeye_semantic_cluster/runs/testing-tt3' target=\"_blank\">testing-tt3</a></strong> to <a href='https://stability.wandb.io/ckadirt/mindeye_semantic_cluster' 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/mindeye_semantic_cluster' target=\"_blank\">https://stability.wandb.io/ckadirt/mindeye_semantic_cluster</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/mindeye_semantic_cluster/runs/testing-tt3' target=\"_blank\">https://stability.wandb.io/ckadirt/mindeye_semantic_cluster/runs/testing-tt3</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "if local_rank==0 and wandb_log: # only use main process for wandb logging\n",
+    "    import wandb\n",
+    "    wandb_project = 'mindeye_semantic_cluster'\n",
+    "    print(f\"wandb {wandb_project} run {model_name}\")\n",
+    "    # need to configure wandb beforehand in terminal with \"wandb init\"!\n",
+    "    wandb_config = {\n",
+    "      \"model_name\": model_name,\n",
+    "      \"global_batch_size\": global_batch_size,\n",
+    "      \"batch_size\": batch_size,\n",
+    "      \"num_epochs\": num_epochs,\n",
+    "      \"num_sessions\": num_sessions,\n",
+    "      \"num_params\": num_params,\n",
+    "      \"clip_scale\": clip_scale,\n",
+    "      \"prior_scale\": prior_scale,\n",
+    "      \"blur_scale\": blur_scale,\n",
+    "      \"use_image_aug\": use_image_aug,\n",
+    "      \"max_lr\": max_lr,\n",
+    "      \"mixup_pct\": mixup_pct,\n",
+    "      \"num_samples_per_epoch\": num_samples_per_epoch,\n",
+    "      \"num_test\": num_test,\n",
+    "      \"ckpt_interval\": ckpt_interval,\n",
+    "      \"ckpt_saving\": ckpt_saving,\n",
+    "      \"seed\": seed,\n",
+    "      \"distributed\": distributed,\n",
+    "      \"num_devices\": num_devices,\n",
+    "      \"world_size\": world_size,\n",
+    "      \"train_url\": train_url,\n",
+    "      \"test_url\": test_url,\n",
+    "    }\n",
+    "    print(\"wandb_config:\\n\",wandb_config)\n",
+    "    print(\"wandb_id:\",model_name)\n",
+    "    wandb.login(host='https://stability.wandb.io')\n",
+    "    wandb.init(\n",
+    "        id=model_name,\n",
+    "        project=wandb_project,\n",
+    "        name=model_name,\n",
+    "        config=wandb_config,\n",
+    "        resume=\"allow\",\n",
+    "    )\n",
+    "else:\n",
+    "    wandb_log = False"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d5690151-2131-4918-b750-e869cbd1a8a8",
+   "metadata": {},
+   "source": [
+    "# Main"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "12de6387-6e18-4e4b-b5ce-a847d625330a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "epoch = 0\n",
+    "losses, test_losses, lrs = [], [], []\n",
+    "best_test_loss = 1e9\n",
+    "torch.cuda.empty_cache()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "607a7c7b-fe5e-41a4-80bf-d2814b3a57cc",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# load multisubject stage1 ckpt if set\n",
+    "if multisubject_ckpt is not None:\n",
+    "    load_ckpt(\"last\",outdir=multisubject_ckpt,load_lr=False,load_optimizer=False,load_epoch=False,strict=False,multisubj_loading=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "99f09f76-4481-4133-b09a-a22b10dbc0c4",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "train_dls = [train_dl[f'subj0{s}'] for s in subj_list]\n",
+    "\n",
+    "model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot = accelerator.prepare(model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot)\n",
+    "# leaving out test_dl since we will only have local_rank 0 device do evals"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "e16fa583-064c-4cfa-9f47-08d08f53f504",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1875\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(num_iterations_per_epoch)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "id": "60be0d5f-3e94-4612-9373-61b53d836393",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "testing-tt3 starting with epoch 4 / 150\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   | 0/146 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([16, 256, 1664])\n",
+      "torch.Size([16, 41])\n",
+      "['photo of tennis' 'photo of bedroom' 'photo of surfer'\n",
+      " 'photo of group of people' 'photo of cat' 'photo of vehicle'\n",
+      " 'photo of baseball' 'photo of bird' 'photo of train'\n",
+      " 'photo of group of people' 'photo of umbrella' 'photo of food'\n",
+      " 'photo of giraffe' 'photo of clocktower' 'photo of food'\n",
+      " 'photo of sports']\n",
+      "tensor([35.,  5., 33., 20., 10., 39.,  2.,  7., 37., 20., 38., 17., 19., 11.,\n",
+      "        17., 31.])\n",
+      "torch.Size([16, 256, 1664])\n",
+      "torch.Size([16, 41])\n",
+      "['photo of food' 'photo of clocktower' 'photo of bedroom'\n",
+      " 'photo of baseball' 'photo of group of people' 'photo of umbrella'\n",
+      " 'photo of surfer' 'photo of food' 'photo of train' 'photo of bird'\n",
+      " 'photo of tennis' 'photo of sports' 'photo of group of people'\n",
+      " 'photo of cat' 'photo of vehicle' 'photo of giraffe']\n",
+      "tensor([17., 11.,  5.,  2., 20., 38., 33., 17., 37.,  7., 35., 31., 20., 10.,\n",
+      "        39., 19.])\n",
+      "torch.Size([16, 256, 1664])\n",
+      "torch.Size([16, 41])\n",
+      "['photo of bedroom' 'photo of tennis' 'photo of clocktower'\n",
+      " 'photo of umbrella' 'photo of group of people' 'photo of group of people'\n",
+      " 'photo of surfer' 'photo of food' 'photo of giraffe' 'photo of cat'\n",
+      " 'photo of bird' 'photo of food' 'photo of train' 'photo of sports'\n",
+      " 'photo of vehicle' 'photo of baseball']\n",
+      "tensor([ 5., 35., 11., 38., 20., 20., 33., 17., 19., 10.,  7., 17., 37., 31.,\n",
+      "        39.,  2.])\n",
+      "torch.Size([16, 256, 1664])\n",
+      "torch.Size([16, 41])\n",
+      "['photo of surfer' 'photo of food' 'photo of umbrella' 'photo of tennis'\n",
+      " 'photo of bird' 'photo of train' 'photo of group of people'\n",
+      " 'photo of group of people' 'photo of clocktower' 'photo of sports'\n",
+      " 'photo of cat' 'photo of vehicle' 'photo of giraffe' 'photo of bedroom'\n",
+      " 'photo of food' 'photo of baseball']\n",
+      "tensor([33., 17., 38., 35.,  7., 37., 20., 20., 11., 31., 10., 39., 19.,  5.,\n",
+      "        17.,  2.])\n",
+      "torch.Size([16, 256, 1664])\n",
+      "torch.Size([16, 41])\n",
+      "['photo of food' 'photo of clocktower' 'photo of bedroom'\n",
+      " 'photo of baseball' 'photo of tennis' 'photo of giraffe'\n",
+      " 'photo of vehicle' 'photo of group of people' 'photo of group of people'\n",
+      " 'photo of cat' 'photo of food' 'photo of umbrella' 'photo of sports'\n",
+      " 'photo of train' 'photo of bird' 'photo of surfer']\n",
+      "tensor([17., 11.,  5.,  2., 35., 19., 39., 20., 20., 10., 17., 38., 31., 37.,\n",
+      "         7., 33.])\n",
+      "torch.Size([16, 256, 1664])\n",
+      "torch.Size([16, 41])\n",
+      "['photo of giraffe' 'photo of food' 'photo of baseball'\n",
+      " 'photo of clocktower' 'photo of cat' 'photo of vehicle'\n",
+      " 'photo of bedroom' 'photo of sports' 'photo of surfer'\n",
+      " 'photo of group of people' 'photo of tennis' 'photo of food'\n",
+      " 'photo of umbrella' 'photo of bird' 'photo of group of people'\n",
+      " 'photo of train']\n",
+      "tensor([19., 17.,  2., 11., 10., 39.,  5., 31., 33., 20., 35., 17., 38.,  7.,\n",
+      "        20., 37.])\n",
+      "torch.Size([16, 256, 1664])\n",
+      "torch.Size([16, 41])\n",
+      "['photo of group of people' 'photo of baseball' 'photo of group of people'\n",
+      " 'photo of bird' 'photo of train' 'photo of food' 'photo of vehicle'\n",
+      " 'photo of cat' 'photo of sports' 'photo of giraffe' 'photo of surfer'\n",
+      " 'photo of umbrella' 'photo of food' 'photo of tennis'\n",
+      " 'photo of clocktower' 'photo of bedroom']\n",
+      "tensor([20.,  2., 20.,  7., 37., 17., 39., 10., 31., 19., 33., 38., 17., 35.,\n",
+      "        11.,  5.])\n",
+      "torch.Size([16, 256, 1664])\n",
+      "torch.Size([16, 41])\n",
+      "['photo of clocktower' 'photo of group of people' 'photo of bedroom'\n",
+      " 'photo of giraffe' 'photo of surfer' 'photo of vehicle' 'photo of cat'\n",
+      " 'photo of tennis' 'photo of food' 'photo of group of people'\n",
+      " 'photo of train' 'photo of bird' 'photo of sports' 'photo of baseball'\n",
+      " 'photo of umbrella' 'photo of food']\n",
+      "tensor([11., 20.,  5., 19., 33., 39., 10., 35., 17., 20., 37.,  7., 31.,  2.,\n",
+      "        38., 17.])\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   | 0/146 [00:34<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([16, 256, 1664])\n",
+      "torch.Size([16, 41])\n",
+      "['photo of vehicle' 'photo of surfer' 'photo of tennis'\n",
+      " 'photo of group of people' 'photo of food' 'photo of group of people'\n",
+      " 'photo of sports' 'photo of food' 'photo of giraffe' 'photo of baseball'\n",
+      " 'photo of umbrella' 'photo of cat' 'photo of bedroom'\n",
+      " 'photo of clocktower' 'photo of train' 'photo of bird']\n",
+      "tensor([39., 33., 35., 20., 17., 20., 31., 17., 19.,  2., 38., 10.,  5., 11.,\n",
+      "        37.,  7.])\n",
+      "torch.Size([16, 256, 1664])\n",
+      "torch.Size([16, 41])\n",
+      "['photo of cat' 'photo of baseball' 'photo of vehicle' 'photo of sports'\n",
+      " 'photo of bird' 'photo of surfer' 'photo of food' 'photo of food'\n",
+      " 'photo of clocktower' 'photo of umbrella' 'photo of group of people'\n",
+      " 'photo of bedroom' 'photo of train' 'photo of tennis' 'photo of giraffe'\n",
+      " 'photo of group of people']\n",
+      "tensor([10.,  2., 39., 31.,  7., 33., 17., 17., 11., 38., 20.,  5., 37., 35.,\n",
+      "        19., 20.])\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    },
+    {
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[28], line 90\u001b[0m\n\u001b[1;32m     88\u001b[0m voxel_list \u001b[38;5;241m=\u001b[39m [voxel_iters[\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msubj0\u001b[39m\u001b[38;5;132;01m{\u001b[39;00ms\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m_iter\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtrain_i\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mdetach()\u001b[38;5;241m.\u001b[39mto(device) \u001b[38;5;28;01mfor\u001b[39;00m s \u001b[38;5;129;01min\u001b[39;00m subj_list]\n\u001b[1;32m     89\u001b[0m image \u001b[38;5;241m=\u001b[39m image_iters[train_i]\u001b[38;5;241m.\u001b[39mdetach()\n\u001b[0;32m---> 90\u001b[0m image \u001b[38;5;241m=\u001b[39m \u001b[43mimage\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     92\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m use_image_aug: \n\u001b[1;32m     93\u001b[0m     image \u001b[38;5;241m=\u001b[39m img_augment(image)\n",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
+     ]
+    }
+   ],
+   "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",
+    "l1 = nn.L1Loss()\n",
+    "soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))\n",
+    "\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",
+    "    recon_cossim = 0.\n",
+    "    test_recon_cossim = 0.\n",
+    "    recon_mse = 0.\n",
+    "    test_recon_mse = 0.\n",
+    "\n",
+    "    loss_clip_total = 0.\n",
+    "    loss_blurry_total = 0.\n",
+    "    loss_blurry_cont_total = 0.\n",
+    "    test_loss_clip_total = 0.\n",
+    "    \n",
+    "    loss_prior_total = 0.\n",
+    "    test_loss_prior_total = 0.\n",
+    "    \n",
+    "    loss_RR_total = 0.\n",
+    "    test_loss_RR_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",
+    "    class_precisions_1 = 0\n",
+    "    test_class_precisions_1 = 0\n",
+    "\n",
+    "    class_precisions_5 = 0\n",
+    "    test_class_precisions_5 = 0\n",
+    "\n",
+    "    class_precisions_10 = 0\n",
+    "    test_class_precisions_10 = 0\n",
+    "\n",
+    "    # pre-load all batches for this epoch (it's MUCH faster to pre-load in bulk than to separate loading per batch)\n",
+    "    voxel_iters = {} # empty dict because diff subjects have differing # of voxels\n",
+    "    image_iters = torch.zeros(num_iterations_per_epoch, batch_size*len(subj_list), 3, 224, 224).float()\n",
+    "    annot_iters = {}\n",
+    "    perm_iters, betas_iters, select_iters = {}, {}, {}\n",
+    "    images_indexes = {}\n",
+    "    for s, train_dl in enumerate(train_dls):\n",
+    "        with torch.cuda.amp.autocast(dtype=data_type):\n",
+    "            iter = -1\n",
+    "            for behav0, past_behav0, future_behav0, old_behav0 in train_dl: \n",
+    "                # Load images to cpu from hdf5 (requires sorted indexing)\n",
+    "                image_idx = behav0[:,0,0].cpu().long().numpy()\n",
+    "\n",
+    "                image0, image_sorted_idx = np.unique(image_idx, return_index=True)                \n",
+    "                if len(image0) != len(image_idx): # hdf5 cant handle duplicate indexing\n",
+    "                    continue\n",
+    "                iter += 1\n",
+    "                image0 = torch.tensor(images[image0], dtype=data_type)\n",
+    "                image_iters[iter,s*batch_size:s*batch_size+batch_size] = image0\n",
+    "                images_indexes[f\"subj0{s}_iter{iter}\"] = image_sorted_idx\n",
+    "                \n",
+    "                # Load voxels for current batch, matching above indexing\n",
+    "                voxel_idx = behav0[:,0,5].cpu().long().numpy()\n",
+    "                voxel_sorted_idx = voxel_idx[image_sorted_idx]\n",
+    "                voxel0 = voxels[f'subj0{subj_list[s]}'][voxel_sorted_idx]\n",
+    "                voxel0 = torch.Tensor(voxel0).unsqueeze(1)\n",
+    "\n",
+    "                if epoch < int(mixup_pct * num_epochs):\n",
+    "                    voxel0, perm, betas, select = utils.mixco(voxel0)\n",
+    "                    perm_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = perm\n",
+    "                    betas_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = betas\n",
+    "                    select_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = select\n",
+    "\n",
+    "                voxel_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = voxel0\n",
+    "\n",
+    "                if iter >= num_iterations_per_epoch-1:\n",
+    "                    break\n",
+    "\n",
+    "    # you now have voxel_iters and image_iters with num_iterations_per_epoch batches each\n",
+    "    for train_i in range(num_iterations_per_epoch):\n",
+    "        with torch.cuda.amp.autocast(dtype=data_type):\n",
+    "            optimizer.zero_grad()\n",
+    "            loss=0.\n",
+    "\n",
+    "            voxel_list = [voxel_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "            image = image_iters[train_i].detach()\n",
+    "            image = image.to(device)\n",
+    "\n",
+    "            if use_image_aug: \n",
+    "                image = img_augment(image)\n",
+    "\n",
+    "            clip_target = clip_img_embedder(image)\n",
+    "            assert not torch.any(torch.isnan(clip_target))\n",
+    "\n",
+    "            if epoch < int(mixup_pct * num_epochs):\n",
+    "                perm_list = [perm_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "                perm = torch.cat(perm_list, dim=0)\n",
+    "                betas_list = [betas_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "                betas = torch.cat(betas_list, dim=0)\n",
+    "                select_list = [select_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "                select = torch.cat(select_list, dim=0)\n",
+    "\n",
+    "            voxel_ridge_list = [model.ridge(voxel_list[si],si) for si,s in enumerate(subj_list)]\n",
+    "            voxel_ridge = torch.cat(voxel_ridge_list, dim=0)\n",
+    "\n",
+    "            backbone, clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)\n",
+    "\n",
+    "            if clip_scale>0:\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 use_prior:\n",
+    "            #     loss_prior, prior_out = model.diffusion_prior(text_embed=backbone, image_embed=clip_target)\n",
+    "            #     loss_prior_total += loss_prior.item()\n",
+    "            #     loss_prior *= prior_scale\n",
+    "            #     loss += loss_prior\n",
+    "\n",
+    "            #     recon_cossim += nn.functional.cosine_similarity(prior_out, clip_target).mean().item()\n",
+    "            #     recon_mse += mse(prior_out, clip_target).item()\n",
+    "            print(backbone.shape)\n",
+    "            logits = model.RRClassifier(backbone.flatten(1))\n",
+    "            print(logits.shape)\n",
+    "            print(semantic_cluster[images_indexes[f\"subj0{s}_iter{train_i}\"]])\n",
+    "            print(torch.Tensor([semantic_cluster_dict[i] for i in semantic_cluster[images_indexes[f\"subj0{s}_iter{train_i}\"]]]))\n",
+    "            #print(logits.shape, torch.argmax(semantic_cluster_onehot[images_indexes[f\"subj0{s}_iter{train_i}\"]], dim=1).shape)\n",
+    "            #print(logits, torch.argmax(semantic_cluster_onehot[images_indexes[f\"subj0{s}_iter{train_i}\"]], dim=1))\n",
+    "            loss_RR = nn.functional.cross_entropy(logits, torch.argmax(semantic_cluster_onehot[images_indexes[f\"subj0{s}_iter{train_i}\"]], dim=1).to(logits.device))\n",
+    "            #print(\"backbone.shape\",backbone.shape, \"clip_voxels.shape\",clip_voxels.shape, \"blurry_image_enc_[0].shape\",blurry_image_enc_[0].shape, \"blurry_image_enc_[1].shape\",blurry_image_enc_[1].shape)\n",
+    "            #something \n",
+    "\n",
+    "            loss_RR_total += loss_RR.item()\n",
+    "            loss += loss_RR\n",
+    "\n",
+    "            if (torch.rand(1) < 0.03).item():\n",
+    "                print(\"loss_RR\", loss_RR.item())\n",
+    "\n",
+    "            if clip_scale>0:\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_clip_total += loss_clip.item()\n",
+    "                loss_clip *= clip_scale\n",
+    "                # loss += loss_clip\n",
+    "\n",
+    "            if blurry_recon:     \n",
+    "                image_enc_pred, transformer_feats = blurry_image_enc_\n",
+    "\n",
+    "                image_enc = autoenc.encode(2*image-1).latent_dist.mode() * 0.18215\n",
+    "                loss_blurry = l1(image_enc_pred, image_enc)\n",
+    "                loss_blurry_total += loss_blurry.item()\n",
+    "\n",
+    "                if epoch < int(mixup_pct * num_epochs):\n",
+    "                    image_enc_shuf = image_enc[perm]\n",
+    "                    betas_shape = [-1] + [1]*(len(image_enc.shape)-1)\n",
+    "                    image_enc[select] = image_enc[select] * betas[select].reshape(*betas_shape) + \\\n",
+    "                        image_enc_shuf[select] * (1 - betas[select]).reshape(*betas_shape)\n",
+    "\n",
+    "                image_norm = (image - mean)/std\n",
+    "                image_aug = (blur_augs(image) - mean)/std\n",
+    "                _, cnx_embeds = cnx(image_norm)\n",
+    "                _, cnx_aug_embeds = cnx(image_aug)\n",
+    "\n",
+    "                cont_loss = utils.soft_cont_loss(\n",
+    "                    nn.functional.normalize(transformer_feats.reshape(-1, transformer_feats.shape[-1]), dim=-1),\n",
+    "                    nn.functional.normalize(cnx_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),\n",
+    "                    nn.functional.normalize(cnx_aug_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),\n",
+    "                    temp=0.2)\n",
+    "                loss_blurry_cont_total += cont_loss.item()\n",
+    "\n",
+    "                # loss += (loss_blurry + 0.1*cont_loss) * blur_scale #/.18215\n",
+    "\n",
+    "            if clip_scale>0:\n",
+    "                # forward and backward top 1 accuracy        \n",
+    "                labels = torch.arange(len(clip_voxels_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).item()\n",
+    "                bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()\n",
+    "\n",
+    "            if blurry_recon:\n",
+    "                with torch.no_grad():\n",
+    "                    # only doing pixcorr eval on a subset of the samples per batch because its costly & slow to compute autoenc.decode()\n",
+    "                    random_samps = np.random.choice(np.arange(len(image)), size=len(image)//5, replace=False)\n",
+    "                    blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample/ 2 + 0.5).clamp(0,1)\n",
+    "                    pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)\n",
+    "                    blurry_pixcorr += pixcorr.item()\n",
+    "\n",
+    "            class_precisions_1 += classPrecision(logits, torch.argmax(semantic_cluster_onehot[images_indexes[f\"subj0{s}_iter{train_i}\"]], dim=1).to(logits.device))\n",
+    "            class_precisions_5 += classPrecision(logits, torch.argmax(semantic_cluster_onehot[images_indexes[f\"subj0{s}_iter{train_i}\"]], dim=1).to(logits.device), 5)\n",
+    "            class_precisions_10 += classPrecision(logits, torch.argmax(semantic_cluster_onehot[images_indexes[f\"subj0{s}_iter{train_i}\"]], dim=1).to(logits.device), 10)\n",
+    "\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",
+    "            if lr_scheduler_type is not None:\n",
+    "                lr_scheduler.step()\n",
+    "\n",
+    "    model.eval()\n",
+    "    if local_rank==0:\n",
+    "        with torch.no_grad(), torch.cuda.amp.autocast(dtype=data_type): \n",
+    "            for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):  \n",
+    "                # all test samples should be loaded per batch such that test_i should never exceed 0\n",
+    "                assert len(behav) == num_test\n",
+    "\n",
+    "                ## Average same-image repeats ##\n",
+    "                if test_image is None:\n",
+    "                    voxel = voxels[f'subj0{subj}'][behav[:,0,5].cpu().long()].unsqueeze(1)\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 len(locs)==1:\n",
+    "                            locs = locs.repeat(3)\n",
+    "                        elif len(locs)==2:\n",
+    "                            locs = locs.repeat(2)[:3]\n",
+    "                        assert len(locs)==3\n",
+    "                        if test_image is None:\n",
+    "                            test_image = torch.Tensor(images[im][None])\n",
+    "                            test_voxel = voxel[locs][None]\n",
+    "                        else:\n",
+    "                            test_image = torch.vstack((test_image, torch.Tensor(images[im][None])))\n",
+    "                            test_voxel = torch.vstack((test_voxel, voxel[locs][None]))\n",
+    "\n",
+    "                loss=0.\n",
+    "                            \n",
+    "                test_indices = torch.arange(len(test_voxel))[:300]\n",
+    "                voxel = test_voxel[test_indices].to(device)\n",
+    "                image = test_image[test_indices].to(device)\n",
+    "                assert len(image) == 300\n",
+    "\n",
+    "                clip_target = clip_img_embedder(image.float())\n",
+    "\n",
+    "                for rep in range(3):\n",
+    "                    voxel_ridge = model.ridge(voxel[:,rep],0) # 0th index of subj_list\n",
+    "                    backbone0, clip_voxels0, blurry_image_enc_ = model.backbone(voxel_ridge)\n",
+    "\n",
+    "                    logits0 = model.RRClassifier(backbone0.flatten(1))\n",
+    "\n",
+    "                    if rep==0:\n",
+    "                        clip_voxels = clip_voxels0\n",
+    "                        backbone = backbone0\n",
+    "                        logits = logits0\n",
+    "                    else:\n",
+    "                        clip_voxels += clip_voxels0\n",
+    "                        backbone += backbone0\n",
+    "                        logits += logits0\n",
+    "                clip_voxels /= 3\n",
+    "                backbone /= 3\n",
+    "                logits /= 3\n",
+    "\n",
+    "                print(logits.shape, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).shape)\n",
+    "                RR_loss = nn.functional.cross_entropy(logits, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).to(logits.device))\n",
+    "                test_loss_RR_total += RR_loss.item()\n",
+    "                loss += RR_loss\n",
+    "\n",
+    "                if clip_scale>0:\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",
+    "                # for some evals, only doing a subset of the samples per batch because of computational cost\n",
+    "                random_samps = np.random.choice(np.arange(len(image)), size=len(image)//5, replace=False)\n",
+    "                \n",
+    "                # if use_prior:\n",
+    "                #     loss_prior, contaminated_prior_out = model.diffusion_prior(text_embed=backbone[random_samps], image_embed=clip_target[random_samps])\n",
+    "                #     test_loss_prior_total += loss_prior.item()\n",
+    "                #     loss_prior *= prior_scale\n",
+    "                #     loss += loss_prior\n",
+    "                    \n",
+    "                if clip_scale>0:\n",
+    "                    loss_clip = utils.soft_clip_loss(\n",
+    "                        clip_voxels_norm,\n",
+    "                        clip_target_norm,\n",
+    "                        temp=.006)\n",
+    "\n",
+    "                    test_loss_clip_total += loss_clip.item()\n",
+    "                    loss_clip = loss_clip * clip_scale\n",
+    "                    loss += loss_clip\n",
+    "\n",
+    "                if blurry_recon:\n",
+    "                    image_enc_pred, _ = blurry_image_enc_\n",
+    "                    blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample / 2 + 0.5).clamp(0,1)\n",
+    "                    pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)\n",
+    "                    test_blurry_pixcorr += pixcorr.item()\n",
+    "\n",
+    "                if clip_scale>0:\n",
+    "                    # forward and backward top 1 accuracy        \n",
+    "                    labels = torch.arange(len(clip_voxels_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).item()\n",
+    "                    test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()\n",
+    "\n",
+    "                test_class_precisions_1 += classPrecision(logits, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).to(logits.device))\n",
+    "                test_class_precisions_5 += classPrecision(logits, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).to(logits.device), 5)\n",
+    "                test_class_precisions_10 += classPrecision(logits, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).to(logits.device), 10)\n",
+    "\n",
+    "                \n",
+    "                utils.check_loss(loss)                \n",
+    "                test_losses.append(loss.item())\n",
+    "\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",
+    "                \"train/loss_blurry_cont_total\": loss_blurry_cont_total / (train_i + 1),\n",
+    "                \"test/loss_clip_total\": test_loss_clip_total / (test_i + 1),\n",
+    "                \"train/blurry_pixcorr\": blurry_pixcorr / (train_i + 1),\n",
+    "                \"test/blurry_pixcorr\": test_blurry_pixcorr / (test_i + 1),\n",
+    "                \"train/recon_cossim\": recon_cossim / (train_i + 1),\n",
+    "                \"test/recon_cossim\": test_recon_cossim / (test_i + 1),\n",
+    "                \"train/recon_mse\": recon_mse / (train_i + 1),\n",
+    "                \"test/recon_mse\": test_recon_mse / (test_i + 1),\n",
+    "                \"train/loss_prior\": loss_prior_total / (train_i + 1),\n",
+    "                \"test/loss_prior\": test_loss_prior_total / (test_i + 1),\n",
+    "                \"train/loss_RR\": loss_RR_total / (train_i + 1),\n",
+    "                \"test/loss_RR\": test_loss_RR_total / (test_i + 1),\n",
+    "                \"train/class_precisions_1\": class_precisions_1 / (train_i + 1),\n",
+    "                \"test/class_precisions_1\": test_class_precisions_1 / (test_i + 1),\n",
+    "                \"train/class_precisions_5\": class_precisions_5 / (train_i + 1),\n",
+    "                \"test/class_precisions_5\": test_class_precisions_5 / (test_i + 1),\n",
+    "                \"train/class_precisions_10\": class_precisions_10 / (train_i + 1),\n",
+    "                \"test/class_precisions_10\": test_class_precisions_10 / (test_i + 1),\n",
+    "                }\n",
+    "\n",
+    "            # if finished training, save jpg recons if they exist\n",
+    "            if (epoch == num_epochs-1) or (epoch % ckpt_interval == 0):\n",
+    "                if blurry_recon:    \n",
+    "                    image_enc = autoenc.encode(2*image[:4]-1).latent_dist.mode() * 0.18215\n",
+    "                    # transform blurry recon latents to images and plot it\n",
+    "                    fig, axes = plt.subplots(1, 8, figsize=(10, 4))\n",
+    "                    jj=-1\n",
+    "                    for j in [0,1,2,3]:\n",
+    "                        jj+=1\n",
+    "                        axes[jj].imshow(utils.torch_to_Image((autoenc.decode(image_enc[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))\n",
+    "                        axes[jj].axis('off')\n",
+    "                        jj+=1\n",
+    "                        axes[jj].imshow(utils.torch_to_Image((autoenc.decode(image_enc_pred[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))\n",
+    "                        axes[jj].axis('off')\n",
+    "\n",
+    "                    if wandb_log:\n",
+    "                        logs[f\"test/blur_recons\"] = wandb.Image(fig, caption=f\"epoch{epoch:03d}\")\n",
+    "                        plt.close()\n",
+    "                    else:\n",
+    "                        plt.show()\n",
+    "\n",
+    "            progress_bar.set_postfix(**logs)\n",
+    "\n",
+    "            if wandb_log: wandb.log(logs)\n",
+    "            \n",
+    "    # Save model checkpoint and reconstruct\n",
+    "    if (ckpt_saving) and (epoch % ckpt_interval == 0):\n",
+    "        save_ckpt(f'last')\n",
+    "\n",
+    "    # wait for other GPUs to catch up if needed\n",
+    "    accelerator.wait_for_everyone()\n",
+    "    torch.cuda.empty_cache()\n",
+    "\n",
+    "print(\"\\n===Finished!===\\n\")\n",
+    "if ckpt_saving:\n",
+    "    save_ckpt(f'last')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "id": "4211dd56-6a12-493e-8ccf-9b5229195cbd",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'photo of a person': 0,\n",
+       " 'photo of airplane': 1,\n",
+       " 'photo of baseball': 2,\n",
+       " 'photo of bathroom': 3,\n",
+       " 'photo of bear': 4,\n",
+       " 'photo of bedroom': 5,\n",
+       " 'photo of bike': 6,\n",
+       " 'photo of bird': 7,\n",
+       " 'photo of boat': 8,\n",
+       " 'photo of bus': 9,\n",
+       " 'photo of cat': 10,\n",
+       " 'photo of clocktower': 11,\n",
+       " 'photo of computer': 12,\n",
+       " 'photo of cow': 13,\n",
+       " 'photo of dog': 14,\n",
+       " 'photo of elephant': 15,\n",
+       " 'photo of flower': 16,\n",
+       " 'photo of food': 17,\n",
+       " 'photo of fruits': 18,\n",
+       " 'photo of giraffe': 19,\n",
+       " 'photo of group of people': 20,\n",
+       " 'photo of horse': 21,\n",
+       " 'photo of hydrant': 22,\n",
+       " 'photo of living room': 23,\n",
+       " 'photo of person eating': 24,\n",
+       " 'photo of pizza': 25,\n",
+       " 'photo of sheep': 26,\n",
+       " 'photo of skate': 27,\n",
+       " 'photo of ski': 28,\n",
+       " 'photo of sky': 29,\n",
+       " 'photo of soccer': 30,\n",
+       " 'photo of sports': 31,\n",
+       " 'photo of stop sign': 32,\n",
+       " 'photo of surfer': 33,\n",
+       " 'photo of sweets': 34,\n",
+       " 'photo of tennis': 35,\n",
+       " 'photo of toy': 36,\n",
+       " 'photo of train': 37,\n",
+       " 'photo of umbrella': 38,\n",
+       " 'photo of vehicle': 39,\n",
+       " 'photo of zebra': 40}"
+      ]
+     },
+     "execution_count": 27,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "semantic_cluster_dict"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a7e81ae3-171f-40ad-a3e8-24bee4472325",
+   "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": "45d5bc17-0914-43a5-bb4e-d98f0ba238f0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import wandb\n",
+    "wandb.login()"
+   ]
+  }
+ ],
+ "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.11.9"
+  },
+  "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
+}

MindEyeV2/src/TrainCluster.ipynb

@@ -0,0 +1,1785 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b0f0f4f3",
+   "metadata": {},
+   "source": [
+    "# Import packages & functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "5bad764b-45c1-45ce-a716-8d055e09821a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "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 string\n",
+    "import h5py\n",
+    "from tqdm import tqdm\n",
+    "import webdataset as wds\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torchvision import transforms\n",
+    "from accelerate import Accelerator\n",
+    "\n",
+    "# SDXL unCLIP requires code from https://github.com/Stability-AI/generative-models/tree/main\n",
+    "sys.path.append('generative_models/')\n",
+    "import sgm\n",
+    "from generative_models.sgm.modules.encoders.modules import FrozenOpenCLIPImageEmbedder # bigG embedder\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"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "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",
+    "data_type = torch.float16 # change depending on your mixed_precision\n",
+    "num_devices = torch.cuda.device_count()\n",
+    "if num_devices==0: num_devices = 1\n",
+    "\n",
+    "# First use \"accelerate config\" in terminal and setup using deepspeed stage 2 with CPU offloading!\n",
+    "accelerator = Accelerator(split_batches=False, mixed_precision=\"fp16\")\n",
+    "if utils.is_interactive(): # set batch size here if using interactive notebook instead of submitting job\n",
+    "    global_batch_size = batch_size = 8\n",
+    "else:\n",
+    "    global_batch_size = os.environ[\"GLOBAL_BATCH_SIZE\"]\n",
+    "    batch_size = int(os.environ[\"GLOBAL_BATCH_SIZE\"]) // num_devices"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "b767ab6f-d4a9-47a5-b3bf-f56bf6760c0c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "PID of this process = 1132958\n",
+      "device: cuda\n",
+      "Distributed environment: DistributedType.NO\n",
+      "Num processes: 1\n",
+      "Process index: 0\n",
+      "Local process index: 0\n",
+      "Device: cuda\n",
+      "\n",
+      "Mixed precision type: fp16\n",
+      "\n",
+      "distributed = False num_devices = 1 local rank = 0 world size = 1 data_type = torch.float16\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(\"PID of this process =\",os.getpid())\n",
+    "device = accelerator.device\n",
+    "print(\"device:\",device)\n",
+    "world_size = accelerator.state.num_processes\n",
+    "distributed = not accelerator.state.distributed_type == 'NO'\n",
+    "num_devices = torch.cuda.device_count()\n",
+    "if num_devices==0 or not distributed: num_devices = 1\n",
+    "num_workers = num_devices\n",
+    "print(accelerator.state)\n",
+    "\n",
+    "print(\"distributed =\",distributed, \"num_devices =\", num_devices, \"local rank =\", local_rank, \"world size =\", world_size, \"data_type =\", data_type)\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": 4,
+   "id": "2b61fec7-72a0-4b67-86da-1375f1d9fbd3",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "model_name: testing2\n",
+      "--data_path=/weka/proj-medarc/shared/mindeyev2_dataset                     --cache_dir=/weka/proj-medarc/shared/cache                     --model_name=testing2                     --no-multi_subject --subj=1 --batch_size=8 --num_sessions=40                     --hidden_dim=1024 --clip_scale=1.                     --no-blurry_recon --blur_scale=.5                      --use_prior --prior_scale=30                     --n_blocks=4 --max_lr=3e-5 --mixup_pct=.33 --num_epochs=150 --no-use_image_aug                     --ckpt_interval=999 --no-ckpt_saving --wandb_log\n"
+     ]
+    }
+   ],
+   "source": [
+    "# if running this interactively, can specify jupyter_args here for argparser to use\n",
+    "if utils.is_interactive():\n",
+    "    model_name = \"testing2\"\n",
+    "    print(\"model_name:\", model_name)\n",
+    "    \n",
+    "    # global_batch_size and batch_size should already be defined in the 2nd cell block\n",
+    "    jupyter_args = f\"--data_path=/weka/proj-medarc/shared/mindeyev2_dataset \\\n",
+    "                    --cache_dir=/weka/proj-medarc/shared/cache \\\n",
+    "                    --model_name={model_name} \\\n",
+    "                    --no-multi_subject --subj=1 --batch_size={batch_size} --num_sessions=40 \\\n",
+    "                    --hidden_dim=1024 --clip_scale=1. \\\n",
+    "                    --no-blurry_recon --blur_scale=.5  \\\n",
+    "                    --use_prior --prior_scale=30 \\\n",
+    "                    --n_blocks=4 --max_lr=3e-5 --mixup_pct=.33 --num_epochs=150 --no-use_image_aug \\\n",
+    "                    --ckpt_interval=999 --no-ckpt_saving --wandb_log\"\n",
+    "    # --multisubject_ckpt=../train_logs/multisubject_subj01_1024_24bs_nolow\n",
+    "\n",
+    "    print(jupyter_args)\n",
+    "    jupyter_args = jupyter_args.split()\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": 5,
+   "id": "2028bdf0-2f41-46d9-b6e7-86b870dbf16c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "subj_list [1] num_sessions 40\n"
+     ]
+    }
+   ],
+   "source": [
+    "parser = argparse.ArgumentParser(description=\"Model Training Configuration\")\n",
+    "parser.add_argument(\n",
+    "    \"--model_name\", type=str, default=\"testing2\",\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=os.getcwd(),\n",
+    "    help=\"Path to where NSD data is stored / where to download it to\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--cache_dir\", type=str, default=os.getcwd(),\n",
+    "    help=\"Path to where misc. files downloaded from huggingface are stored. Defaults to current src directory.\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--subj\",type=int, default=1, choices=[1,2,3,4,5,6,7,8],\n",
+    "    help=\"Validate on which subject?\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--multisubject_ckpt\", type=str, default=None,\n",
+    "    help=\"Path to pre-trained multisubject model to finetune a single subject from. multisubject must be False.\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--num_sessions\", type=int, default=1,\n",
+    "    help=\"Number of training sessions to include\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--use_prior\",action=argparse.BooleanOptionalAction,default=True,\n",
+    "    help=\"whether to train diffusion prior (True) or just rely on retrieval part of the pipeline (False)\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--batch_size\", type=int, default=16,\n",
+    "    help=\"Batch size can be increased by 10x if only training retreival submodule and not diffusion prior\",\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",
+    "    \"--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",
+    "    \"--blurry_recon\",action=argparse.BooleanOptionalAction,default=True,\n",
+    "    help=\"whether to output blurry reconstructions\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--blur_scale\",type=float,default=.5,\n",
+    "    help=\"multiply loss from blurry recons by this number\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--clip_scale\",type=float,default=1.,\n",
+    "    help=\"multiply contrastive loss by this number\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--prior_scale\",type=float,default=30,\n",
+    "    help=\"multiply diffusion prior loss by this\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--use_image_aug\",action=argparse.BooleanOptionalAction,default=False,\n",
+    "    help=\"whether to use image augmentation\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--num_epochs\",type=int,default=150,\n",
+    "    help=\"number of epochs of training\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--multi_subject\",action=argparse.BooleanOptionalAction,default=False,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--new_test\",action=argparse.BooleanOptionalAction,default=True,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--n_blocks\",type=int,default=4,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--hidden_dim\",type=int,default=1024,\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-5,\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",
+    "# seed all random functions\n",
+    "utils.seed_everything(seed)\n",
+    "\n",
+    "outdir = os.path.abspath(f'../train_logs/{model_name}')\n",
+    "if not os.path.exists(outdir) and ckpt_saving:\n",
+    "    os.makedirs(outdir,exist_ok=True)\n",
+    "    \n",
+    "if use_image_aug or blurry_recon:\n",
+    "    import kornia\n",
+    "    from kornia.augmentation.container import AugmentationSequential\n",
+    "if use_image_aug:\n",
+    "    img_augment = AugmentationSequential(\n",
+    "        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1, p=0.3),\n",
+    "        same_on_batch=False,\n",
+    "        data_keys=[\"input\"],\n",
+    "    )\n",
+    "    \n",
+    "if multi_subject:\n",
+    "    subj_list = np.arange(1,9)\n",
+    "    subj_list = subj_list[subj_list != subj]\n",
+    "else:\n",
+    "    subj_list = [subj]\n",
+    "\n",
+    "print(\"subj_list\", subj_list, \"num_sessions\", num_sessions)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "2fb18bf3-27f5-470c-be05-002215d391b9",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "3e-05"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "max_lr"
+   ]
+  },
+  {
+   "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": [
+    "### Creating wds dataloader, preload betas and all 73k possible images"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "aefe7c27-ab39-4b2c-90f4-480f4087b7ab",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "dividing batch size by subj_list, which will then be concatenated across subj during training...\n",
+      "batch_size = 8 num_iterations_per_epoch = 3750 num_samples_per_epoch = 30000\n"
+     ]
+    }
+   ],
+   "source": [
+    "def my_split_by_node(urls): return urls\n",
+    "num_voxels_list = []\n",
+    "\n",
+    "if multi_subject:\n",
+    "    nsessions_allsubj=np.array([40, 40, 32, 30, 40, 32, 40, 30])\n",
+    "    num_samples_per_epoch = (750*40) // num_devices \n",
+    "else:\n",
+    "    num_samples_per_epoch = (750*num_sessions) // num_devices \n",
+    "\n",
+    "print(\"dividing batch size by subj_list, which will then be concatenated across subj during training...\") \n",
+    "batch_size = batch_size // len(subj_list)\n",
+    "\n",
+    "num_iterations_per_epoch = num_samples_per_epoch // (batch_size*len(subj_list))\n",
+    "\n",
+    "print(\"batch_size =\", batch_size, \"num_iterations_per_epoch =\",num_iterations_per_epoch, \"num_samples_per_epoch =\",num_samples_per_epoch)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "81084834-035f-4465-ad59-59e6b806a2f5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Training with 40 sessions\n",
+      "/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar\n",
+      "num_voxels for subj01: 15724\n",
+      "Loaded all subj train dls and betas!\n",
+      "\n",
+      "/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar\n",
+      "Loaded test dl for subj1!\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "train_data = {}\n",
+    "train_dl = {}\n",
+    "num_voxels = {}\n",
+    "voxels = {}\n",
+    "for s in subj_list:\n",
+    "    print(f\"Training with {num_sessions} sessions\")\n",
+    "    if multi_subject:\n",
+    "        train_url = f\"{data_path}/wds/subj0{s}/train/\" + \"{0..\" + f\"{nsessions_allsubj[s-1]-1}\" + \"}.tar\"\n",
+    "    else:\n",
+    "        train_url = f\"{data_path}/wds/subj0{s}/train/\" + \"{0..\" + f\"{num_sessions-1}\" + \"}.tar\"\n",
+    "    print(train_url)\n",
+    "    \n",
+    "    train_data[f'subj0{s}'] = wds.WebDataset(train_url,resampled=True,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[f'subj0{s}'] = torch.utils.data.DataLoader(train_data[f'subj0{s}'], batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)\n",
+    "\n",
+    "    f = h5py.File(f'{data_path}/betas_all_subj0{s}_fp32_renorm.hdf5', 'r')\n",
+    "    betas = f['betas'][:]\n",
+    "    betas = torch.Tensor(betas).to(\"cpu\").to(data_type)\n",
+    "    num_voxels_list.append(betas[0].shape[-1])\n",
+    "    num_voxels[f'subj0{s}'] = betas[0].shape[-1]\n",
+    "    voxels[f'subj0{s}'] = betas\n",
+    "    print(f\"num_voxels for subj0{s}: {num_voxels[f'subj0{s}']}\")\n",
+    "\n",
+    "print(\"Loaded all subj train dls and betas!\\n\")\n",
+    "\n",
+    "# Validate only on one subject\n",
+    "if multi_subject: \n",
+    "    subj = subj_list[0] # cant validate on the actual held out person so picking first in subj_list\n",
+    "if not new_test: # using old test set from before full dataset released (used in original MindEye paper)\n",
+    "    if subj==3:\n",
+    "        num_test=2113\n",
+    "    elif subj==4:\n",
+    "        num_test=1985\n",
+    "    elif subj==6:\n",
+    "        num_test=2113\n",
+    "    elif subj==8:\n",
+    "        num_test=1985\n",
+    "    else:\n",
+    "        num_test=2770\n",
+    "    test_url = f\"{data_path}/wds/subj0{subj}/test/\" + \"0.tar\"\n",
+    "elif new_test: # using larger test set from after full dataset released\n",
+    "    if subj==3:\n",
+    "        num_test=2371\n",
+    "    elif subj==4:\n",
+    "        num_test=2188\n",
+    "    elif subj==6:\n",
+    "        num_test=2371\n",
+    "    elif subj==8:\n",
+    "        num_test=2188\n",
+    "    else:\n",
+    "        num_test=3000\n",
+    "    test_url = f\"{data_path}/wds/subj0{subj}/new_test/\" + \"0.tar\"\n",
+    "print(test_url)\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=num_test, shuffle=False, drop_last=True, pin_memory=True)\n",
+    "print(f\"Loaded test dl for subj{subj}!\\n\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "c13b4b84-094c-4b5b-bace-26c155aa6181",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Loaded all 73k possible NSD images to cpu! (73000, 3, 224, 224)\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Load 73k NSD images\n",
+    "f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')\n",
+    "images = f['images']\n",
+    "print(\"Loaded all 73k possible NSD images to cpu!\", 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": 10,
+   "id": "b0420dc0-199e-4c1a-857d-b1747058b467",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "clip_img_embedder = FrozenOpenCLIPImageEmbedder(\n",
+    "    arch=\"ViT-bigG-14\",\n",
+    "    version=\"laion2b_s39b_b160k\",\n",
+    "    output_tokens=True,\n",
+    "    only_tokens=True,\n",
+    ")\n",
+    "clip_img_embedder.to(device)\n",
+    "\n",
+    "clip_seq_dim = 256\n",
+    "clip_emb_dim = 1664"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5b79bd38-6990-4504-8d45-4a68d57d8885",
+   "metadata": {},
+   "source": [
+    "### SD VAE"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "01baff79-8114-482b-b115-6f05aa8ad691",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "if blurry_recon:\n",
+    "    from diffusers import AutoencoderKL    \n",
+    "    autoenc = AutoencoderKL(\n",
+    "        down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D'],\n",
+    "        up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D'],\n",
+    "        block_out_channels=[128, 256, 512, 512],\n",
+    "        layers_per_block=2,\n",
+    "        sample_size=256,\n",
+    "    )\n",
+    "    ckpt = torch.load(f'{cache_dir}/sd_image_var_autoenc.pth')\n",
+    "    autoenc.load_state_dict(ckpt)\n",
+    "    \n",
+    "    autoenc.eval()\n",
+    "    autoenc.requires_grad_(False)\n",
+    "    autoenc.to(device)\n",
+    "    utils.count_params(autoenc)\n",
+    "    \n",
+    "    from autoencoder.convnext import ConvnextXL\n",
+    "    cnx = ConvnextXL(f'{cache_dir}/convnext_xlarge_alpha0.75_fullckpt.pth')\n",
+    "    cnx.requires_grad_(False)\n",
+    "    cnx.eval()\n",
+    "    cnx.to(device)\n",
+    "    \n",
+    "    mean = torch.tensor([0.485, 0.456, 0.406]).to(device).reshape(1,3,1,1)\n",
+    "    std = torch.tensor([0.228, 0.224, 0.225]).to(device).reshape(1,3,1,1)\n",
+    "    \n",
+    "    blur_augs = AugmentationSequential(\n",
+    "        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.8),\n",
+    "        kornia.augmentation.RandomGrayscale(p=0.1),\n",
+    "        kornia.augmentation.RandomSolarize(p=0.1),\n",
+    "        kornia.augmentation.RandomResizedCrop((224,224), scale=(.9,.9), ratio=(1,1), p=1.0),\n",
+    "        data_keys=[\"input\"],\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "260e5e4a-f697-4b2c-88fc-01f6a54886c0",
+   "metadata": {},
+   "source": [
+    "### MindEye modules"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "c44c271b-173f-472e-b059-a2eda0f4c4c5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "MindEyeModule()"
+      ]
+     },
+     "execution_count": 12,
+     "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": 13,
+   "id": "038a5d61-4769-40b9-a004-f4e7b5b38bb0",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "16,102,400 total\n",
+      "16,102,400 trainable\n",
+      "param counts:\n",
+      "16,102,400 total\n",
+      "16,102,400 trainable\n",
+      "torch.Size([2, 1, 15724]) torch.Size([2, 1, 1024])\n"
+     ]
+    }
+   ],
+   "source": [
+    "class RidgeRegression(torch.nn.Module):\n",
+    "    # make sure to add weight_decay when initializing optimizer to enable regularization\n",
+    "    def __init__(self, input_sizes, out_features): \n",
+    "        super(RidgeRegression, self).__init__()\n",
+    "        self.out_features = out_features\n",
+    "        self.linears = torch.nn.ModuleList([\n",
+    "                torch.nn.Linear(input_size, out_features) for input_size in input_sizes\n",
+    "            ])\n",
+    "    def forward(self, x, subj_idx):\n",
+    "        out = self.linears[subj_idx](x[:,0]).unsqueeze(1)\n",
+    "        return out\n",
+    "        \n",
+    "class IndividRidgeRegression(torch.nn.Module):\n",
+    "    def __init__(self, input_size, out_features):\n",
+    "        super(IndividRidgeRegression, self).__init__()\n",
+    "        self.out_features = out_features\n",
+    "        self.linear = torch.nn.Linear(input_size, out_features)\n",
+    "    def forward(self, x):\n",
+    "        out = self.linear(x)\n",
+    "        return out\n",
+    "    \n",
+    "model.ridge = RidgeRegression(num_voxels_list, out_features=hidden_dim)\n",
+    "utils.count_params(model.ridge)\n",
+    "utils.count_params(model)\n",
+    "\n",
+    "# test on subject 1 with fake data\n",
+    "b = torch.randn((2,1,num_voxels_list[0]))\n",
+    "print(b.shape, model.ridge(b,0).shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "7b8de65a-6d3b-4248-bea9-9b6f4d562321",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "453,360,280 total\n",
+      "453,360,280 trainable\n",
+      "param counts:\n",
+      "469,462,680 total\n",
+      "469,462,680 trainable\n",
+      "b.shape torch.Size([2, 1, 1024])\n",
+      "torch.Size([2, 256, 1664]) torch.Size([2, 256, 1664]) torch.Size([1]) torch.Size([1])\n"
+     ]
+    }
+   ],
+   "source": [
+    "from models import BrainNetwork\n",
+    "model.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, seq_len=1, n_blocks=n_blocks,\n",
+    "                          clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, \n",
+    "                          blurry_recon=blurry_recon, clip_scale=clip_scale)\n",
+    "utils.count_params(model.backbone)\n",
+    "utils.count_params(model)\n",
+    "\n",
+    "# test that the model works on some fake data\n",
+    "b = torch.randn((2,1,hidden_dim))\n",
+    "print(\"b.shape\",b.shape)\n",
+    "\n",
+    "backbone_, clip_, blur_ = model.backbone(b)\n",
+    "print(backbone_.shape, clip_.shape, blur_[0].shape, blur_[1].shape)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e63b26a9-3a36-4638-be3c-dbf28705dd76",
+   "metadata": {},
+   "source": [
+    "### Load semantic clusters"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "2950add9-3828-4d55-bbf7-0d9f3519de5c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "semantic_cluster_onehot.shape torch.Size([73024, 41])\n",
+      "num_seman_clusters 41\n"
+     ]
+    }
+   ],
+   "source": [
+    "path_semantic_names = \"/weka/proj-medarc/shared/mindeyev2_dataset/semantic_cluster_names.npy\"\n",
+    "path_semantic_cluster = \"/weka/proj-medarc/shared/mindeyev2_dataset/COCO_73k_semantic_cluster.npy\"\n",
+    "semantic_cluster_names = np.load(path_semantic_names)\n",
+    "semantic_cluster = np.load(path_semantic_cluster)\n",
+    "possible_semantic_clusters = np.unique(semantic_cluster)\n",
+    "\n",
+    "# one-hot encode semantic clusters\n",
+    "# move possible_semantic_clusters to numbers and create a dictionary\n",
+    "semantic_cluster_dict = {cluster: i for i, cluster in enumerate(possible_semantic_clusters)}\n",
+    "semantic_cluster_onehot = torch.zeros((len(semantic_cluster), len(possible_semantic_clusters)))\n",
+    "for i, cluster in enumerate(semantic_cluster):\n",
+    "    semantic_cluster_onehot[i, semantic_cluster_dict[cluster]] = 1\n",
+    "\n",
+    "\n",
+    "print(\"semantic_cluster_onehot.shape\", semantic_cluster_onehot.shape)\n",
+    "\n",
+    "num_seman_clusters = len(np.unique(semantic_cluster))\n",
+    "print(\"num_seman_clusters\", num_seman_clusters)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b397c0d7-52a3-4153-823b-c27d2eb3eeba",
+   "metadata": {},
+   "source": [
+    "### Adding the ridge regression to the class"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "69965344-9346-4592-9cc5-e537e31d5fce",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "param counts:\n",
+      "17,465,385 total\n",
+      "17,465,385 trainable\n",
+      "param counts:\n",
+      "486,928,065 total\n",
+      "486,928,065 trainable\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "486928065"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# if use_prior:\n",
+    "#     from models import *\n",
+    "\n",
+    "#     # setup diffusion prior network\n",
+    "#     out_dim = clip_emb_dim\n",
+    "#     depth = 6\n",
+    "#     dim_head = 52\n",
+    "#     heads = clip_emb_dim//52 # heads * dim_head = clip_emb_dim\n",
+    "#     timesteps = 100\n",
+    "\n",
+    "#     prior_network = PriorNetwork(\n",
+    "#             dim=out_dim,\n",
+    "#             depth=depth,\n",
+    "#             dim_head=dim_head,\n",
+    "#             heads=heads,\n",
+    "#             causal=False,\n",
+    "#             num_tokens = clip_seq_dim,\n",
+    "#             learned_query_mode=\"pos_emb\"\n",
+    "#         )\n",
+    "\n",
+    "#     model.diffusion_prior = BrainDiffusionPrior(\n",
+    "#         net=prior_network,\n",
+    "#         image_embed_dim=out_dim,\n",
+    "#         condition_on_text_encodings=False,\n",
+    "#         timesteps=timesteps,\n",
+    "#         cond_drop_prob=0.2,\n",
+    "#         image_embed_scale=None,\n",
+    "#     )\n",
+    "    \n",
+    "#     utils.count_params(model.diffusion_prior)\n",
+    "#     utils.count_params(model)\n",
+    "\n",
+    "model.RRClassifier = IndividRidgeRegression(clip_emb_dim*clip_seq_dim, out_features=num_seman_clusters)\n",
+    "utils.count_params(model.RRClassifier)\n",
+    "utils.count_params(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ec25271a-2209-400c-8026-df3b8ddc1eef",
+   "metadata": {},
+   "source": [
+    "### Setup optimizer / lr / ckpt saving"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "e14d0482-dc42-43b9-9ce1-953c32f2c9c1",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "total_steps 562500\n",
+      "\n",
+      "Done with model preparations!\n",
+      "param counts:\n",
+      "486,928,065 total\n",
+      "486,928,065 trainable\n"
+     ]
+    }
+   ],
+   "source": [
+    "no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n",
+    "\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",
+    "# if use_prior:\n",
+    "#     opt_grouped_parameters.extend([\n",
+    "#         {'params': [p for n, p in model.diffusion_prior.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.diffusion_prior.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}\n",
+    "#     ])\n",
+    "opt_grouped_parameters.extend([\n",
+    "    {'params': [p for n, p in model.RRClassifier.named_parameters()], 'weight_decay': 1e-2},\n",
+    "])\n",
+    "\n",
+    "optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr)\n",
+    "\n",
+    "if lr_scheduler_type == 'linear':\n",
+    "    lr_scheduler = torch.optim.lr_scheduler.LinearLR(\n",
+    "        optimizer,\n",
+    "        total_iters=int(np.floor(num_epochs*num_iterations_per_epoch)),\n",
+    "        last_epoch=-1\n",
+    "    )\n",
+    "elif lr_scheduler_type == 'cycle':\n",
+    "    total_steps=int(np.floor(num_epochs*num_iterations_per_epoch))\n",
+    "    print(\"total_steps\", total_steps)\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",
+    "    if accelerator.is_main_process:\n",
+    "        unwrapped_model = accelerator.unwrap_model(model)\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",
+    "    print(f\"\\n---saved {outdir}/{tag} ckpt!---\\n\")\n",
+    "\n",
+    "def load_ckpt(tag,load_lr=True,load_optimizer=True,load_epoch=True,strict=True,outdir=outdir,multisubj_loading=False): \n",
+    "    print(f\"\\n---loading {outdir}/{tag}.pth ckpt---\\n\")\n",
+    "    checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')\n",
+    "    state_dict = checkpoint['model_state_dict']\n",
+    "    if multisubj_loading: # remove incompatible ridge layer that will otherwise error\n",
+    "        state_dict.pop('ridge.linears.0.weight',None)\n",
+    "    model.load_state_dict(state_dict, strict=strict)\n",
+    "    if load_epoch:\n",
+    "        globals()[\"epoch\"] = checkpoint['epoch']\n",
+    "        print(\"Epoch\",epoch)\n",
+    "    if load_optimizer:\n",
+    "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
+    "    if load_lr:\n",
+    "        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n",
+    "    del checkpoint\n",
+    "\n",
+    "print(\"\\nDone with model preparations!\")\n",
+    "num_params = utils.count_params(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "983f458b-35b8-49f2-b6db-80296cece730",
+   "metadata": {},
+   "source": [
+    "# Weights and Biases"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "0a25a662-daa8-4de9-9233-8364800fcb6b",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "wandb mindeye_semantic_cluster run testing2\n",
+      "wandb_config:\n",
+      " {'model_name': 'testing2', 'global_batch_size': 8, 'batch_size': 8, 'num_epochs': 150, 'num_sessions': 40, 'num_params': 486928065, 'clip_scale': 1.0, 'prior_scale': 30.0, 'blur_scale': 0.5, 'use_image_aug': False, 'max_lr': 3e-05, 'mixup_pct': 0.33, 'num_samples_per_epoch': 30000, 'num_test': 3000, 'ckpt_interval': 999, 'ckpt_saving': False, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1, 'train_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar', 'test_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar'}\n",
+      "wandb_id: testing2\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"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "wandb version 0.17.2 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.17.1"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Run data is saved locally in <code>/weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_133545-testing2</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Resuming run <strong><a href='https://stability.wandb.io/ckadirt/mindeye_semantic_cluster/runs/testing2' target=\"_blank\">testing2</a></strong> to <a href='https://stability.wandb.io/ckadirt/mindeye_semantic_cluster' 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/mindeye_semantic_cluster' target=\"_blank\">https://stability.wandb.io/ckadirt/mindeye_semantic_cluster</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/mindeye_semantic_cluster/runs/testing2' target=\"_blank\">https://stability.wandb.io/ckadirt/mindeye_semantic_cluster/runs/testing2</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "if local_rank==0 and wandb_log: # only use main process for wandb logging\n",
+    "    import wandb\n",
+    "    wandb_project = 'mindeye_semantic_cluster'\n",
+    "    print(f\"wandb {wandb_project} run {model_name}\")\n",
+    "    # need to configure wandb beforehand in terminal with \"wandb init\"!\n",
+    "    wandb_config = {\n",
+    "      \"model_name\": model_name,\n",
+    "      \"global_batch_size\": global_batch_size,\n",
+    "      \"batch_size\": batch_size,\n",
+    "      \"num_epochs\": num_epochs,\n",
+    "      \"num_sessions\": num_sessions,\n",
+    "      \"num_params\": num_params,\n",
+    "      \"clip_scale\": clip_scale,\n",
+    "      \"prior_scale\": prior_scale,\n",
+    "      \"blur_scale\": blur_scale,\n",
+    "      \"use_image_aug\": use_image_aug,\n",
+    "      \"max_lr\": max_lr,\n",
+    "      \"mixup_pct\": mixup_pct,\n",
+    "      \"num_samples_per_epoch\": num_samples_per_epoch,\n",
+    "      \"num_test\": num_test,\n",
+    "      \"ckpt_interval\": ckpt_interval,\n",
+    "      \"ckpt_saving\": ckpt_saving,\n",
+    "      \"seed\": seed,\n",
+    "      \"distributed\": distributed,\n",
+    "      \"num_devices\": num_devices,\n",
+    "      \"world_size\": world_size,\n",
+    "      \"train_url\": train_url,\n",
+    "      \"test_url\": test_url,\n",
+    "    }\n",
+    "    print(\"wandb_config:\\n\",wandb_config)\n",
+    "    print(\"wandb_id:\",model_name)\n",
+    "    wandb.login(host='https://stability.wandb.io')\n",
+    "    wandb.init(\n",
+    "        id=model_name,\n",
+    "        project=wandb_project,\n",
+    "        name=model_name,\n",
+    "        config=wandb_config,\n",
+    "        resume=\"allow\",\n",
+    "    )\n",
+    "else:\n",
+    "    wandb_log = False"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d5690151-2131-4918-b750-e869cbd1a8a8",
+   "metadata": {},
+   "source": [
+    "# Main"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "12de6387-6e18-4e4b-b5ce-a847d625330a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "epoch = 0\n",
+    "losses, test_losses, lrs = [], [], []\n",
+    "best_test_loss = 1e9\n",
+    "torch.cuda.empty_cache()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "607a7c7b-fe5e-41a4-80bf-d2814b3a57cc",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# load multisubject stage1 ckpt if set\n",
+    "if multisubject_ckpt is not None:\n",
+    "    load_ckpt(\"last\",outdir=multisubject_ckpt,load_lr=False,load_optimizer=False,load_epoch=False,strict=False,multisubj_loading=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "99f09f76-4481-4133-b09a-a22b10dbc0c4",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "train_dls = [train_dl[f'subj0{s}'] for s in subj_list]\n",
+    "\n",
+    "model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot = accelerator.prepare(model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot)\n",
+    "# leaving out test_dl since we will only have local_rank 0 device do evals"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "e16fa583-064c-4cfa-9f47-08d08f53f504",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "3750\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(num_iterations_per_epoch)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "60be0d5f-3e94-4612-9373-61b53d836393",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "testing2 starting with epoch 0 / 150\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   | 0/150 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([300, 41]) torch.Size([300])\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  1%|████▏                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                               | 1/150 [06:36<16:25:13, 396.74s/it, test/blurry_pixcorr=0, test/loss=13, test/loss_RR=10.7, test/loss_clip_total=2.26, test/loss_prior=0, test/num_steps=1, test/recon_cossim=0, test/recon_mse=0, test/test_bwd_pct_correct=0.37, test/test_fwd_pct_correct=0.703, train/blurry_pixcorr=0, train/bwd_pct_correct=0.402, train/fwd_pct_correct=0.445, train/loss=3.54, train/loss_RR=2.21, train/loss_blurry_cont_total=0, train/loss_blurry_total=0, train/loss_clip_total=1.33, train/loss_prior=0, train/lr=1.56e-5, train/num_steps=3750, train/recon_cossim=0, train/recon_mse=0]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([300, 41]) torch.Size([300])\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  1%|████████▎                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                           | 2/150 [13:23<16:32:38, 402.42s/it, test/blurry_pixcorr=0, test/loss=42.5, test/loss_RR=41.1, test/loss_clip_total=1.37, test/loss_prior=0, test/num_steps=2, test/recon_cossim=0, test/recon_mse=0, test/test_bwd_pct_correct=0.613, test/test_fwd_pct_correct=0.827, train/blurry_pixcorr=0, train/bwd_pct_correct=0.717, train/fwd_pct_correct=0.708, train/loss=3.95, train/loss_RR=3.5, train/loss_blurry_cont_total=0, train/loss_blurry_total=0, train/loss_clip_total=0.448, train/loss_prior=0, train/lr=3e-5, train/num_steps=7500, train/recon_cossim=0, train/recon_mse=0]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([300, 41]) torch.Size([300])\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  2%|████████████▌                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                       | 3/150 [19:12<15:26:12, 378.04s/it, test/blurry_pixcorr=0, test/loss=76.7, test/loss_RR=75.5, test/loss_clip_total=1.18, test/loss_prior=0, test/num_steps=3, test/recon_cossim=0, test/recon_mse=0, test/test_bwd_pct_correct=0.72, test/test_fwd_pct_correct=0.89, train/blurry_pixcorr=0, train/bwd_pct_correct=0.746, train/fwd_pct_correct=0.725, train/loss=4.13, train/loss_RR=3.75, train/loss_blurry_cont_total=0, train/loss_blurry_total=0, train/loss_clip_total=0.376, train/loss_prior=0, train/lr=3e-5, train/num_steps=11250, train/recon_cossim=0, train/recon_mse=0]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([300, 41]) torch.Size([300])\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  3%|████████████████▋                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                  | 4/150 [23:52<13:46:10, 339.53s/it, test/blurry_pixcorr=0, test/loss=58.4, test/loss_RR=57.2, test/loss_clip_total=1.18, test/loss_prior=0, test/num_steps=4, test/recon_cossim=0, test/recon_mse=0, test/test_bwd_pct_correct=0.817, test/test_fwd_pct_correct=0.913, train/blurry_pixcorr=0, train/bwd_pct_correct=0.758, train/fwd_pct_correct=0.73, train/loss=3.77, train/loss_RR=3.44, train/loss_blurry_cont_total=0, train/loss_blurry_total=0, train/loss_clip_total=0.335, train/loss_prior=0, train/lr=3e-5, train/num_steps=15000, train/recon_cossim=0, train/recon_mse=0]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([300, 41]) torch.Size([300])\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  3%|████████████████████▉                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              | 5/150 [28:32<12:48:34, 318.03s/it, test/blurry_pixcorr=0, test/loss=58.6, test/loss_RR=57.5, test/loss_clip_total=1.12, test/loss_prior=0, test/num_steps=5, test/recon_cossim=0, test/recon_mse=0, test/test_bwd_pct_correct=0.797, test/test_fwd_pct_correct=0.903, train/blurry_pixcorr=0, train/bwd_pct_correct=0.763, train/fwd_pct_correct=0.733, train/loss=3.73, train/loss_RR=3.41, train/loss_blurry_cont_total=0, train/loss_blurry_total=0, train/loss_clip_total=0.32, train/loss_prior=0, train/lr=3e-5, train/num_steps=18750, train/recon_cossim=0, train/recon_mse=0]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([300, 41]) torch.Size([300])\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  4%|████████████████████████▉                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                      | 6/150 [33:12<12:12:10, 305.07s/it, test/blurry_pixcorr=0, test/loss=44.9, test/loss_RR=43.7, test/loss_clip_total=1.14, test/loss_prior=0, test/num_steps=6, test/recon_cossim=0, test/recon_mse=0, test/test_bwd_pct_correct=0.833, test/test_fwd_pct_correct=0.897, train/blurry_pixcorr=0, train/bwd_pct_correct=0.773, train/fwd_pct_correct=0.735, train/loss=3.67, train/loss_RR=3.37, train/loss_blurry_cont_total=0, train/loss_blurry_total=0, train/loss_clip_total=0.302, train/loss_prior=0, train/lr=2.99e-5, train/num_steps=22500, train/recon_cossim=0, train/recon_mse=0]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([300, 41]) torch.Size([300])\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  5%|█████████████████████████████                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                  | 7/150 [37:52<11:47:39, 296.92s/it, test/blurry_pixcorr=0, test/loss=65.1, test/loss_RR=64.1, test/loss_clip_total=0.981, test/loss_prior=0, test/num_steps=7, test/recon_cossim=0, test/recon_mse=0, test/test_bwd_pct_correct=0.867, test/test_fwd_pct_correct=0.93, train/blurry_pixcorr=0, train/bwd_pct_correct=0.774, train/fwd_pct_correct=0.738, train/loss=3.68, train/loss_RR=3.38, train/loss_blurry_cont_total=0, train/loss_blurry_total=0, train/loss_clip_total=0.294, train/loss_prior=0, train/lr=2.99e-5, train/num_steps=26250, train/recon_cossim=0, train/recon_mse=0]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([300, 41]) torch.Size([300])\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  5%|█████████████████████████████████▍                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                 | 8/150 [42:32<11:29:38, 291.40s/it, test/blurry_pixcorr=0, test/loss=102, test/loss_RR=101, test/loss_clip_total=1.01, test/loss_prior=0, test/num_steps=8, test/recon_cossim=0, test/recon_mse=0, test/test_bwd_pct_correct=0.823, test/test_fwd_pct_correct=0.927, train/blurry_pixcorr=0, train/bwd_pct_correct=0.771, train/fwd_pct_correct=0.735, train/loss=3.6, train/loss_RR=3.31, train/loss_blurry_cont_total=0, train/loss_blurry_total=0, train/loss_clip_total=0.291, train/loss_prior=0, train/lr=2.99e-5, train/num_steps=3e+4, train/recon_cossim=0, train/recon_mse=0]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([300, 41]) torch.Size([300])\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  6%|█████████████████████████████████████▌                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                           | 9/150 [47:13<11:17:18, 288.22s/it, test/blurry_pixcorr=0, test/loss=40.6, test/loss_RR=39.5, test/loss_clip_total=1.03, test/loss_prior=0, test/num_steps=9, test/recon_cossim=0, test/recon_mse=0, test/test_bwd_pct_correct=0.85, test/test_fwd_pct_correct=0.92, train/blurry_pixcorr=0, train/bwd_pct_correct=0.777, train/fwd_pct_correct=0.739, train/loss=3.63, train/loss_RR=3.34, train/loss_blurry_cont_total=0, train/loss_blurry_total=0, train/loss_clip_total=0.282, train/loss_prior=0, train/lr=2.98e-5, train/num_steps=33750, train/recon_cossim=0, train/recon_mse=0]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([300, 41]) torch.Size([300])\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  7%|█████████████████████████████████████████▎                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                  | 10/150 [1:21:00<18:54:10, 486.07s/it, test/blurry_pixcorr=0, test/loss=106, test/loss_RR=105, test/loss_clip_total=0.948, test/loss_prior=0, test/num_steps=10, test/recon_cossim=0, test/recon_mse=0, test/test_bwd_pct_correct=0.897, test/test_fwd_pct_correct=0.923, train/blurry_pixcorr=0, train/bwd_pct_correct=0.773, train/fwd_pct_correct=0.733, train/loss=3.61, train/loss_RR=3.33, train/loss_blurry_cont_total=0, train/loss_blurry_total=0, train/loss_clip_total=0.284, train/loss_prior=0, train/lr=2.98e-5, train/num_steps=37500, train/recon_cossim=0, train/recon_mse=0]\n"
+     ]
+    },
+    {
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[23], line 188\u001b[0m\n\u001b[1;32m    185\u001b[0m accelerator\u001b[38;5;241m.\u001b[39mbackward(loss)\n\u001b[1;32m    186\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mstep()\n\u001b[0;32m--> 188\u001b[0m losses\u001b[38;5;241m.\u001b[39mappend(\u001b[43mloss\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitem\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m    189\u001b[0m lrs\u001b[38;5;241m.\u001b[39mappend(optimizer\u001b[38;5;241m.\u001b[39mparam_groups[\u001b[38;5;241m0\u001b[39m][\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mlr\u001b[39m\u001b[38;5;124m'\u001b[39m])\n\u001b[1;32m    191\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m lr_scheduler_type \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
+     ]
+    }
+   ],
+   "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",
+    "l1 = nn.L1Loss()\n",
+    "soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))\n",
+    "\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",
+    "    recon_cossim = 0.\n",
+    "    test_recon_cossim = 0.\n",
+    "    recon_mse = 0.\n",
+    "    test_recon_mse = 0.\n",
+    "\n",
+    "    loss_clip_total = 0.\n",
+    "    loss_blurry_total = 0.\n",
+    "    loss_blurry_cont_total = 0.\n",
+    "    test_loss_clip_total = 0.\n",
+    "    \n",
+    "    loss_prior_total = 0.\n",
+    "    test_loss_prior_total = 0.\n",
+    "    \n",
+    "    loss_RR_total = 0.\n",
+    "    test_loss_RR_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",
+    "    # pre-load all batches for this epoch (it's MUCH faster to pre-load in bulk than to separate loading per batch)\n",
+    "    voxel_iters = {} # empty dict because diff subjects have differing # of voxels\n",
+    "    image_iters = torch.zeros(num_iterations_per_epoch, batch_size*len(subj_list), 3, 224, 224).float()\n",
+    "    annot_iters = {}\n",
+    "    perm_iters, betas_iters, select_iters = {}, {}, {}\n",
+    "    images_indexes = {}\n",
+    "    for s, train_dl in enumerate(train_dls):\n",
+    "        with torch.cuda.amp.autocast(dtype=data_type):\n",
+    "            iter = -1\n",
+    "            for behav0, past_behav0, future_behav0, old_behav0 in train_dl: \n",
+    "                # Load images to cpu from hdf5 (requires sorted indexing)\n",
+    "                image_idx = behav0[:,0,0].cpu().long().numpy()\n",
+    "\n",
+    "                image0, image_sorted_idx = np.unique(image_idx, return_index=True)                \n",
+    "                if len(image0) != len(image_idx): # hdf5 cant handle duplicate indexing\n",
+    "                    continue\n",
+    "                iter += 1\n",
+    "                image0 = torch.tensor(images[image0], dtype=data_type)\n",
+    "                image_iters[iter,s*batch_size:s*batch_size+batch_size] = image0\n",
+    "                images_indexes[f\"subj0{s}_iter{iter}\"] = image_sorted_idx\n",
+    "                \n",
+    "                # Load voxels for current batch, matching above indexing\n",
+    "                voxel_idx = behav0[:,0,5].cpu().long().numpy()\n",
+    "                voxel_sorted_idx = voxel_idx[image_sorted_idx]\n",
+    "                voxel0 = voxels[f'subj0{subj_list[s]}'][voxel_sorted_idx]\n",
+    "                voxel0 = torch.Tensor(voxel0).unsqueeze(1)\n",
+    "\n",
+    "                if epoch < int(mixup_pct * num_epochs):\n",
+    "                    voxel0, perm, betas, select = utils.mixco(voxel0)\n",
+    "                    perm_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = perm\n",
+    "                    betas_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = betas\n",
+    "                    select_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = select\n",
+    "\n",
+    "                voxel_iters[f\"subj0{subj_list[s]}_iter{iter}\"] = voxel0\n",
+    "\n",
+    "                if iter >= num_iterations_per_epoch-1:\n",
+    "                    break\n",
+    "\n",
+    "    # you now have voxel_iters and image_iters with num_iterations_per_epoch batches each\n",
+    "    for train_i in range(num_iterations_per_epoch):\n",
+    "        with torch.cuda.amp.autocast(dtype=data_type):\n",
+    "            optimizer.zero_grad()\n",
+    "            loss=0.\n",
+    "\n",
+    "            voxel_list = [voxel_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "            image = image_iters[train_i].detach()\n",
+    "            image = image.to(device)\n",
+    "\n",
+    "            if use_image_aug: \n",
+    "                image = img_augment(image)\n",
+    "\n",
+    "            clip_target = clip_img_embedder(image)\n",
+    "            assert not torch.any(torch.isnan(clip_target))\n",
+    "\n",
+    "            if epoch < int(mixup_pct * num_epochs):\n",
+    "                perm_list = [perm_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "                perm = torch.cat(perm_list, dim=0)\n",
+    "                betas_list = [betas_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "                betas = torch.cat(betas_list, dim=0)\n",
+    "                select_list = [select_iters[f\"subj0{s}_iter{train_i}\"].detach().to(device) for s in subj_list]\n",
+    "                select = torch.cat(select_list, dim=0)\n",
+    "\n",
+    "            voxel_ridge_list = [model.ridge(voxel_list[si],si) for si,s in enumerate(subj_list)]\n",
+    "            voxel_ridge = torch.cat(voxel_ridge_list, dim=0)\n",
+    "\n",
+    "            backbone, clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)\n",
+    "\n",
+    "            if clip_scale>0:\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 use_prior:\n",
+    "            #     loss_prior, prior_out = model.diffusion_prior(text_embed=backbone, image_embed=clip_target)\n",
+    "            #     loss_prior_total += loss_prior.item()\n",
+    "            #     loss_prior *= prior_scale\n",
+    "            #     loss += loss_prior\n",
+    "\n",
+    "            #     recon_cossim += nn.functional.cosine_similarity(prior_out, clip_target).mean().item()\n",
+    "            #     recon_mse += mse(prior_out, clip_target).item()\n",
+    "\n",
+    "            logits = model.RRClassifier(backbone.flatten(1))\n",
+    "            #print(logits.shape, torch.argmax(semantic_cluster_onehot[images_indexes[f\"subj0{s}_iter{train_i}\"]], dim=1).shape)\n",
+    "            #print(logits, torch.argmax(semantic_cluster_onehot[images_indexes[f\"subj0{s}_iter{train_i}\"]], dim=1))\n",
+    "            loss_RR = nn.functional.cross_entropy(logits, torch.argmax(semantic_cluster_onehot[images_indexes[f\"subj0{s}_iter{train_i}\"]], dim=1).to(logits.device))\n",
+    "            #print(\"backbone.shape\",backbone.shape, \"clip_voxels.shape\",clip_voxels.shape, \"blurry_image_enc_[0].shape\",blurry_image_enc_[0].shape, \"blurry_image_enc_[1].shape\",blurry_image_enc_[1].shape)\n",
+    "            #something \n",
+    "\n",
+    "            loss_RR_total += loss_RR.item()\n",
+    "            loss += loss_RR\n",
+    "\n",
+    "            if clip_scale>0:\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_clip_total += loss_clip.item()\n",
+    "                loss_clip *= clip_scale\n",
+    "                loss += loss_clip\n",
+    "\n",
+    "            if blurry_recon:     \n",
+    "                image_enc_pred, transformer_feats = blurry_image_enc_\n",
+    "\n",
+    "                image_enc = autoenc.encode(2*image-1).latent_dist.mode() * 0.18215\n",
+    "                loss_blurry = l1(image_enc_pred, image_enc)\n",
+    "                loss_blurry_total += loss_blurry.item()\n",
+    "\n",
+    "                if epoch < int(mixup_pct * num_epochs):\n",
+    "                    image_enc_shuf = image_enc[perm]\n",
+    "                    betas_shape = [-1] + [1]*(len(image_enc.shape)-1)\n",
+    "                    image_enc[select] = image_enc[select] * betas[select].reshape(*betas_shape) + \\\n",
+    "                        image_enc_shuf[select] * (1 - betas[select]).reshape(*betas_shape)\n",
+    "\n",
+    "                image_norm = (image - mean)/std\n",
+    "                image_aug = (blur_augs(image) - mean)/std\n",
+    "                _, cnx_embeds = cnx(image_norm)\n",
+    "                _, cnx_aug_embeds = cnx(image_aug)\n",
+    "\n",
+    "                cont_loss = utils.soft_cont_loss(\n",
+    "                    nn.functional.normalize(transformer_feats.reshape(-1, transformer_feats.shape[-1]), dim=-1),\n",
+    "                    nn.functional.normalize(cnx_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),\n",
+    "                    nn.functional.normalize(cnx_aug_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),\n",
+    "                    temp=0.2)\n",
+    "                loss_blurry_cont_total += cont_loss.item()\n",
+    "\n",
+    "                loss += (loss_blurry + 0.1*cont_loss) * blur_scale #/.18215\n",
+    "\n",
+    "            if clip_scale>0:\n",
+    "                # forward and backward top 1 accuracy        \n",
+    "                labels = torch.arange(len(clip_voxels_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).item()\n",
+    "                bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()\n",
+    "\n",
+    "            if blurry_recon:\n",
+    "                with torch.no_grad():\n",
+    "                    # only doing pixcorr eval on a subset of the samples per batch because its costly & slow to compute autoenc.decode()\n",
+    "                    random_samps = np.random.choice(np.arange(len(image)), size=len(image)//5, replace=False)\n",
+    "                    blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample/ 2 + 0.5).clamp(0,1)\n",
+    "                    pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)\n",
+    "                    blurry_pixcorr += pixcorr.item()\n",
+    "\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",
+    "            if lr_scheduler_type is not None:\n",
+    "                lr_scheduler.step()\n",
+    "\n",
+    "    model.eval()\n",
+    "    if local_rank==0:\n",
+    "        with torch.no_grad(), torch.cuda.amp.autocast(dtype=data_type): \n",
+    "            for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):  \n",
+    "                # all test samples should be loaded per batch such that test_i should never exceed 0\n",
+    "                assert len(behav) == num_test\n",
+    "\n",
+    "                ## Average same-image repeats ##\n",
+    "                if test_image is None:\n",
+    "                    voxel = voxels[f'subj0{subj}'][behav[:,0,5].cpu().long()].unsqueeze(1)\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 len(locs)==1:\n",
+    "                            locs = locs.repeat(3)\n",
+    "                        elif len(locs)==2:\n",
+    "                            locs = locs.repeat(2)[:3]\n",
+    "                        assert len(locs)==3\n",
+    "                        if test_image is None:\n",
+    "                            test_image = torch.Tensor(images[im][None])\n",
+    "                            test_voxel = voxel[locs][None]\n",
+    "                        else:\n",
+    "                            test_image = torch.vstack((test_image, torch.Tensor(images[im][None])))\n",
+    "                            test_voxel = torch.vstack((test_voxel, voxel[locs][None]))\n",
+    "\n",
+    "                loss=0.\n",
+    "                            \n",
+    "                test_indices = torch.arange(len(test_voxel))[:300]\n",
+    "                voxel = test_voxel[test_indices].to(device)\n",
+    "                image = test_image[test_indices].to(device)\n",
+    "                assert len(image) == 300\n",
+    "\n",
+    "                clip_target = clip_img_embedder(image.float())\n",
+    "\n",
+    "                for rep in range(3):\n",
+    "                    voxel_ridge = model.ridge(voxel[:,rep],0) # 0th index of subj_list\n",
+    "                    backbone0, clip_voxels0, blurry_image_enc_ = model.backbone(voxel_ridge)\n",
+    "\n",
+    "                    logits0 = model.RRClassifier(backbone0.flatten(1))\n",
+    "\n",
+    "                    if rep==0:\n",
+    "                        clip_voxels = clip_voxels0\n",
+    "                        backbone = backbone0\n",
+    "                        logits = logits0\n",
+    "                    else:\n",
+    "                        clip_voxels += clip_voxels0\n",
+    "                        backbone += backbone0\n",
+    "                        logits += logits0\n",
+    "                clip_voxels /= 3\n",
+    "                backbone /= 3\n",
+    "                logits /= 3\n",
+    "\n",
+    "                print(logits.shape, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).shape)\n",
+    "                RR_loss = nn.functional.cross_entropy(logits, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).to(logits.device))\n",
+    "                test_loss_RR_total += RR_loss.item()\n",
+    "                loss += RR_loss\n",
+    "\n",
+    "                if clip_scale>0:\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",
+    "                # for some evals, only doing a subset of the samples per batch because of computational cost\n",
+    "                random_samps = np.random.choice(np.arange(len(image)), size=len(image)//5, replace=False)\n",
+    "                \n",
+    "                # if use_prior:\n",
+    "                #     loss_prior, contaminated_prior_out = model.diffusion_prior(text_embed=backbone[random_samps], image_embed=clip_target[random_samps])\n",
+    "                #     test_loss_prior_total += loss_prior.item()\n",
+    "                #     loss_prior *= prior_scale\n",
+    "                #     loss += loss_prior\n",
+    "                    \n",
+    "                if clip_scale>0:\n",
+    "                    loss_clip = utils.soft_clip_loss(\n",
+    "                        clip_voxels_norm,\n",
+    "                        clip_target_norm,\n",
+    "                        temp=.006)\n",
+    "\n",
+    "                    test_loss_clip_total += loss_clip.item()\n",
+    "                    loss_clip = loss_clip * clip_scale\n",
+    "                    loss += loss_clip\n",
+    "\n",
+    "                if blurry_recon:\n",
+    "                    image_enc_pred, _ = blurry_image_enc_\n",
+    "                    blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample / 2 + 0.5).clamp(0,1)\n",
+    "                    pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)\n",
+    "                    test_blurry_pixcorr += pixcorr.item()\n",
+    "\n",
+    "                if clip_scale>0:\n",
+    "                    # forward and backward top 1 accuracy        \n",
+    "                    labels = torch.arange(len(clip_voxels_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).item()\n",
+    "                    test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1).item()\n",
+    "                \n",
+    "                utils.check_loss(loss)                \n",
+    "                test_losses.append(loss.item())\n",
+    "\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",
+    "                \"train/loss_blurry_cont_total\": loss_blurry_cont_total / (train_i + 1),\n",
+    "                \"test/loss_clip_total\": test_loss_clip_total / (test_i + 1),\n",
+    "                \"train/blurry_pixcorr\": blurry_pixcorr / (train_i + 1),\n",
+    "                \"test/blurry_pixcorr\": test_blurry_pixcorr / (test_i + 1),\n",
+    "                \"train/recon_cossim\": recon_cossim / (train_i + 1),\n",
+    "                \"test/recon_cossim\": test_recon_cossim / (test_i + 1),\n",
+    "                \"train/recon_mse\": recon_mse / (train_i + 1),\n",
+    "                \"test/recon_mse\": test_recon_mse / (test_i + 1),\n",
+    "                \"train/loss_prior\": loss_prior_total / (train_i + 1),\n",
+    "                \"test/loss_prior\": test_loss_prior_total / (test_i + 1),\n",
+    "                \"train/loss_RR\": loss_RR_total / (train_i + 1),\n",
+    "                \"test/loss_RR\": test_loss_RR_total / (test_i + 1),\n",
+    "                }\n",
+    "\n",
+    "            # if finished training, save jpg recons if they exist\n",
+    "            if (epoch == num_epochs-1) or (epoch % ckpt_interval == 0):\n",
+    "                if blurry_recon:    \n",
+    "                    image_enc = autoenc.encode(2*image[:4]-1).latent_dist.mode() * 0.18215\n",
+    "                    # transform blurry recon latents to images and plot it\n",
+    "                    fig, axes = plt.subplots(1, 8, figsize=(10, 4))\n",
+    "                    jj=-1\n",
+    "                    for j in [0,1,2,3]:\n",
+    "                        jj+=1\n",
+    "                        axes[jj].imshow(utils.torch_to_Image((autoenc.decode(image_enc[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))\n",
+    "                        axes[jj].axis('off')\n",
+    "                        jj+=1\n",
+    "                        axes[jj].imshow(utils.torch_to_Image((autoenc.decode(image_enc_pred[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))\n",
+    "                        axes[jj].axis('off')\n",
+    "\n",
+    "                    if wandb_log:\n",
+    "                        logs[f\"test/blur_recons\"] = wandb.Image(fig, caption=f\"epoch{epoch:03d}\")\n",
+    "                        plt.close()\n",
+    "                    else:\n",
+    "                        plt.show()\n",
+    "\n",
+    "            progress_bar.set_postfix(**logs)\n",
+    "\n",
+    "            if wandb_log: wandb.log(logs)\n",
+    "            \n",
+    "    # Save model checkpoint and reconstruct\n",
+    "    if (ckpt_saving) and (epoch % ckpt_interval == 0):\n",
+    "        save_ckpt(f'last')\n",
+    "\n",
+    "    # wait for other GPUs to catch up if needed\n",
+    "    accelerator.wait_for_everyone()\n",
+    "    torch.cuda.empty_cache()\n",
+    "\n",
+    "print(\"\\n===Finished!===\\n\")\n",
+    "if ckpt_saving:\n",
+    "    save_ckpt(f'last')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a7e81ae3-171f-40ad-a3e8-24bee4472325",
+   "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": "45d5bc17-0914-43a5-bb4e-d98f0ba238f0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import wandb\n",
+    "wandb.login()"
+   ]
+  }
+ ],
+ "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.11.9"
+  },
+  "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
+}

MindEyeV2/src/TrainCluster.py

@@ -0,0 +1,1078 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# # Import packages & functions
+
+# In[1]:
+
+
+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 matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from accelerate import Accelerator
+
+# SDXL unCLIP requires code from https://github.com/Stability-AI/generative-models/tree/main
+sys.path.append('generative_models/')
+import sgm
+from generative_models.sgm.modules.encoders.modules import FrozenOpenCLIPImageEmbedder # bigG embedder
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+
+
+# In[2]:
+
+
+### 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)  
+
+data_type = torch.float16 # change depending on your mixed_precision
+num_devices = torch.cuda.device_count()
+if num_devices==0: num_devices = 1
+
+# First use "accelerate config" in terminal and setup using deepspeed stage 2 with CPU offloading!
+accelerator = Accelerator(split_batches=False, mixed_precision="fp16")
+if utils.is_interactive(): # set batch size here if using interactive notebook instead of submitting job
+    global_batch_size = batch_size = 8
+else:
+    global_batch_size = os.environ["GLOBAL_BATCH_SIZE"]
+    batch_size = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices
+
+
+# In[3]:
+
+
+print("PID of this process =",os.getpid())
+device = accelerator.device
+print("device:",device)
+world_size = accelerator.state.num_processes
+distributed = not accelerator.state.distributed_type == 'NO'
+num_devices = torch.cuda.device_count()
+if num_devices==0 or not distributed: num_devices = 1
+num_workers = num_devices
+print(accelerator.state)
+
+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[4]:
+
+
+# if running this interactively, can specify jupyter_args here for argparser to use
+if utils.is_interactive():
+    model_name = "testing2"
+    print("model_name:", model_name)
+    
+    # global_batch_size and batch_size should already be defined in the 2nd cell block
+    jupyter_args = f"--data_path=/weka/proj-medarc/shared/mindeyev2_dataset \
+                    --cache_dir=/weka/proj-medarc/shared/cache \
+                    --model_name={model_name} \
+                    --no-multi_subject --subj=1 --batch_size={batch_size} --num_sessions=40 \
+                    --hidden_dim=1024 --clip_scale=1. \
+                    --no-blurry_recon --blur_scale=.5  \
+                    --use_prior --prior_scale=30 \
+                    --n_blocks=4 --max_lr=3e-5 --mixup_pct=.33 --num_epochs=150 --no-use_image_aug \
+                    --ckpt_interval=999 --no-ckpt_saving --wandb_log"
+    # --multisubject_ckpt=../train_logs/multisubject_subj01_1024_24bs_nolow
+
+    print(jupyter_args)
+    jupyter_args = jupyter_args.split()
+    
+    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[5]:
+
+
+parser = argparse.ArgumentParser(description="Model Training Configuration")
+parser.add_argument(
+    "--model_name", type=str, default="testing2",
+    help="name of model, used for ckpt saving and wandb logging (if enabled)",
+)
+parser.add_argument(
+    "--data_path", type=str, default=os.getcwd(),
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--cache_dir", type=str, default=os.getcwd(),
+    help="Path to where misc. files downloaded from huggingface are stored. Defaults to current src directory.",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,3,4,5,6,7,8],
+    help="Validate on which subject?",
+)
+parser.add_argument(
+    "--multisubject_ckpt", type=str, default=None,
+    help="Path to pre-trained multisubject model to finetune a single subject from. multisubject must be False.",
+)
+parser.add_argument(
+    "--num_sessions", type=int, default=1,
+    help="Number of training sessions to include",
+)
+parser.add_argument(
+    "--use_prior",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to train diffusion prior (True) or just rely on retrieval part of the pipeline (False)",
+)
+parser.add_argument(
+    "--batch_size", type=int, default=16,
+    help="Batch size can be increased by 10x if only training retreival submodule and not diffusion prior",
+)
+parser.add_argument(
+    "--wandb_log",action=argparse.BooleanOptionalAction,default=False,
+    help="whether to log to wandb",
+)
+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(
+    "--blur_scale",type=float,default=.5,
+    help="multiply loss from blurry recons by this number",
+)
+parser.add_argument(
+    "--clip_scale",type=float,default=1.,
+    help="multiply contrastive loss by this number",
+)
+parser.add_argument(
+    "--prior_scale",type=float,default=30,
+    help="multiply diffusion prior loss by this",
+)
+parser.add_argument(
+    "--use_image_aug",action=argparse.BooleanOptionalAction,default=False,
+    help="whether to use image augmentation",
+)
+parser.add_argument(
+    "--num_epochs",type=int,default=150,
+    help="number of epochs of training",
+)
+parser.add_argument(
+    "--multi_subject",action=argparse.BooleanOptionalAction,default=False,
+)
+parser.add_argument(
+    "--new_test",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--n_blocks",type=int,default=4,
+)
+parser.add_argument(
+    "--hidden_dim",type=int,default=1024,
+)
+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-5,
+)
+
+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)
+    
+# seed all random functions
+utils.seed_everything(seed)
+
+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 or blurry_recon:
+    import kornia
+    from kornia.augmentation.container import AugmentationSequential
+if use_image_aug:
+    img_augment = AugmentationSequential(
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1, p=0.3),
+        same_on_batch=False,
+        data_keys=["input"],
+    )
+    
+if multi_subject:
+    subj_list = np.arange(1,9)
+    subj_list = subj_list[subj_list != subj]
+else:
+    subj_list = [subj]
+
+print("subj_list", subj_list, "num_sessions", num_sessions)
+
+
+# In[6]:
+
+
+max_lr
+
+
+# # Prep data, models, and dataloaders
+
+# ### Creating wds dataloader, preload betas and all 73k possible images
+
+# In[7]:
+
+
+def my_split_by_node(urls): return urls
+num_voxels_list = []
+
+if multi_subject:
+    nsessions_allsubj=np.array([40, 40, 32, 30, 40, 32, 40, 30])
+    num_samples_per_epoch = (750*40) // num_devices 
+else:
+    num_samples_per_epoch = (750*num_sessions) // num_devices 
+
+print("dividing batch size by subj_list, which will then be concatenated across subj during training...") 
+batch_size = batch_size // len(subj_list)
+
+num_iterations_per_epoch = num_samples_per_epoch // (batch_size*len(subj_list))
+
+print("batch_size =", batch_size, "num_iterations_per_epoch =",num_iterations_per_epoch, "num_samples_per_epoch =",num_samples_per_epoch)
+
+
+# In[8]:
+
+
+train_data = {}
+train_dl = {}
+num_voxels = {}
+voxels = {}
+for s in subj_list:
+    print(f"Training with {num_sessions} sessions")
+    if multi_subject:
+        train_url = f"{data_path}/wds/subj0{s}/train/" + "{0.." + f"{nsessions_allsubj[s-1]-1}" + "}.tar"
+    else:
+        train_url = f"{data_path}/wds/subj0{s}/train/" + "{0.." + f"{num_sessions-1}" + "}.tar"
+    print(train_url)
+    
+    train_data[f'subj0{s}'] = wds.WebDataset(train_url,resampled=True,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[f'subj0{s}'] = torch.utils.data.DataLoader(train_data[f'subj0{s}'], batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+    f = h5py.File(f'{data_path}/betas_all_subj0{s}_fp32_renorm.hdf5', 'r')
+    betas = f['betas'][:]
+    betas = torch.Tensor(betas).to("cpu").to(data_type)
+    num_voxels_list.append(betas[0].shape[-1])
+    num_voxels[f'subj0{s}'] = betas[0].shape[-1]
+    voxels[f'subj0{s}'] = betas
+    print(f"num_voxels for subj0{s}: {num_voxels[f'subj0{s}']}")
+
+print("Loaded all subj train dls and betas!\n")
+
+# Validate only on one subject
+if multi_subject: 
+    subj = subj_list[0] # cant validate on the actual held out person so picking first in subj_list
+if not new_test: # using old test set from before full dataset released (used in original MindEye paper)
+    if subj==3:
+        num_test=2113
+    elif subj==4:
+        num_test=1985
+    elif subj==6:
+        num_test=2113
+    elif subj==8:
+        num_test=1985
+    else:
+        num_test=2770
+    test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+elif new_test: # using larger test set from after full dataset released
+    if subj==3:
+        num_test=2371
+    elif subj==4:
+        num_test=2188
+    elif subj==6:
+        num_test=2371
+    elif subj==8:
+        num_test=2188
+    else:
+        num_test=3000
+    test_url = f"{data_path}/wds/subj0{subj}/new_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=num_test, shuffle=False, drop_last=True, pin_memory=True)
+print(f"Loaded test dl for subj{subj}!\n")
+
+
+# In[9]:
+
+
+# Load 73k NSD images
+f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+images = f['images']
+print("Loaded all 73k possible NSD images to cpu!", images.shape)
+
+
+# ## Load models
+
+# ### CLIP image embeddings  model
+
+# In[10]:
+
+
+clip_img_embedder = FrozenOpenCLIPImageEmbedder(
+    arch="ViT-bigG-14",
+    version="laion2b_s39b_b160k",
+    output_tokens=True,
+    only_tokens=True,
+)
+clip_img_embedder.to(device)
+
+clip_seq_dim = 256
+clip_emb_dim = 1664
+
+
+# ### SD VAE
+
+# In[11]:
+
+
+if blurry_recon:
+    from diffusers import AutoencoderKL    
+    autoenc = AutoencoderKL(
+        down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D'],
+        up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D'],
+        block_out_channels=[128, 256, 512, 512],
+        layers_per_block=2,
+        sample_size=256,
+    )
+    ckpt = torch.load(f'{cache_dir}/sd_image_var_autoenc.pth')
+    autoenc.load_state_dict(ckpt)
+    
+    autoenc.eval()
+    autoenc.requires_grad_(False)
+    autoenc.to(device)
+    utils.count_params(autoenc)
+    
+    from autoencoder.convnext import ConvnextXL
+    cnx = ConvnextXL(f'{cache_dir}/convnext_xlarge_alpha0.75_fullckpt.pth')
+    cnx.requires_grad_(False)
+    cnx.eval()
+    cnx.to(device)
+    
+    mean = torch.tensor([0.485, 0.456, 0.406]).to(device).reshape(1,3,1,1)
+    std = torch.tensor([0.228, 0.224, 0.225]).to(device).reshape(1,3,1,1)
+    
+    blur_augs = AugmentationSequential(
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.8),
+        kornia.augmentation.RandomGrayscale(p=0.1),
+        kornia.augmentation.RandomSolarize(p=0.1),
+        kornia.augmentation.RandomResizedCrop((224,224), scale=(.9,.9), ratio=(1,1), p=1.0),
+        data_keys=["input"],
+    )
+
+
+# ### MindEye modules
+
+# In[12]:
+
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+model
+
+
+# In[13]:
+
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer to enable regularization
+    def __init__(self, input_sizes, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linears = torch.nn.ModuleList([
+                torch.nn.Linear(input_size, out_features) for input_size in input_sizes
+            ])
+    def forward(self, x, subj_idx):
+        out = self.linears[subj_idx](x[:,0]).unsqueeze(1)
+        return out
+        
+class IndividRidgeRegression(torch.nn.Module):
+    def __init__(self, input_size, out_features):
+        super(IndividRidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linear = torch.nn.Linear(input_size, out_features)
+    def forward(self, x):
+        out = self.linear(x)
+        return out
+    
+model.ridge = RidgeRegression(num_voxels_list, out_features=hidden_dim)
+utils.count_params(model.ridge)
+utils.count_params(model)
+
+# test on subject 1 with fake data
+b = torch.randn((2,1,num_voxels_list[0]))
+print(b.shape, model.ridge(b,0).shape)
+
+
+# In[14]:
+
+
+from models import BrainNetwork
+model.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, seq_len=1, n_blocks=n_blocks,
+                          clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, 
+                          blurry_recon=blurry_recon, clip_scale=clip_scale)
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+# test that the model works on some fake data
+b = torch.randn((2,1,hidden_dim))
+print("b.shape",b.shape)
+
+backbone_, clip_, blur_ = model.backbone(b)
+print(backbone_.shape, clip_.shape, blur_[0].shape, blur_[1].shape)
+
+
+# ### Load semantic clusters
+
+# In[15]:
+
+
+path_semantic_names = "/weka/proj-medarc/shared/mindeyev2_dataset/semantic_cluster_names.npy"
+path_semantic_cluster = "/weka/proj-medarc/shared/mindeyev2_dataset/COCO_73k_semantic_cluster.npy"
+semantic_cluster_names = np.load(path_semantic_names)
+semantic_cluster = np.load(path_semantic_cluster)
+possible_semantic_clusters = np.unique(semantic_cluster)
+
+# one-hot encode semantic clusters
+# move possible_semantic_clusters to numbers and create a dictionary
+semantic_cluster_dict = {cluster: i for i, cluster in enumerate(possible_semantic_clusters)}
+semantic_cluster_onehot = torch.zeros((len(semantic_cluster), len(possible_semantic_clusters)))
+for i, cluster in enumerate(semantic_cluster):
+    semantic_cluster_onehot[i, semantic_cluster_dict[cluster]] = 1
+
+
+print("semantic_cluster_onehot.shape", semantic_cluster_onehot.shape)
+
+num_seman_clusters = len(np.unique(semantic_cluster))
+print("num_seman_clusters", num_seman_clusters)
+
+
+# ### Adding the ridge regression to the class
+
+# In[16]:
+
+
+# if use_prior:
+#     from models import *
+
+#     # setup diffusion prior network
+#     out_dim = clip_emb_dim
+#     depth = 6
+#     dim_head = 52
+#     heads = clip_emb_dim//52 # heads * dim_head = clip_emb_dim
+#     timesteps = 100
+
+#     prior_network = PriorNetwork(
+#             dim=out_dim,
+#             depth=depth,
+#             dim_head=dim_head,
+#             heads=heads,
+#             causal=False,
+#             num_tokens = clip_seq_dim,
+#             learned_query_mode="pos_emb"
+#         )
+
+#     model.diffusion_prior = BrainDiffusionPrior(
+#         net=prior_network,
+#         image_embed_dim=out_dim,
+#         condition_on_text_encodings=False,
+#         timesteps=timesteps,
+#         cond_drop_prob=0.2,
+#         image_embed_scale=None,
+#     )
+    
+#     utils.count_params(model.diffusion_prior)
+#     utils.count_params(model)
+
+model.RRClassifier = IndividRidgeRegression(clip_emb_dim*clip_seq_dim, out_features=num_seman_clusters)
+utils.count_params(model.RRClassifier)
+utils.count_params(model)
+
+
+# ### Setup optimizer / lr / ckpt saving
+
+# 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},
+]
+# if use_prior:
+#     opt_grouped_parameters.extend([
+#         {'params': [p for n, p in model.diffusion_prior.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},
+#         {'params': [p for n, p in model.diffusion_prior.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
+#     ])
+opt_grouped_parameters.extend([
+    {'params': [p for n, p in model.RRClassifier.named_parameters()], 'weight_decay': 1e-2},
+])
+
+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_iterations_per_epoch)),
+        last_epoch=-1
+    )
+elif lr_scheduler_type == 'cycle':
+    total_steps=int(np.floor(num_epochs*num_iterations_per_epoch))
+    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'
+    if accelerator.is_main_process:
+        unwrapped_model = accelerator.unwrap_model(model)
+        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)
+    print(f"\n---saved {outdir}/{tag} ckpt!---\n")
+
+def load_ckpt(tag,load_lr=True,load_optimizer=True,load_epoch=True,strict=True,outdir=outdir,multisubj_loading=False): 
+    print(f"\n---loading {outdir}/{tag}.pth ckpt---\n")
+    checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+    state_dict = checkpoint['model_state_dict']
+    if multisubj_loading: # remove incompatible ridge layer that will otherwise error
+        state_dict.pop('ridge.linears.0.weight',None)
+    model.load_state_dict(state_dict, strict=strict)
+    if load_epoch:
+        globals()["epoch"] = checkpoint['epoch']
+        print("Epoch",epoch)
+    if load_optimizer:
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    if load_lr:
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    del checkpoint
+
+print("\nDone with model preparations!")
+num_params = utils.count_params(model)
+
+
+# # Weights and Biases
+
+# In[18]:
+
+
+if local_rank==0 and wandb_log: # only use main process for wandb logging
+    import wandb
+    wandb_project = 'mindeye_semantic_cluster'
+    print(f"wandb {wandb_project} run {model_name}")
+    # need to configure wandb beforehand in terminal with "wandb init"!
+    wandb_config = {
+      "model_name": model_name,
+      "global_batch_size": global_batch_size,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "num_sessions": num_sessions,
+      "num_params": num_params,
+      "clip_scale": clip_scale,
+      "prior_scale": prior_scale,
+      "blur_scale": blur_scale,
+      "use_image_aug": use_image_aug,
+      "max_lr": max_lr,
+      "mixup_pct": mixup_pct,
+      "num_samples_per_epoch": num_samples_per_epoch,
+      "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)
+    print("wandb_id:",model_name)
+    wandb.login(host='https://stability.wandb.io')
+    wandb.init(
+        id=model_name,
+        project=wandb_project,
+        name=model_name,
+        config=wandb_config,
+        resume="allow",
+    )
+else:
+    wandb_log = False
+
+
+# # Main
+
+# In[19]:
+
+
+epoch = 0
+losses, test_losses, lrs = [], [], []
+best_test_loss = 1e9
+torch.cuda.empty_cache()
+
+
+# In[20]:
+
+
+# load multisubject stage1 ckpt if set
+if multisubject_ckpt is not None:
+    load_ckpt("last",outdir=multisubject_ckpt,load_lr=False,load_optimizer=False,load_epoch=False,strict=False,multisubj_loading=True)
+
+
+# In[21]:
+
+
+train_dls = [train_dl[f'subj0{s}'] for s in subj_list]
+
+model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot = accelerator.prepare(model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot)
+# leaving out test_dl since we will only have local_rank 0 device do evals
+
+
+# In[22]:
+
+
+print(num_iterations_per_epoch)
+
+
+# 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()
+soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))
+
+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.
+    
+    recon_cossim = 0.
+    test_recon_cossim = 0.
+    recon_mse = 0.
+    test_recon_mse = 0.
+
+    loss_clip_total = 0.
+    loss_blurry_total = 0.
+    loss_blurry_cont_total = 0.
+    test_loss_clip_total = 0.
+    
+    loss_prior_total = 0.
+    test_loss_prior_total = 0.
+    
+    loss_RR_total = 0.
+    test_loss_RR_total = 0.
+
+    blurry_pixcorr = 0.
+    test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1
+
+    # pre-load all batches for this epoch (it's MUCH faster to pre-load in bulk than to separate loading per batch)
+    voxel_iters = {} # empty dict because diff subjects have differing # of voxels
+    image_iters = torch.zeros(num_iterations_per_epoch, batch_size*len(subj_list), 3, 224, 224).float()
+    annot_iters = {}
+    perm_iters, betas_iters, select_iters = {}, {}, {}
+    images_indexes = {}
+    for s, train_dl in enumerate(train_dls):
+        with torch.cuda.amp.autocast(dtype=data_type):
+            iter = -1
+            for behav0, past_behav0, future_behav0, old_behav0 in train_dl: 
+                # Load images to cpu from hdf5 (requires sorted indexing)
+                image_idx = behav0[:,0,0].cpu().long().numpy()
+
+                image0, image_sorted_idx = np.unique(image_idx, return_index=True)                
+                if len(image0) != len(image_idx): # hdf5 cant handle duplicate indexing
+                    continue
+                iter += 1
+                image0 = torch.tensor(images[image0], dtype=data_type)
+                image_iters[iter,s*batch_size:s*batch_size+batch_size] = image0
+                images_indexes[f"subj0{s}_iter{iter}"] = image_sorted_idx
+                
+                # Load voxels for current batch, matching above indexing
+                voxel_idx = behav0[:,0,5].cpu().long().numpy()
+                voxel_sorted_idx = voxel_idx[image_sorted_idx]
+                voxel0 = voxels[f'subj0{subj_list[s]}'][voxel_sorted_idx]
+                voxel0 = torch.Tensor(voxel0).unsqueeze(1)
+
+                if epoch < int(mixup_pct * num_epochs):
+                    voxel0, perm, betas, select = utils.mixco(voxel0)
+                    perm_iters[f"subj0{subj_list[s]}_iter{iter}"] = perm
+                    betas_iters[f"subj0{subj_list[s]}_iter{iter}"] = betas
+                    select_iters[f"subj0{subj_list[s]}_iter{iter}"] = select
+
+                voxel_iters[f"subj0{subj_list[s]}_iter{iter}"] = voxel0
+
+                if iter >= num_iterations_per_epoch-1:
+                    break
+
+    # you now have voxel_iters and image_iters with num_iterations_per_epoch batches each
+    for train_i in range(num_iterations_per_epoch):
+        with torch.cuda.amp.autocast(dtype=data_type):
+            optimizer.zero_grad()
+            loss=0.
+
+            voxel_list = [voxel_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list]
+            image = image_iters[train_i].detach()
+            image = image.to(device)
+
+            if use_image_aug: 
+                image = img_augment(image)
+
+            clip_target = clip_img_embedder(image)
+            assert not torch.any(torch.isnan(clip_target))
+
+            if epoch < int(mixup_pct * num_epochs):
+                perm_list = [perm_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list]
+                perm = torch.cat(perm_list, dim=0)
+                betas_list = [betas_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list]
+                betas = torch.cat(betas_list, dim=0)
+                select_list = [select_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list]
+                select = torch.cat(select_list, dim=0)
+
+            voxel_ridge_list = [model.ridge(voxel_list[si],si) for si,s in enumerate(subj_list)]
+            voxel_ridge = torch.cat(voxel_ridge_list, dim=0)
+
+            backbone, clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+
+            if clip_scale>0:
+                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 use_prior:
+            #     loss_prior, prior_out = model.diffusion_prior(text_embed=backbone, image_embed=clip_target)
+            #     loss_prior_total += loss_prior.item()
+            #     loss_prior *= prior_scale
+            #     loss += loss_prior
+
+            #     recon_cossim += nn.functional.cosine_similarity(prior_out, clip_target).mean().item()
+            #     recon_mse += mse(prior_out, clip_target).item()
+
+            logits = model.RRClassifier(backbone.flatten(1))
+            #print(logits.shape, torch.argmax(semantic_cluster_onehot[images_indexes[f"subj0{s}_iter{train_i}"]], dim=1).shape)
+            #print(logits, torch.argmax(semantic_cluster_onehot[images_indexes[f"subj0{s}_iter{train_i}"]], dim=1))
+            loss_RR = nn.functional.cross_entropy(logits, torch.argmax(semantic_cluster_onehot[images_indexes[f"subj0{s}_iter{train_i}"]], dim=1).to(logits.device))
+            #print("backbone.shape",backbone.shape, "clip_voxels.shape",clip_voxels.shape, "blurry_image_enc_[0].shape",blurry_image_enc_[0].shape, "blurry_image_enc_[1].shape",blurry_image_enc_[1].shape)
+            #something 
+
+            loss_RR_total += loss_RR.item()
+            loss += loss_RR
+
+            if clip_scale>0:
+                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:     
+                image_enc_pred, transformer_feats = blurry_image_enc_
+
+                image_enc = autoenc.encode(2*image-1).latent_dist.mode() * 0.18215
+                loss_blurry = l1(image_enc_pred, image_enc)
+                loss_blurry_total += loss_blurry.item()
+
+                if epoch < int(mixup_pct * num_epochs):
+                    image_enc_shuf = image_enc[perm]
+                    betas_shape = [-1] + [1]*(len(image_enc.shape)-1)
+                    image_enc[select] = image_enc[select] * betas[select].reshape(*betas_shape) + \
+                        image_enc_shuf[select] * (1 - betas[select]).reshape(*betas_shape)
+
+                image_norm = (image - mean)/std
+                image_aug = (blur_augs(image) - mean)/std
+                _, cnx_embeds = cnx(image_norm)
+                _, cnx_aug_embeds = cnx(image_aug)
+
+                cont_loss = utils.soft_cont_loss(
+                    nn.functional.normalize(transformer_feats.reshape(-1, transformer_feats.shape[-1]), dim=-1),
+                    nn.functional.normalize(cnx_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),
+                    nn.functional.normalize(cnx_aug_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),
+                    temp=0.2)
+                loss_blurry_cont_total += cont_loss.item()
+
+                loss += (loss_blurry + 0.1*cont_loss) * blur_scale #/.18215
+
+            if clip_scale>0:
+                # forward and backward top 1 accuracy        
+                labels = torch.arange(len(clip_voxels_norm)).to(clip_voxels_norm.device) 
+                fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item()
+                bwd_percent_correct += utils.topk(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(image)), size=len(image)//5, replace=False)
+                    blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample/ 2 + 0.5).clamp(0,1)
+                    pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)
+                    blurry_pixcorr += pixcorr.item()
+
+            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[f'subj0{subj}'][behav[:,0,5].cpu().long()].unsqueeze(1)
+                    
+                    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 len(locs)==1:
+                            locs = locs.repeat(3)
+                        elif len(locs)==2:
+                            locs = locs.repeat(2)[:3]
+                        assert len(locs)==3
+                        if test_image is None:
+                            test_image = torch.Tensor(images[im][None])
+                            test_voxel = voxel[locs][None]
+                        else:
+                            test_image = torch.vstack((test_image, torch.Tensor(images[im][None])))
+                            test_voxel = torch.vstack((test_voxel, voxel[locs][None]))
+
+                loss=0.
+                            
+                test_indices = torch.arange(len(test_voxel))[:300]
+                voxel = test_voxel[test_indices].to(device)
+                image = test_image[test_indices].to(device)
+                assert len(image) == 300
+
+                clip_target = clip_img_embedder(image.float())
+
+                for rep in range(3):
+                    voxel_ridge = model.ridge(voxel[:,rep],0) # 0th index of subj_list
+                    backbone0, clip_voxels0, blurry_image_enc_ = model.backbone(voxel_ridge)
+
+                    logits0 = model.RRClassifier(backbone0.flatten(1))
+
+                    if rep==0:
+                        clip_voxels = clip_voxels0
+                        backbone = backbone0
+                        logits = logits0
+                    else:
+                        clip_voxels += clip_voxels0
+                        backbone += backbone0
+                        logits += logits0
+                clip_voxels /= 3
+                backbone /= 3
+                logits /= 3
+
+                print(logits.shape, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).shape)
+                RR_loss = nn.functional.cross_entropy(logits, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).to(logits.device))
+                test_loss_RR_total += RR_loss.item()
+                loss += RR_loss
+
+                if clip_scale>0:
+                    clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+                    clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+                
+                # for some evals, only doing a subset of the samples per batch because of computational cost
+                random_samps = np.random.choice(np.arange(len(image)), size=len(image)//5, replace=False)
+                
+                # if use_prior:
+                #     loss_prior, contaminated_prior_out = model.diffusion_prior(text_embed=backbone[random_samps], image_embed=clip_target[random_samps])
+                #     test_loss_prior_total += loss_prior.item()
+                #     loss_prior *= prior_scale
+                #     loss += loss_prior
+                    
+                if clip_scale>0:
+                    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:
+                    image_enc_pred, _ = blurry_image_enc_
+                    blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample / 2 + 0.5).clamp(0,1)
+                    pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)
+                    test_blurry_pixcorr += pixcorr.item()
+
+                if clip_scale>0:
+                    # forward and backward top 1 accuracy        
+                    labels = torch.arange(len(clip_voxels_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())
+
+            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),
+                "train/loss_blurry_cont_total": loss_blurry_cont_total / (train_i + 1),
+                "test/loss_clip_total": test_loss_clip_total / (test_i + 1),
+                "train/blurry_pixcorr": blurry_pixcorr / (train_i + 1),
+                "test/blurry_pixcorr": test_blurry_pixcorr / (test_i + 1),
+                "train/recon_cossim": recon_cossim / (train_i + 1),
+                "test/recon_cossim": test_recon_cossim / (test_i + 1),
+                "train/recon_mse": recon_mse / (train_i + 1),
+                "test/recon_mse": test_recon_mse / (test_i + 1),
+                "train/loss_prior": loss_prior_total / (train_i + 1),
+                "test/loss_prior": test_loss_prior_total / (test_i + 1),
+                "train/loss_RR": loss_RR_total / (train_i + 1),
+                "test/loss_RR": test_loss_RR_total / (test_i + 1),
+                }
+
+            # if finished training, save jpg recons if they exist
+            if (epoch == num_epochs-1) or (epoch % ckpt_interval == 0):
+                if blurry_recon:    
+                    image_enc = autoenc.encode(2*image[:4]-1).latent_dist.mode() * 0.18215
+                    # 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(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(image_enc_pred[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                        axes[jj].axis('off')
+
+                    if wandb_log:
+                        logs[f"test/blur_recons"] = wandb.Image(fig, caption=f"epoch{epoch:03d}")
+                        plt.close()
+                    else:
+                        plt.show()
+
+            progress_bar.set_postfix(**logs)
+
+            if wandb_log: wandb.log(logs)
+            
+    # Save model checkpoint and reconstruct
+    if (ckpt_saving) and (epoch % ckpt_interval == 0):
+        save_ckpt(f'last')
+
+    # wait for other GPUs to catch up if needed
+    accelerator.wait_for_everyone()
+    torch.cuda.empty_cache()
+
+print("\n===Finished!===\n")
+if ckpt_saving:
+    save_ckpt(f'last')
+
+
+# In[ ]:
+
+
+plt.plot(losses)
+plt.show()
+plt.plot(test_losses)
+plt.show()
+
+
+# In[ ]:
+
+
+import wandb
+wandb.login()
+

MindEyeV2/src/accel.slurm

@@ -0,0 +1,43 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --qos=normal
+#SBATCH --partition=p5
+#SBATCH --job-name=eyeRR
+#SBATCH --nodes=1           
+#SBATCH --gres=gpu:1
+#SBATCH --time=48:00:00          # total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err        # first create a "slurms" folder in current directory to store logs
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=medarc
+#SBATCH --no-requeue
+#SBATCH --exclusive
+
+# Make sure you activate your fmri environment created from src/setup.sh
+cd /weka/proj-fmri/ckadirt/MindEyeV2/src
+source /admin/home-ckadirt/fmri/bin/activate
+
+# The following line converts your jupyter notebook into a python script runnable with Slurm
+jupyter nbconvert TrainCluster-Copy2.ipynb --to python
+
+export NUM_GPUS=1  # Set to equal gres=gpu:#!
+export BATCH_SIZE=21 # 21 for multisubject / 24 for singlesubject (orig. paper used 42 for multisubject / 24 for singlesubject)
+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=1 #$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo WORLD_SIZE=${COUNT_NODE}
+
+# multisubject pretraining
+model_name="rr-test-top"
+echo model_name=${model_name}
+# accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --mixed_precision=fp16 TrainCluster.py --data_path=/weka/proj-medarc/shared/mindeyev2_dataset --cache_dir=/weka/proj-medarc/shared/cache --model_name=${model_name} --multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log
+python TrainCluster-Copy2.py --data_path=/weka/proj-medarc/shared/mindeyev2_dataset --cache_dir=/weka/proj-medarc/shared/cache --model_name=${model_name} --no-multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-5 --mixup_pct=.33 --num_epochs=150 --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log
+# singlesubject finetuning
+#model_name="finetuned_subj01_40sess"
+#echo model_name=${model_name}
+#accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --mixed_precision=fp16 Train.py --data_path=/weka/proj-fmri/shared/mindeyev2_dataset --cache_dir=/weka/proj-fmri/shared/cache --model_name=${model_name} --no-multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --use_prior --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log --multisubject_ckpt=../train_logs/multisubject_excludingsubj01_40sess

MindEyeV2/src/accelSM.slurm

@@ -0,0 +1,42 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --qos=normal
+#SBATCH --partition=p5
+#SBATCH --job-name=sm_p_msub_0.33
+#SBATCH --nodes=1           
+#SBATCH --gres=gpu:1
+#SBATCH --time=48:00:00          # total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err        # first create a "slurms" folder in current directory to store logs
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=medarc
+#SBATCH --no-requeue
+
+# Make sure you activate your fmri environment created from src/setup.sh
+cd /weka/proj-fmri/ckadirt/MindEyeV2/src
+source /admin/home-ckadirt/fmri/bin/activate
+
+# The following line converts your jupyter notebook into a python script runnable with Slurm
+jupyter nbconvert train_semantic_cluster.ipynb --to python
+
+export NUM_GPUS=1  # Set to equal gres=gpu:#!
+export BATCH_SIZE=21 # 21 for multisubject / 24 for singlesubject (orig. paper used 42 for multisubject / 24 for singlesubject)
+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=1 #$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo WORLD_SIZE=${COUNT_NODE}
+
+# multisubject pretraining
+model_name="semantic_cluster_1.2_average_after_wd-2_prior_multi"
+echo model_name=${model_name}
+# accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --mixed_precision=fp16 TrainCluster.py --data_path=/weka/proj-medarc/shared/mindeyev2_dataset --cache_dir=/weka/proj-medarc/shared/cache --model_name=${model_name} --multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log
+python Untitled1.py --data_path=/weka/proj-medarc/shared/mindeyev2_dataset --cache_dir=/weka/proj-medarc/shared/cache --model_name=${model_name} --multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-5 --mixup_pct=.33 --num_epochs=150 --prior_scale=30 --clip_scale=1 --no-blurry_recon --use_prior --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log
+# singlesubject finetuning
+#model_name="finetuned_subj01_40sess"
+#echo model_name=${model_name}
+#accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --mixed_precision=fp16 Train.py --data_path=/weka/proj-fmri/shared/mindeyev2_dataset --cache_dir=/weka/proj-fmri/shared/cache --model_name=${model_name} --no-multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --use_prior --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log --multisubject_ckpt=../train_logs/multisubject_excludingsubj01_40sess

MindEyeV2/src/accel_recon_imagery.slurm

@@ -0,0 +1,58 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --partition=p5
+#SBATCH --job-name=me2_imagery
+#SBATCH --ntasks-per-node=1
+#SBATCH --nodes=1
+#SBATCH --gres=gpu:1
+#SBATCH --gpus-per-task=1       # Set to equal gres=gpu:#
+#SBATCH --time=72:00:00         # Total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err        # Ensure "slurms" directory exists for logs
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=medarc
+#SBATCH --no-requeue
+#SBATCH --qos=idle
+#SBATCH --dependency=afterok:534007
+
+source ~/.bashrc
+source /admin/home-ckadirt/fmri/bin/activate
+cd /weka/proj-fmri/ckadirt/MindEyeV2/src/
+
+
+# Set up networking environment variables to avoid port conflicts
+export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000))
+export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+export MASTER_ADDR=$(echo "$HOSTNAMES" | head -n 1)
+echo "MASTER_ADDR=${MASTER_ADDR}"
+echo "MASTER_PORT=${MASTER_PORT}"
+
+export MODULEPATH=/opt/amazon/modules/modulefiles:$MODULEPATH
+
+jupyter nbconvert enhanced_recon_inference.ipynb --to python
+jupyter nbconvert recon_inference.ipynb --to python
+
+# jupyter nbconvert final_evaluations_mi_multi_icnn.ipynb --to python
+
+
+for subj in 1; do
+    for num_sessions in 2 3 5 10 20 40; do
+        model_name="final_subj0${subj}_pretrained_${num_sessions}sess_24bs"
+        echo ${model_name}
+        echo "new_sessions"
+
+        for mode in "imagery" "vision"; do #
+                
+            python recon_inference.py \
+                    --data_path=/weka/proj-medarc/shared/mindeyev2_dataset \
+                    --cache_dir=/weka/proj-medarc/shared/mindeyev2_dataset \
+                    --model_name="${model_name}" --subj=1 \
+                    --hidden_dim=4096 --n_blocks=4 --new_test \
+                    --mode="${mode}" \
+                    --imagery_data_path=/weka/proj-medarc/shared/umn-imagery 
+
+            python enhanced_recon_inference.py \
+                    --model_name="${model_name}" --subj=1 --mode="${mode}"
+
+            done
+        done
+    done

MindEyeV2/src/accel_recon_shared1000.slurm

@@ -0,0 +1,53 @@
+#!/bin/bash
+#SBATCH --account=fmri
+#SBATCH --partition=p5
+#SBATCH --job-name=me2_shared1000
+#SBATCH --ntasks-per-node=1
+#SBATCH --nodes=1
+#SBATCH --gres=gpu:1
+#SBATCH --gpus-per-task=1       # Set to equal gres=gpu:#
+#SBATCH --time=72:00:00         # Total run time limit (HH:MM:SS)
+#SBATCH -e slurms/%j.err        # Ensure "slurms" directory exists for logs
+#SBATCH -o slurms/%j.out
+#SBATCH --comment=medarc
+#SBATCH --no-requeue
+#SBATCH --qos=normal
+
+source ~/.bashrc
+source /admin/home-ckadirt/fmri/bin/activate
+cd /weka/proj-fmri/ckadirt/MindEyeV2/src/
+
+
+# Set up networking environment variables to avoid port conflicts
+export MASTER_PORT=$((RANDOM % (19000 - 11000 + 1) + 11000))
+export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+export MASTER_ADDR=$(echo "$HOSTNAMES" | head -n 1)
+echo "MASTER_ADDR=${MASTER_ADDR}"
+echo "MASTER_PORT=${MASTER_PORT}"
+
+export MODULEPATH=/opt/amazon/modules/modulefiles:$MODULEPATH
+
+jupyter nbconvert enhanced_recon_inference_old.ipynb --to python
+jupyter nbconvert recon_inference_old.ipynb --to python
+
+# jupyter nbconvert final_evaluations_mi_multi_icnn.ipynb --to python
+
+
+for subj in 1; do
+    for num_sessions in 3; do
+        model_name="final_subj0${subj}_pretrained_${num_sessions}sess_24bs"
+        echo ${model_name}
+        echo "new_sessions"
+
+                
+        python recon_inference_old.py \
+                --data_path=/weka/proj-medarc/shared/mindeyev2_dataset \
+                --cache_dir=/weka/proj-medarc/shared/mindeyev2_dataset \
+                --model_name="${model_name}" --subj=1 \
+                --hidden_dim=4096 --n_blocks=4 --new_test
+
+        python enhanced_recon_inference_old.py \
+                --model_name="${model_name}" --subj=1
+
+        done
+    done

MindEyeV2/src/autoencoder/convnext.py

@@ -0,0 +1,239 @@
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.layers import trunc_normal_, DropPath
+
+
+class Block(nn.Module):
+    r""" ConvNeXt Block. There are two equivalent implementations:
+    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
+    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
+    We use (2) as we find it slightly faster in PyTorch
+    
+    Args:
+        dim (int): Number of input channels.
+        drop_path (float): Stochastic depth rate. Default: 0.0
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+    """
+
+    def __init__(self, dim, drop_path=0.0, layer_scale_init_value=1e-6):
+        super().__init__()
+        self.dwconv = nn.Conv2d(
+            dim, dim, kernel_size=7, padding=3, groups=dim
+        )  # depthwise conv
+        self.norm = LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(
+            dim, 4 * dim
+        )  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
+            if layer_scale_init_value > 0
+            else None
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+
+        x = input + self.drop_path(x)
+        return x
+
+
+class ConvNeXt(nn.Module):
+    r""" ConvNeXt
+        A PyTorch impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+    Args:
+        in_chans (int): Number of input image channels. Default: 3
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(
+        self,
+        in_chans=3,
+        num_classes=1000,
+        depths=[3, 3, 9, 3],
+        dims=[96, 192, 384, 768],
+        drop_path_rate=0.1,
+        layer_scale_init_value=0.0,
+        head_init_scale=1.0,
+    ):
+        super().__init__()
+
+        self.downsample_layers = (
+            nn.ModuleList()
+        )  # stem and 3 intermediate downsampling conv layers
+        stem = nn.Sequential(
+            nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
+            LayerNorm(dims[0], eps=1e-6, data_format="channels_first"),
+        )
+        self.downsample_layers.append(stem)
+        for i in range(3):
+            downsample_layer = nn.Sequential(
+                LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
+                nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = (
+            nn.ModuleList()
+        )  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(4):
+            stage = nn.Sequential(
+                *[
+                    Block(
+                        dim=dims[i],
+                        drop_path=dp_rates[cur + j],
+                        layer_scale_init_value=layer_scale_init_value,
+                    )
+                    for j in range(depths[i])
+                ]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+        self.norm = nn.LayerNorm(dims[-1], eps=1e-6)  # final norm layer
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(m.weight, std=0.02)
+            nn.init.constant_(m.bias, 0)
+
+    def forward_features(self, x):
+        for i in range(4):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+        x = x.flatten(2, 3).permute(0, 2, 1)
+        x = self.norm(x)
+        x_avg = x.mean(dim=1)
+        return x, x_avg
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        return x
+
+
+class LayerNorm(nn.Module):
+    r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. 
+    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with 
+    shape (batch_size, height, width, channels) while channels_first corresponds to inputs 
+    with shape (batch_size, channels, height, width).
+    """
+
+    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        if self.data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError
+        self.normalized_shape = (normalized_shape,)
+
+    def forward(self, x):
+        if self.data_format == "channels_last":
+            return F.layer_norm(
+                x, self.normalized_shape, self.weight, self.bias, self.eps
+            )
+        elif self.data_format == "channels_first":
+            u = x.mean(1, keepdim=True)
+            s = (x - u).pow(2).mean(1, keepdim=True)
+            x = (x - u) / torch.sqrt(s + self.eps)
+            x = self.weight[:, None, None] * x + self.bias[:, None, None]
+            return x
+
+
+def convnext_tiny(**kwargs):
+    model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
+    return model, 768
+
+
+def convnext_small(**kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], **kwargs)
+    return model, 768
+
+
+def convnext_base(**kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)
+    return model, 1024
+
+
+def convnext_large(**kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536], **kwargs)
+    return model, 1536
+
+
+def convnext_xlarge(**kwargs):
+    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048], **kwargs)
+    return model, 2048
+
+def MLP(mlp, embedding, norm_layer):
+    mlp_spec = f"{embedding}-{mlp}"
+    layers = []
+    f = list(map(int, mlp_spec.split("-")))
+    for i in range(len(f) - 2):
+        layers.append(nn.Linear(f[i], f[i + 1]))
+        if norm_layer == "batch_norm":
+            layers.append(nn.BatchNorm1d(f[i + 1]))
+        elif norm_layer == "layer_norm":
+            layers.append(nn.LayerNorm(f[i + 1]))
+        layers.append(nn.ReLU(True))
+    layers.append(nn.Linear(f[-2], f[-1], bias=False))
+    return nn.Sequential(*layers)
+
+class ConvnextXL(torch.nn.Module):
+    def __init__(self, path):
+        super().__init__()
+        self.backbone, _ = convnext_xlarge()
+        self.maps_projector = MLP('512-512-512', 2048, 'layer_norm')
+        self.init_weights(path)
+
+    def init_weights(self, path):
+        sd = torch.load(path, map_location='cpu')['model']
+        for k in list(sd.keys()):
+            if 'module' in k:
+                sd[k.replace('module.', '')] = sd[k]
+                del sd[k]
+        try:
+            self.load_state_dict(sd, strict=False)
+        except:
+            pass
+    
+    def forward(self, x):
+        p, _ = self.backbone(x)  # b, hw, 2048
+        return p, self.maps_projector(p)
+
+
+if __name__== "__main__":
+    m = ConvnextXL('../train_logs/models/convnext_xlarge_alpha0.75_fullckpt.pth')
+    preds = m(torch.randn(2, 3, 512, 512))

MindEyeV2/src/enhanced_recon_inference.ipynb

@@ -0,0 +1,891 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "6b18f6a3-cc4f-437e-9756-c99fc6a5fad4",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "device: cuda\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/admin/home-ckadirt/mindeye/lib/python3.11/site-packages/accelerate/accelerator.py:457: FutureWarning: Passing the following arguments to `Accelerator` is deprecated and will be removed in version 1.0 of Accelerate: dict_keys(['split_batches']). Please pass an `accelerate.DataLoaderConfiguration` instead: \n",
+      "dataloader_config = DataLoaderConfiguration(split_batches=False)\n",
+      "  warnings.warn(\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 string\n",
+    "import h5py\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torchvision import transforms\n",
+    "from accelerate import Accelerator, DeepSpeedPlugin\n",
+    "\n",
+    "# SDXL unCLIP requires code from https://github.com/Stability-AI/generative-models/tree/main\n",
+    "sys.path.append('generative_models/')\n",
+    "import sgm\n",
+    "from generative_models.sgm.modules.encoders.modules import FrozenOpenCLIPImageEmbedder, FrozenCLIPEmbedder, FrozenOpenCLIPEmbedder2\n",
+    "from generative_models.sgm.models.diffusion import DiffusionEngine\n",
+    "from generative_models.sgm.util import append_dims\n",
+    "from omegaconf import OmegaConf\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",
+    "from models import *\n",
+    "\n",
+    "accelerator = Accelerator(split_batches=False, mixed_precision=\"fp16\")\n",
+    "device = accelerator.device\n",
+    "print(\"device:\",device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "20cdb696-1d6e-4b73-951b-b0cd1dda219a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "model_name: final_subj01_pretrained_375trials_24bs\n",
+      "--model_name=final_subj01_pretrained_375trials_24bs --subj=1 --mode=imagery\n"
+     ]
+    }
+   ],
+   "source": [
+    "# if running this interactively, can specify jupyter_args here for argparser to use\n",
+    "if utils.is_interactive():\n",
+    "    model_name = \"final_subj01_pretrained_375trials_24bs\"\n",
+    "    print(\"model_name:\", model_name)\n",
+    "    mode = \"imagery\"\n",
+    "    # global_batch_size and batch_size should already be defined in the above cells\n",
+    "    # other variables can be specified in the following string:\n",
+    "    jupyter_args = f\"--model_name={model_name} --subj=1 --mode={mode}\"\n",
+    "    print(jupyter_args)\n",
+    "    jupyter_args = jupyter_args.split()\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": 3,
+   "id": "4b31d7c0-f5bd-4a19-a8be-7a3a165d79b6",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                                                                                                             | 0/18 [00:00<?, ?it/s]/admin/home-ckadirt/mindeye/lib/python3.11/site-packages/torchvision/transforms/functional.py:1603: UserWarning: The default value of the antialias parameter of all the resizing transforms (Resize(), RandomResizedCrop(), etc.) will change from None to True in v0.17, in order to be consistent across the PIL and Tensor backends. To suppress this warning, directly pass antialias=True (recommended, future default), antialias=None (current default, which means False for Tensors and True for PIL), or antialias=False (only works on Tensors - PIL will still use antialiasing). This also applies if you are using the inference transforms from the models weights: update the call to weights.transforms(antialias=True).\n",
+      "  warnings.warn(\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████��███████████████| 18/18 [00:03<00:00,  4.95it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "final_subj01_pretrained_375trials_24bs\n",
+      "torch.Size([18, 3, 425, 425]) torch.Size([18, 10, 3, 768, 768]) torch.Size([18, 10, 256, 1664]) torch.Size([18, 10, 3, 768, 768]) (18, 10)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "parser = argparse.ArgumentParser(description=\"Model Training Configuration\")\n",
+    "parser.add_argument(\n",
+    "    \"--model_name\", type=str, default=\"testing\",\n",
+    "    help=\"will load ckpt for model found in ../train_logs/model_name\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--subj\",type=int, default=1, choices=[1,2,3,4,5,6,7,8],\n",
+    "    help=\"Evaluate on which subject?\",\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--seed\",type=int,default=42,\n",
+    ")\n",
+    "parser.add_argument(\n",
+    "    \"--mode\",type=str,default=\"vision\",choices=[\"vision\",\"imagery\",\"shared1000\"],\n",
+    "    help=\"Evaluate on which modality?\",\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",
+    "# seed all random functions\n",
+    "utils.seed_everything(seed)\n",
+    "\n",
+    "# make output directory\n",
+    "os.makedirs(\"evals\",exist_ok=True)\n",
+    "os.makedirs(f\"evals/{model_name}\",exist_ok=True)\n",
+    "\n",
+    "# Some of these files are downloadable from huggingface: https://huggingface.co/datasets/pscotti/mindeyev2/tree/main/evals\n",
+    "# The others are obtained from running recon_inference.ipynb first with your desired model\n",
+    "all_images = torch.load(f\"evals/all_images.pt\")\n",
+    "rep_all_recons = torch.load(f\"evals/{model_name}/{model_name}_all_recons_{mode}.pt\")\n",
+    "rep_all_clipvoxels = torch.load(f\"evals/{model_name}/{model_name}_all_clipvoxels_{mode}.pt\")\n",
+    "rep_all_blurryrecons = torch.load(f\"evals/{model_name}/{model_name}_all_blurryrecons_{mode}.pt\")\n",
+    "rep_all_predcaptions = torch.load(f\"evals/{model_name}/{model_name}_all_predcaptions_{mode}.pt\")\n",
+    "\n",
+    "\n",
+    "resized_recons = None\n",
+    "resized_blurryrecons = None\n",
+    "for i in tqdm(range(len(all_images))):\n",
+    "    resized_recons_ = transforms.Resize((768,768))(rep_all_recons[i]).float()\n",
+    "    resized_blurryrecons_ = transforms.Resize((768,768))(rep_all_blurryrecons[i]).float()\n",
+    "    if resized_recons is None:\n",
+    "        resized_recons = resized_recons_.unsqueeze(0)\n",
+    "        resized_blurryrecons = resized_blurryrecons_.unsqueeze(0)\n",
+    "    else:\n",
+    "        resized_recons = torch.cat([resized_recons, resized_recons_.unsqueeze(0)],0)\n",
+    "        resized_blurryrecons = torch.cat([resized_blurryrecons, resized_blurryrecons_.unsqueeze(0)],0)\n",
+    "\n",
+    "rep_all_blurryrecons = resized_blurryrecons\n",
+    "rep_all_recons = resized_recons\n",
+    "\n",
+    "# reshape the predcaptions from 180 to 18,10\n",
+    "from einops import rearrange\n",
+    "if rep_all_predcaptions.shape[0] != resized_recons.shape[0]:\n",
+    "    rep_all_predcaptions = rearrange(rep_all_predcaptions, '(c r) -> r c', r=resized_recons.shape[0], c=resized_recons.shape[1])\n",
+    "\n",
+    "print(model_name)\n",
+    "print(all_images.shape, rep_all_recons.shape, rep_all_clipvoxels.shape, rep_all_blurryrecons.shape, rep_all_predcaptions.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "24bdd667-0862-4561-b432-9fa7543df863",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "SpatialTransformer: Found context dims [2048] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [2048, 2048] now.\n",
+      "SpatialTransformer: Found context dims [2048] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [2048, 2048] now.\n",
+      "SpatialTransformer: Found context dims [2048] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048] now.\n",
+      "SpatialTransformer: Found context dims [2048] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048] now.\n",
+      "SpatialTransformer: Found context dims [2048] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048] now.\n",
+      "SpatialTransformer: Found context dims [2048] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048] now.\n",
+      "SpatialTransformer: Found context dims [2048] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048] now.\n",
+      "SpatialTransformer: Found context dims [2048] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048] now.\n",
+      "SpatialTransformer: Found context dims [2048] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [2048, 2048] now.\n",
+      "SpatialTransformer: Found context dims [2048] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [2048, 2048] now.\n",
+      "SpatialTransformer: Found context dims [2048] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [2048, 2048] now.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Initialized embedder #0: FrozenCLIPEmbedder with 123060480 params. Trainable: False\n",
+      "Initialized embedder #1: FrozenOpenCLIPEmbedder2 with 694659841 params. Trainable: False\n",
+      "Initialized embedder #2: ConcatTimestepEmbedderND with 0 params. Trainable: False\n",
+      "Initialized embedder #3: ConcatTimestepEmbedderND with 0 params. Trainable: False\n",
+      "Initialized embedder #4: ConcatTimestepEmbedderND with 0 params. Trainable: False\n",
+      "Restored from /weka/proj-medarc/shared/mindeyev2_dataset/zavychromaxl_v30.safetensors with 1 missing and 1 unexpected keys\n",
+      "Missing Keys: ['denoiser.sigmas']\n",
+      "Unexpected Keys: ['conditioner.embedders.0.transformer.text_model.embeddings.position_ids']\n",
+      "crossattn torch.Size([1, 77, 2048])\n",
+      "vector_suffix torch.Size([1, 1536])\n",
+      "---\n",
+      "crossattn_uc torch.Size([1, 77, 2048])\n",
+      "vector_uc torch.Size([1, 2816])\n"
+     ]
+    }
+   ],
+   "source": [
+    "config = OmegaConf.load(\"generative_models/configs/unclip6.yaml\")\n",
+    "config = OmegaConf.to_container(config, resolve=True)\n",
+    "unclip_params = config[\"model\"][\"params\"]\n",
+    "sampler_config = unclip_params[\"sampler_config\"]\n",
+    "sampler_config['params']['num_steps'] = 38\n",
+    "config = OmegaConf.load(\"generative_models/configs/inference/sd_xl_base.yaml\")\n",
+    "config = OmegaConf.to_container(config, resolve=True)\n",
+    "refiner_params = config[\"model\"][\"params\"]\n",
+    "\n",
+    "network_config = refiner_params[\"network_config\"]\n",
+    "denoiser_config = refiner_params[\"denoiser_config\"]\n",
+    "first_stage_config = refiner_params[\"first_stage_config\"]\n",
+    "conditioner_config = refiner_params[\"conditioner_config\"]\n",
+    "scale_factor = refiner_params[\"scale_factor\"]\n",
+    "disable_first_stage_autocast = refiner_params[\"disable_first_stage_autocast\"]\n",
+    "\n",
+    "# base_ckpt_path = '/weka/robin/projects/stable-research/checkpoints/sd_xl_base_1.0.safetensors'\n",
+    "base_ckpt_path = '/weka/proj-medarc/shared/mindeyev2_dataset/zavychromaxl_v30.safetensors'\n",
+    "base_engine = DiffusionEngine(network_config=network_config,\n",
+    "                       denoiser_config=denoiser_config,\n",
+    "                       first_stage_config=first_stage_config,\n",
+    "                       conditioner_config=conditioner_config,\n",
+    "                       sampler_config=sampler_config, # using the one defined by the unclip\n",
+    "                       scale_factor=scale_factor,\n",
+    "                       disable_first_stage_autocast=disable_first_stage_autocast,\n",
+    "                       ckpt_path=base_ckpt_path)\n",
+    "base_engine.eval().requires_grad_(False)\n",
+    "base_engine.to(device)\n",
+    "\n",
+    "base_text_embedder1 = FrozenCLIPEmbedder(\n",
+    "    layer=conditioner_config['params']['emb_models'][0]['params']['layer'],\n",
+    "    layer_idx=conditioner_config['params']['emb_models'][0]['params']['layer_idx'],\n",
+    ")\n",
+    "base_text_embedder1.to(device)\n",
+    "\n",
+    "base_text_embedder2 = FrozenOpenCLIPEmbedder2(\n",
+    "    arch=conditioner_config['params']['emb_models'][1]['params']['arch'],\n",
+    "    version=conditioner_config['params']['emb_models'][1]['params']['version'],\n",
+    "    freeze=conditioner_config['params']['emb_models'][1]['params']['freeze'],\n",
+    "    layer=conditioner_config['params']['emb_models'][1]['params']['layer'],\n",
+    "    always_return_pooled=conditioner_config['params']['emb_models'][1]['params']['always_return_pooled'],\n",
+    "    legacy=conditioner_config['params']['emb_models'][1]['params']['legacy'],\n",
+    ")\n",
+    "base_text_embedder2.to(device)\n",
+    "\n",
+    "batch={\"txt\": \"\",\n",
+    "      \"original_size_as_tuple\": torch.ones(1, 2).to(device) * 768,\n",
+    "      \"crop_coords_top_left\": torch.zeros(1, 2).to(device),\n",
+    "      \"target_size_as_tuple\": torch.ones(1, 2).to(device) * 1024}\n",
+    "out = base_engine.conditioner(batch)\n",
+    "crossattn = out[\"crossattn\"].to(device)\n",
+    "vector_suffix = out[\"vector\"][:,-1536:].to(device)\n",
+    "print(\"crossattn\", crossattn.shape)\n",
+    "print(\"vector_suffix\", vector_suffix.shape)\n",
+    "print(\"---\")\n",
+    "\n",
+    "batch_uc={\"txt\": \"painting, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, deformed, ugly, blurry, bad anatomy, bad proportions, extra limbs, cloned face, skinny, glitchy, double torso, extra arms, extra hands, mangled fingers, missing lips, ugly face, distorted face, extra legs, anime\",\n",
+    "      \"original_size_as_tuple\": torch.ones(1, 2).to(device) * 768,\n",
+    "      \"crop_coords_top_left\": torch.zeros(1, 2).to(device),\n",
+    "      \"target_size_as_tuple\": torch.ones(1, 2).to(device) * 1024}\n",
+    "out = base_engine.conditioner(batch_uc)\n",
+    "crossattn_uc = out[\"crossattn\"].to(device)\n",
+    "vector_uc = out[\"vector\"].to(device)\n",
+    "print(\"crossattn_uc\", crossattn_uc.shape)\n",
+    "print(\"vector_uc\", vector_uc.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "07f437d1-9b8e-4b13-85ad-d45062a5ce09",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "if utils.is_interactive(): plotting=False\n",
+    "plotting=False\n",
+    "num_samples = 1 # PS: I tried increasing this to 16 and picking highest cosine similarity like we did in MindEye1, it didnt seem to increase eval performance!\n",
+    "img2img_timepoint = 13 # 9 # higher number means more reliance on prompt, less reliance on matching the conditioning image\n",
+    "base_engine.sampler.guider.scale = 5 # 5 # cfg\n",
+    "def denoiser(x, sigma, c): return base_engine.denoiser(base_engine.model, x, sigma, c)\n",
+    "\n",
+    "if plotting or num_samples>1:\n",
+    "    clip_img_embedder = FrozenOpenCLIPImageEmbedder(\n",
+    "        arch=\"ViT-bigG-14\",\n",
+    "        version=\"laion2b_s39b_b160k\",\n",
+    "        output_tokens=True,\n",
+    "        only_tokens=True,\n",
+    "    )\n",
+    "    clip_img_embedder.to(device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "939e1cbb-5836-48c2-87d8-3e493e950011",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a3e32a1951204ca79d0ab0117da33406",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/18 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "130c932f43fd477f98c3f1787f881b92",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/admin/home-ckadirt/mindeye/lib/python3.11/site-packages/torch/utils/checkpoint.py:429: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
+      "  warnings.warn(\n",
+      "/admin/home-ckadirt/mindeye/lib/python3.11/site-packages/torch/utils/checkpoint.py:61: UserWarning: None of the inputs have requires_grad=True. Gradients will be None\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "383921dd7d074af9b6330eb68f112c27",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5c4eb7e9e2594cfbbfac8035b762fa5a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "20d79a05a5a643ee9063b5239a549290",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "45d9979238d04909889ff11f458adb65",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0cbf925ddc334a9ea11c8ce310d155ec",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b58c54fcd2924f7db8e9f1d2ea425fdf",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "03790a6b73bc4b549fb35c6ecde64a06",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "30b4e55f835943f29c1d8d3778fc6b1b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "073db9cbbae84e76b9e1acabd2dec807",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "026d94daa9404792b3952efe47b693d1",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bb82c763bf54416aadb250c2dfb5c4aa",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a419df891da9468682f97224281d040e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ab3fbf9138c94f0d8b9ecea922208ced",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "efcc94ccbb944073b43db9d2ca41502f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c129efc8496e45b7971051efab9195b2",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "52da5ff9d7894bf8a18523d14cbabcdf",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f45c31dc95a84e6ab3a37f893c8e5da7",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "all_enhancedrecons torch.Size([10, 3, 256, 256])\n",
+      "final_enhancedrecons torch.Size([18, 10, 3, 256, 256])\n",
+      "saved evals/final_subj01_pretrained_375trials_24bs/final_subj01_pretrained_375trials_24bs_all_enhancedrecons_imagery.pt\n"
+     ]
+    }
+   ],
+   "source": [
+    "from tqdm.auto import tqdm\n",
+    "\n",
+    "final_enhancedrecons = None\n",
+    "\n",
+    "for rep in tqdm(range(len(rep_all_recons))):\n",
+    "    all_enhancedrecons = None\n",
+    "    all_recons = rep_all_recons[rep]\n",
+    "    all_clipvoxels = rep_all_clipvoxels[rep]\n",
+    "    all_blurryrecons = rep_all_blurryrecons[rep]\n",
+    "    all_predcaptions = rep_all_predcaptions[rep]\n",
+    "\n",
+    "    for img_idx in tqdm(range(len(all_recons))):\n",
+    "        with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.float16), base_engine.ema_scope():\n",
+    "            base_engine.sampler.num_steps = 25\n",
+    "            \n",
+    "            image = all_recons[[img_idx]]\n",
+    "            if plotting:\n",
+    "                print(\"blur pixcorr:\",utils.pixcorr(all_blurryrecons[[img_idx]].float(), all_images[[img_idx]].float()))\n",
+    "                print(\"blur cossim:\",nn.functional.cosine_similarity(clip_img_embedder(utils.resize(all_blurryrecons[[img_idx]].float(),256).to(device)).flatten(1), \n",
+    "                                                            clip_img_embedder(utils.resize(all_images[[img_idx]].float(),224).to(device)).flatten(1)))\n",
+    "\n",
+    "                print(\"recon pixcorr:\",utils.pixcorr(image,all_images[[img_idx]].float()))\n",
+    "                print(\"recon cossim:\",nn.functional.cosine_similarity(clip_img_embedder(utils.resize(image,224).to(device)).flatten(1), \n",
+    "                                                            clip_img_embedder(utils.resize(all_images[[img_idx]].float(),224).to(device)).flatten(1)))\n",
+    "            \n",
+    "            image = image.to(device)\n",
+    "            prompt = all_predcaptions[[img_idx]][0]\n",
+    "            # prompt = \"\"\n",
+    "            if plotting: \n",
+    "                print(\"prompt:\",prompt)\n",
+    "                plt.imshow(transforms.ToPILImage()(all_blurryrecons[img_idx].float()))\n",
+    "                plt.show()\n",
+    "                plt.imshow(transforms.ToPILImage()(all_recons[img_idx].float()))\n",
+    "                plt.show()\n",
+    "                plt.imshow(transforms.ToPILImage()(image[0]))\n",
+    "                plt.show()\n",
+    "\n",
+    "            # z = torch.randn(num_samples,4,96,96).to(device)\n",
+    "            assert image.shape[-1]==768\n",
+    "            z = base_engine.encode_first_stage(image*2-1).repeat(num_samples,1,1,1)\n",
+    "\n",
+    "            openai_clip_text = base_text_embedder1(prompt)\n",
+    "            clip_text_tokenized, clip_text_emb  = base_text_embedder2(prompt)\n",
+    "            clip_text_emb = torch.hstack((clip_text_emb, vector_suffix))\n",
+    "            clip_text_tokenized = torch.cat((openai_clip_text, clip_text_tokenized),dim=-1)\n",
+    "            c = {\"crossattn\": clip_text_tokenized.repeat(num_samples,1,1), \"vector\": clip_text_emb.repeat(num_samples,1)}\n",
+    "            uc = {\"crossattn\": crossattn_uc.repeat(num_samples,1,1), \"vector\": vector_uc.repeat(num_samples,1)}\n",
+    "\n",
+    "            noise = torch.randn_like(z)\n",
+    "            sigmas = base_engine.sampler.discretization(base_engine.sampler.num_steps).to(device)\n",
+    "            init_z = (z + noise * append_dims(sigmas[-img2img_timepoint], z.ndim)) / torch.sqrt(1.0 + sigmas[0] ** 2.0)\n",
+    "            sigmas = sigmas[-img2img_timepoint:].repeat(num_samples,1)\n",
+    "\n",
+    "            base_engine.sampler.num_steps = sigmas.shape[-1] - 1\n",
+    "            noised_z, _, _, _, c, uc = base_engine.sampler.prepare_sampling_loop(init_z, cond=c, uc=uc, \n",
+    "                                                                num_steps=base_engine.sampler.num_steps)\n",
+    "            for timestep in range(base_engine.sampler.num_steps):\n",
+    "                noised_z = base_engine.sampler.sampler_step(sigmas[:,timestep],\n",
+    "                                                            sigmas[:,timestep+1],\n",
+    "                                                            denoiser, noised_z, cond=c, uc=uc, gamma=0)\n",
+    "            samples_z_base = noised_z\n",
+    "            samples_x = base_engine.decode_first_stage(samples_z_base)\n",
+    "            samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)\n",
+    "\n",
+    "            # find best sample\n",
+    "            if plotting==False and num_samples==1:\n",
+    "                samples = samples[0]\n",
+    "            else:\n",
+    "                sample_cossim = nn.functional.cosine_similarity(clip_img_embedder(utils.resize(samples,224).to(device)).flatten(1), \n",
+    "                                    clip_img_embedder(utils.resize(all_images[[img_idx]].float(),224).to(device)).flatten(1))\n",
+    "                which_sample = torch.argmax(sample_cossim)\n",
+    "                best_cossim = torch.max(sample_cossim)\n",
+    "\n",
+    "                if plotting:\n",
+    "                    print(\"samples\", samples.shape)\n",
+    "                    for n in range(num_samples):\n",
+    "                        recon = transforms.ToPILImage()(samples[n])\n",
+    "                        plt.imshow(recon)\n",
+    "                        plt.show()\n",
+    "                        if (n==which_sample).item(): print(\"CHOSEN ABOVE\")\n",
+    "                        print(\"upsampled pixcorr:\",utils.pixcorr(samples[[n]].cpu(),all_images[[img_idx]].float()))\n",
+    "                        print(\"upsampled cossim:\",nn.functional.cosine_similarity(clip_img_embedder(utils.resize(samples[[n]],224).to(device)).flatten(1), \n",
+    "                                                            clip_img_embedder(utils.resize(all_images[[img_idx]].float(),224).to(device)).flatten(1)))\n",
+    "                    err # dont want to do entire for loop with plotting=True\n",
+    "\n",
+    "                samples = samples[which_sample]\n",
+    "\n",
+    "            samples = samples.cpu()[None]\n",
+    "            if all_enhancedrecons is None:\n",
+    "                all_enhancedrecons = samples\n",
+    "            else:\n",
+    "                all_enhancedrecons = torch.vstack((all_enhancedrecons, samples))\n",
+    "\n",
+    "    all_enhancedrecons = transforms.Resize((256,256))(all_enhancedrecons).float()\n",
+    "    print(\"all_enhancedrecons\", all_enhancedrecons.shape)\n",
+    "    if final_enhancedrecons is None:\n",
+    "        final_enhancedrecons = all_enhancedrecons.unsqueeze(0)\n",
+    "    else:\n",
+    "        final_enhancedrecons = torch.vstack((final_enhancedrecons, all_enhancedrecons.unsqueeze(0)))\n",
+    "            \n",
+    "\n",
+    "# all_enhancedrecons = transforms.Resize((256,256))(all_enhancedrecons).float()\n",
+    "print(\"final_enhancedrecons\", final_enhancedrecons.shape)\n",
+    "torch.save(final_enhancedrecons,f\"evals/{model_name}/{model_name}_all_enhancedrecons_{mode}.pt\")\n",
+    "print(f\"saved evals/{model_name}/{model_name}_all_enhancedrecons_{mode}.pt\")\n",
+    "\n",
+    "if not utils.is_interactive():\n",
+    "    sys.exit(0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "b9a22a3c-e1a2-4453-beb9-101f6c477486",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# plt.imshow(transforms.ToPILImage()(final_enhancedrecons[6,5].float()))\n",
+    "# plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "mindeye",
+   "language": "python",
+   "name": "mindeye"
+  },
+  "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.11.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

MindEyeV2/src/recon_inference.py

@@ -0,0 +1,580 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+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 matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from accelerate import Accelerator
+
+# SDXL unCLIP requires code from https://github.com/Stability-AI/generative-models/tree/main
+sys.path.append('generative_models/')
+import sgm
+from generative_models.sgm.modules.encoders.modules import FrozenOpenCLIPImageEmbedder, FrozenOpenCLIPEmbedder2
+from generative_models.sgm.models.diffusion import DiffusionEngine
+from generative_models.sgm.util import append_dims
+from omegaconf import OmegaConf
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+from models import *
+
+accelerator = Accelerator(split_batches=False, mixed_precision="fp16")
+device = accelerator.device
+print("device:",device)
+
+
+# In[2]:
+
+
+# if running this interactively, can specify jupyter_args here for argparser to use
+if utils.is_interactive():
+    model_name = "final_subj01_pretrained_375trials_24bs"
+    print("model_name:", model_name)
+
+    # other variables can be specified in the following string:
+    jupyter_args = f"--data_path=/weka/proj-medarc/shared/mindeyev2_dataset \
+                    --cache_dir=/weka/proj-medarc/shared/mindeyev2_dataset \
+                    --model_name={model_name} --subj=1 \
+                    --hidden_dim=4096 --n_blocks=4 --new_test \
+                    --mode imagery \
+                    --imagery_data_path=/weka/proj-medarc/shared/umn-imagery "
+    print(jupyter_args)
+    jupyter_args = jupyter_args.split()
+    
+    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[3]:
+
+
+parser = argparse.ArgumentParser(description="Model Training Configuration")
+parser.add_argument(
+    "--model_name", type=str, default="testing",
+    help="will load ckpt for model found in ../train_logs/model_name",
+)
+parser.add_argument(
+    "--data_path", type=str, default=os.getcwd(),
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--cache_dir", type=str, default=os.getcwd(),
+    help="Path to where misc. files downloaded from huggingface are stored. Defaults to current src directory.",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,3,4,5,6,7,8],
+    help="Validate on which subject?",
+)
+parser.add_argument(
+    "--blurry_recon",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--n_blocks",type=int,default=4,
+)
+parser.add_argument(
+    "--hidden_dim",type=int,default=2048,
+)
+parser.add_argument(
+    "--new_test",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--seed",type=int,default=42,
+)
+parser.add_argument(
+    "--mode",type=str,default="vision",choices=["vision","imagery","shared1000"],
+)
+parser.add_argument(
+    "--imagery_data_path", type=str, default=os.getcwd(),
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--gen_rep",type=int,default=10,
+)
+
+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)
+    
+# seed all random functions
+utils.seed_everything(seed)
+
+# make output directory
+os.makedirs("evals",exist_ok=True)
+os.makedirs(f"evals/{model_name}",exist_ok=True)
+
+
+# In[4]:
+
+
+if mode == "synthetic":
+    voxels, all_images = utils.load_nsd_synthetic(subject=subj, average=False, nest=True, data_root = imagery_data_path)
+elif subj > 8:
+    _, _, voxels, all_images = utils.load_imageryrf(subject=subj-8, mode=mode, stimtype="object", average=False, nest=True, split=True, data_root = imagery_data_path)
+elif mode == "shared1000":
+    x_train, valid_nsd_ids_train, x_test, test_nsd_ids = utils.load_nsd(subject=subj, data_path=data_path)
+    voxels = torch.mean(x_test, dim=1, keepdim=True)
+    print(f"Loaded subj {subj} test betas! {voxels.shape}")
+    f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+    images = f['images']
+
+    all_images = torch.zeros((len(test_nsd_ids), 3, 224, 224))
+    for i, idx in enumerate(test_nsd_ids):
+        all_images[i] =  torch.from_numpy(images[idx])
+    del images, f
+    print(f"Filtered down to only the {len(test_nsd_ids)} test images for subject {subj}!")
+else:
+    voxels, all_images = utils.load_nsd_mental_imagery(subject=subj, 
+                                                       mode=mode, 
+                                                       stimtype="all", 
+                                                       average=True, 
+                                                       nest=False,
+                                                       data_root = imagery_data_path)
+print(voxels.shape)
+num_voxels = voxels.shape[-1]
+plotting = False
+
+
+# In[5]:
+
+
+# voxels = {}
+# # Load hdf5 data for betas
+# f = h5py.File(f'{data_path}/betas_all_subj0{subj}_fp32_renorm.hdf5', 'r')
+# betas = f['betas'][:]
+# betas = torch.Tensor(betas).to("cpu")
+# num_voxels = betas[0].shape[-1]
+# voxels[f'subj0{subj}'] = betas
+# print(f"num_voxels for subj0{subj}: {num_voxels}")
+
+# if not new_test: # using old test set from before full dataset released (used in original MindEye paper)
+#     if subj==3:
+#         num_test=2113
+#     elif subj==4:
+#         num_test=1985
+#     elif subj==6:
+#         num_test=2113
+#     elif subj==8:
+#         num_test=1985
+#     else:
+#         num_test=2770
+#     test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+# else: # using larger test set from after full dataset released
+#     if subj==3:
+#         num_test=2371
+#     elif subj==4:
+#         num_test=2188
+#     elif subj==6:
+#         num_test=2371
+#     elif subj==8:
+#         num_test=2188
+#     else:
+#         num_test=3000
+#     test_url = f"{data_path}/wds/subj0{subj}/new_test/" + "0.tar"
+    
+# print(test_url)
+# def my_split_by_node(urls): return urls
+# test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\
+#                     .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=num_test, shuffle=False, drop_last=True, pin_memory=True)
+# print(f"Loaded test dl for subj{subj}!\n")
+
+
+# In[6]:
+
+
+# # Prep images but don't load them all to memory
+# f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+# images = f['images']
+
+# # Prep test voxels and indices of test images
+# test_images_idx = []
+# test_voxels_idx = []
+# for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+#     test_voxels = voxels[f'subj0{subj}'][behav[:,0,5].cpu().long()]
+#     test_voxels_idx = np.append(test_images_idx, behav[:,0,5].cpu().numpy())
+#     test_images_idx = np.append(test_images_idx, behav[:,0,0].cpu().numpy())
+# test_images_idx = test_images_idx.astype(int)
+# test_voxels_idx = test_voxels_idx.astype(int)
+
+# assert (test_i+1) * num_test == len(test_voxels) == len(test_images_idx)
+# print(test_i, len(test_voxels), len(test_images_idx), len(np.unique(test_images_idx)))
+
+
+# In[7]:
+
+
+clip_img_embedder = FrozenOpenCLIPImageEmbedder(
+    arch="ViT-bigG-14",
+    version="laion2b_s39b_b160k",
+    output_tokens=True,
+    only_tokens=True,
+)
+clip_img_embedder.to(device)
+clip_seq_dim = 256
+clip_emb_dim = 1664
+
+if blurry_recon:
+    from diffusers import AutoencoderKL
+    autoenc = AutoencoderKL(
+        down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D'],
+        up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D'],
+        block_out_channels=[128, 256, 512, 512],
+        layers_per_block=2,
+        sample_size=256,
+    )
+    ckpt = torch.load(f'{cache_dir}/sd_image_var_autoenc.pth')
+    autoenc.load_state_dict(ckpt)
+    autoenc.eval()
+    autoenc.requires_grad_(False)
+    autoenc.to(device)
+    utils.count_params(autoenc)
+    
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer to enable regularization
+    def __init__(self, input_sizes, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linears = torch.nn.ModuleList([
+                torch.nn.Linear(input_size, out_features) for input_size in input_sizes
+            ])
+    def forward(self, x, subj_idx):
+        out = self.linears[subj_idx](x[:,0]).unsqueeze(1)
+        return out
+        
+model.ridge = RidgeRegression([num_voxels], out_features=hidden_dim)
+
+from diffusers.models.vae import Decoder
+from models import BrainNetwork
+model.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, seq_len=1, 
+                          clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim) 
+utils.count_params(model.ridge)
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+# setup diffusion prior network
+out_dim = clip_emb_dim
+depth = 6
+dim_head = 52
+heads = clip_emb_dim//52 # heads * dim_head = clip_emb_dim
+timesteps = 100
+
+prior_network = PriorNetwork(
+        dim=out_dim,
+        depth=depth,
+        dim_head=dim_head,
+        heads=heads,
+        causal=False,
+        num_tokens = clip_seq_dim,
+        learned_query_mode="pos_emb"
+    )
+
+model.diffusion_prior = BrainDiffusionPrior(
+    net=prior_network,
+    image_embed_dim=out_dim,
+    condition_on_text_encodings=False,
+    timesteps=timesteps,
+    cond_drop_prob=0.2,
+    image_embed_scale=None,
+)
+model.to(device)
+
+utils.count_params(model.diffusion_prior)
+utils.count_params(model)
+
+# Load pretrained model ckpt
+tag='last'
+outdir = os.path.abspath(f'../train_logs/{model_name}')
+print(f"\n---loading {outdir}/{tag}.pth ckpt---\n")
+try:
+    checkpoint = torch.load(outdir+f'/{tag}.pth', map_location='cpu')
+    state_dict = checkpoint['model_state_dict']
+    model.load_state_dict(state_dict, strict=True)
+    del checkpoint
+except: # probably ckpt is saved using deepspeed format
+    import deepspeed
+    state_dict = deepspeed.utils.zero_to_fp32.get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir=outdir, tag=tag)
+    model.load_state_dict(state_dict, strict=False)
+    del state_dict
+print("ckpt loaded!")
+
+
+# In[8]:
+
+
+# setup text caption networks
+from transformers import AutoProcessor, AutoModelForCausalLM
+from modeling_git import GitForCausalLMClipEmb
+processor = AutoProcessor.from_pretrained("microsoft/git-large-coco")
+clip_text_model = GitForCausalLMClipEmb.from_pretrained("microsoft/git-large-coco")
+clip_text_model.to(device) # if you get OOM running this script, you can switch this to cpu and lower minibatch_size to 4
+clip_text_model.eval().requires_grad_(False)
+clip_text_seq_dim = 257
+clip_text_emb_dim = 1024
+
+class CLIPConverter(torch.nn.Module):
+    def __init__(self):
+        super(CLIPConverter, self).__init__()
+        self.linear1 = nn.Linear(clip_seq_dim, clip_text_seq_dim)
+        self.linear2 = nn.Linear(clip_emb_dim, clip_text_emb_dim)
+    def forward(self, x):
+        x = x.permute(0,2,1)
+        x = self.linear1(x)
+        x = self.linear2(x.permute(0,2,1))
+        return x
+        
+clip_convert = CLIPConverter()
+state_dict = torch.load(f"{cache_dir}/bigG_to_L_epoch8.pth", map_location='cpu')['model_state_dict']
+clip_convert.load_state_dict(state_dict, strict=True)
+clip_convert.to(device) # if you get OOM running this script, you can switch this to cpu and lower minibatch_size to 4
+del state_dict
+
+
+# In[9]:
+
+
+# prep unCLIP
+config = OmegaConf.load("generative_models/configs/unclip6.yaml")
+config = OmegaConf.to_container(config, resolve=True)
+unclip_params = config["model"]["params"]
+network_config = unclip_params["network_config"]
+denoiser_config = unclip_params["denoiser_config"]
+first_stage_config = unclip_params["first_stage_config"]
+conditioner_config = unclip_params["conditioner_config"]
+sampler_config = unclip_params["sampler_config"]
+scale_factor = unclip_params["scale_factor"]
+disable_first_stage_autocast = unclip_params["disable_first_stage_autocast"]
+offset_noise_level = unclip_params["loss_fn_config"]["params"]["offset_noise_level"]
+
+first_stage_config['target'] = 'sgm.models.autoencoder.AutoencoderKL'
+sampler_config['params']['num_steps'] = 38
+
+diffusion_engine = DiffusionEngine(network_config=network_config,
+                       denoiser_config=denoiser_config,
+                       first_stage_config=first_stage_config,
+                       conditioner_config=conditioner_config,
+                       sampler_config=sampler_config,
+                       scale_factor=scale_factor,
+                       disable_first_stage_autocast=disable_first_stage_autocast)
+# set to inference
+diffusion_engine.eval().requires_grad_(False)
+diffusion_engine.to(device)
+
+ckpt_path = f'{cache_dir}/unclip6_epoch0_step110000.ckpt'
+ckpt = torch.load(ckpt_path, map_location='cpu')
+diffusion_engine.load_state_dict(ckpt['state_dict'])
+
+batch={"jpg": torch.randn(1,3,1,1).to(device), # jpg doesnt get used, it's just a placeholder
+      "original_size_as_tuple": torch.ones(1, 2).to(device) * 768,
+      "crop_coords_top_left": torch.zeros(1, 2).to(device)}
+out = diffusion_engine.conditioner(batch)
+vector_suffix = out["vector"].to(device)
+print("vector_suffix", vector_suffix.shape)
+
+
+# In[10]:
+
+
+from tqdm.auto import tqdm
+# get all reconstructions
+model.to(device)
+model.eval().requires_grad_(False)
+
+final_all_recons = None
+final_all_blurryrecons = None
+final_all_predcaptions = []
+final_all_clipvoxels = None
+
+for rep in tqdm(range(gen_rep)):
+    # all_images = None
+    all_blurryrecons = None
+    all_recons = None
+    all_predcaptions = []
+    all_clipvoxels = None
+
+    minibatch_size = 1
+    num_samples_per_image = 1
+    assert num_samples_per_image == 1
+
+    if utils.is_interactive(): plotting=False
+
+    with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.float16):
+        for idx, imagery_voxel in enumerate(voxels):
+            # uniq_imgs = np.unique(test_images_idx)[batch:batch+minibatch_size]
+            # voxel = None
+            # for uniq_img in uniq_imgs:
+            #     locs = np.where(test_images_idx==uniq_img)[0]
+            #     if len(locs)==1:
+            #         locs = locs.repeat(3)
+            #     elif len(locs)==2:
+            #         locs = locs.repeat(2)[:3]
+            #     assert len(locs)==3
+            #     if voxel is None:
+            #         voxel = test_voxels[None,locs] # 1, num_image_repetitions, num_voxels
+            #     else:
+            #         voxel = torch.vstack((voxel, test_voxels[None,locs]))
+            voxel = imagery_voxel.to(device)
+            print(voxel.shape, voxel.shape)
+            for rep in range(1):
+                voxel_ridge = model.ridge(voxel.unsqueeze(0),0) # 0th index of subj_list
+                backbone0, clip_voxels0, blurry_image_enc0 = model.backbone(voxel_ridge)
+                if rep==0:
+                    clip_voxels = clip_voxels0
+                    backbone = backbone0
+                    blurry_image_enc = blurry_image_enc0[0]
+                else:
+                    clip_voxels += clip_voxels0
+                    backbone += backbone0
+                    blurry_image_enc += blurry_image_enc0[0]
+            clip_voxels /= 1
+            backbone /= 1
+            blurry_image_enc /= 1
+                    
+            # Save retrieval submodule outputs
+            if all_clipvoxels is None:
+                all_clipvoxels = clip_voxels.cpu()
+            else:
+                all_clipvoxels = torch.vstack((all_clipvoxels, clip_voxels.cpu()))
+            
+            # Feed voxels through OpenCLIP-bigG diffusion prior
+            prior_out = model.diffusion_prior.p_sample_loop(backbone.shape, 
+                            text_cond = dict(text_embed = backbone), 
+                            cond_scale = 1., timesteps = 20)
+            
+            pred_caption_emb = clip_convert(prior_out)
+            generated_ids = clip_text_model.generate(pixel_values=pred_caption_emb, max_length=20)
+            generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)
+            all_predcaptions = np.hstack((all_predcaptions, generated_caption))
+            print(generated_caption)
+            
+            # Feed diffusion prior outputs through unCLIP
+            for i in range(len(voxel)):
+                samples = utils.unclip_recon(prior_out[[i]],
+                                diffusion_engine,
+                                vector_suffix,
+                                num_samples=num_samples_per_image)
+                if all_recons is None:
+                    all_recons = samples.cpu()
+                else:
+                    all_recons = torch.vstack((all_recons, samples.cpu()))
+                if plotting:
+                    for s in range(num_samples_per_image):
+                        plt.figure(figsize=(2,2))
+                        plt.imshow(transforms.ToPILImage()(samples[s]))
+                        plt.axis('off')
+                        plt.show()
+
+            if blurry_recon:
+                blurred_image = (autoenc.decode(blurry_image_enc/0.18215).sample/ 2 + 0.5).clamp(0,1)
+                
+                for i in range(len(voxel)):
+                    im = torch.Tensor(blurred_image[i])
+                    if all_blurryrecons is None:
+                        all_blurryrecons = im[None].cpu()
+                    else:
+                        all_blurryrecons = torch.vstack((all_blurryrecons, im[None].cpu()))
+                    if plotting:
+                        plt.figure(figsize=(2,2))
+                        plt.imshow(transforms.ToPILImage()(im))
+                        plt.axis('off')
+                        plt.show()
+
+            if plotting: 
+                print(model_name)
+                err # dont actually want to run the whole thing with plotting=True
+
+    # resize outputs before saving
+    imsize = 256
+    all_recons = transforms.Resize((imsize,imsize))(all_recons).float()
+    if blurry_recon: 
+        all_blurryrecons = transforms.Resize((imsize,imsize))(all_blurryrecons).float()
+            
+    # saving
+    print(all_recons.shape)
+    # # You can find the all_images file on huggingface: https://huggingface.co/datasets/pscotti/mindeyev2/tree/main/evals
+    # torch.save(all_images,"evals/all_images.pt") 
+    # if blurry_recon:
+    #     torch.save(all_blurryrecons,f"evals/{model_name}/{model_name}_all_blurryrecons.pt")
+    # torch.save(all_recons,f"evals/{model_name}/{model_name}_all_recons.pt")
+    # torch.save(all_predcaptions,f"evals/{model_name}/{model_name}_all_predcaptions.pt")
+    # torch.save(all_clipvoxels,f"evals/{model_name}/{model_name}_all_clipvoxels.pt")
+
+    if final_all_recons is None:
+        if blurry_recon:
+            final_all_blurryrecons = all_blurryrecons.unsqueeze(1)
+        final_all_recons = all_recons.unsqueeze(1)
+        final_all_predcaptions = all_predcaptions.copy()
+        final_all_clipvoxels = all_clipvoxels.unsqueeze(1)
+    else:
+        if blurry_recon:
+            final_all_blurryrecons = torch.cat((final_all_blurryrecons, all_blurryrecons.unsqueeze(1)), dim=1)
+        final_all_recons = torch.cat((final_all_recons, all_recons.unsqueeze(1)), dim=1)
+        final_all_predcaptions = np.hstack((final_all_predcaptions, all_predcaptions))
+        final_all_clipvoxels = torch.cat((final_all_clipvoxels, all_clipvoxels.unsqueeze(1)), dim=1)
+
+
+# saving
+print(final_all_recons.shape)
+# You can find the all_images file on huggingface: https://huggingface.co/datasets/pscotti/mindeyev2/tree/main/evals
+# torch.save(all_images,"evals/all_images.pt")
+
+
+# if blurry_recon:
+#     torch.save(final_blurryrecons.half(),f"evals/{model_name}/{model_name}_all_blurryrecons_{mode}.pt")
+# torch.save(final_recons.half(),f"evals/{model_name}/{model_name}_all_recons_{mode}.pt")
+
+if blurry_recon:
+    torch.save(final_all_blurryrecons,f"evals/{model_name}/{model_name}_all_blurryrecons_{mode}.pt")
+torch.save(final_all_recons,f"evals/{model_name}/{model_name}_all_recons_{mode}.pt")
+torch.save(final_all_predcaptions,f"evals/{model_name}/{model_name}_all_predcaptions_{mode}.pt")
+torch.save(final_all_clipvoxels,f"evals/{model_name}/{model_name}_all_clipvoxels_{mode}.pt")
+
+
+print(f"saved {model_name} outputs!")
+
+if not utils.is_interactive():
+    sys.exit(0)
+
+
+# In[11]:
+
+
+torch.save(all_images,"evals/all_images.pt")
+

MindEyeV2/src/recon_inference_old.py

@@ -0,0 +1,492 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+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 matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from accelerate import Accelerator
+
+# SDXL unCLIP requires code from https://github.com/Stability-AI/generative-models/tree/main
+sys.path.append('generative_models/')
+import sgm
+from generative_models.sgm.modules.encoders.modules import FrozenOpenCLIPImageEmbedder, FrozenOpenCLIPEmbedder2
+from generative_models.sgm.models.diffusion import DiffusionEngine
+from generative_models.sgm.util import append_dims
+from omegaconf import OmegaConf
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+from models import *
+
+accelerator = Accelerator(split_batches=False, mixed_precision="fp16")
+device = accelerator.device
+print("device:",device)
+plotting = False
+
+
+# In[2]:
+
+
+# if running this interactively, can specify jupyter_args here for argparser to use
+if utils.is_interactive():
+    model_name = "final_subj01_pretrained_3sess_24bs"
+    print("model_name:", model_name)
+
+    # other variables can be specified in the following string:
+    jupyter_args = f"--data_path=/weka/proj-medarc/shared/mindeyev2_dataset \
+                    --cache_dir=/weka/proj-medarc/shared/mindeyev2_dataset \
+                    --model_name={model_name} --subj=1 \
+                    --hidden_dim=4096 --n_blocks=4 --new_test"
+    print(jupyter_args)
+    jupyter_args = jupyter_args.split()
+    
+    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[3]:
+
+
+parser = argparse.ArgumentParser(description="Model Training Configuration")
+parser.add_argument(
+    "--model_name", type=str, default="testing",
+    help="will load ckpt for model found in ../train_logs/model_name",
+)
+parser.add_argument(
+    "--data_path", type=str, default=os.getcwd(),
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--cache_dir", type=str, default=os.getcwd(),
+    help="Path to where misc. files downloaded from huggingface are stored. Defaults to current src directory.",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,3,4,5,6,7,8],
+    help="Validate on which subject?",
+)
+parser.add_argument(
+    "--blurry_recon",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--n_blocks",type=int,default=4,
+)
+parser.add_argument(
+    "--hidden_dim",type=int,default=2048,
+)
+parser.add_argument(
+    "--new_test",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--seed",type=int,default=42,
+)
+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)
+    
+# seed all random functions
+utils.seed_everything(seed)
+
+# make output directory
+os.makedirs("evals",exist_ok=True)
+os.makedirs(f"evals/{model_name}",exist_ok=True)
+
+
+# In[4]:
+
+
+voxels = {}
+# Load hdf5 data for betas
+f = h5py.File(f'{data_path}/betas_all_subj0{subj}_fp32_renorm.hdf5', 'r')
+betas = f['betas'][:]
+betas = torch.Tensor(betas).to("cpu")
+num_voxels = betas[0].shape[-1]
+voxels[f'subj0{subj}'] = betas
+print(f"num_voxels for subj0{subj}: {num_voxels}")
+
+if not new_test: # using old test set from before full dataset released (used in original MindEye paper)
+    if subj==3:
+        num_test=2113
+    elif subj==4:
+        num_test=1985
+    elif subj==6:
+        num_test=2113
+    elif subj==8:
+        num_test=1985
+    else:
+        num_test=2770
+    test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+else: # using larger test set from after full dataset released
+    if subj==3:
+        num_test=2371
+    elif subj==4:
+        num_test=2188
+    elif subj==6:
+        num_test=2371
+    elif subj==8:
+        num_test=2188
+    else:
+        num_test=3000
+    test_url = f"{data_path}/wds/subj0{subj}/new_test/" + "0.tar"
+    
+print(test_url)
+def my_split_by_node(urls): return urls
+test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .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=num_test, shuffle=False, drop_last=True, pin_memory=True)
+print(f"Loaded test dl for subj{subj}!\n")
+
+
+# In[5]:
+
+
+# Prep images but don't load them all to memory
+f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+images = f['images']
+
+# Prep test voxels and indices of test images
+test_images_idx = []
+test_voxels_idx = []
+for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+    test_voxels = voxels[f'subj0{subj}'][behav[:,0,5].cpu().long()]
+    test_voxels_idx = np.append(test_images_idx, behav[:,0,5].cpu().numpy())
+    test_images_idx = np.append(test_images_idx, behav[:,0,0].cpu().numpy())
+test_images_idx = test_images_idx.astype(int)
+test_voxels_idx = test_voxels_idx.astype(int)
+
+assert (test_i+1) * num_test == len(test_voxels) == len(test_images_idx)
+print(test_i, len(test_voxels), len(test_images_idx), len(np.unique(test_images_idx)))
+
+
+# In[6]:
+
+
+clip_img_embedder = FrozenOpenCLIPImageEmbedder(
+    arch="ViT-bigG-14",
+    version="laion2b_s39b_b160k",
+    output_tokens=True,
+    only_tokens=True,
+)
+clip_img_embedder.to(device)
+clip_seq_dim = 256
+clip_emb_dim = 1664
+
+if blurry_recon:
+    from diffusers import AutoencoderKL
+    autoenc = AutoencoderKL(
+        down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D'],
+        up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D'],
+        block_out_channels=[128, 256, 512, 512],
+        layers_per_block=2,
+        sample_size=256,
+    )
+    ckpt = torch.load(f'{cache_dir}/sd_image_var_autoenc.pth')
+    autoenc.load_state_dict(ckpt)
+    autoenc.eval()
+    autoenc.requires_grad_(False)
+    autoenc.to(device)
+    utils.count_params(autoenc)
+    
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer to enable regularization
+    def __init__(self, input_sizes, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linears = torch.nn.ModuleList([
+                torch.nn.Linear(input_size, out_features) for input_size in input_sizes
+            ])
+    def forward(self, x, subj_idx):
+        out = self.linears[subj_idx](x[:,0]).unsqueeze(1)
+        return out
+        
+model.ridge = RidgeRegression([num_voxels], out_features=hidden_dim)
+
+from diffusers.models.vae import Decoder
+from models import BrainNetwork
+model.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, seq_len=1, 
+                          clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim) 
+utils.count_params(model.ridge)
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+# setup diffusion prior network
+out_dim = clip_emb_dim
+depth = 6
+dim_head = 52
+heads = clip_emb_dim//52 # heads * dim_head = clip_emb_dim
+timesteps = 100
+
+prior_network = PriorNetwork(
+        dim=out_dim,
+        depth=depth,
+        dim_head=dim_head,
+        heads=heads,
+        causal=False,
+        num_tokens = clip_seq_dim,
+        learned_query_mode="pos_emb"
+    )
+
+model.diffusion_prior = BrainDiffusionPrior(
+    net=prior_network,
+    image_embed_dim=out_dim,
+    condition_on_text_encodings=False,
+    timesteps=timesteps,
+    cond_drop_prob=0.2,
+    image_embed_scale=None,
+)
+model.to(device)
+
+utils.count_params(model.diffusion_prior)
+utils.count_params(model)
+
+# Load pretrained model ckpt
+tag='last'
+outdir = os.path.abspath(f'../train_logs/{model_name}')
+print(f"\n---loading {outdir}/{tag}.pth ckpt---\n")
+try:
+    checkpoint = torch.load(outdir+f'/{tag}.pth', map_location='cpu')
+    state_dict = checkpoint['model_state_dict']
+    model.load_state_dict(state_dict, strict=True)
+    del checkpoint
+except: # probably ckpt is saved using deepspeed format
+    import deepspeed
+    state_dict = deepspeed.utils.zero_to_fp32.get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir=outdir, tag=tag)
+    model.load_state_dict(state_dict, strict=False)
+    del state_dict
+print("ckpt loaded!")
+
+
+# In[7]:
+
+
+# setup text caption networks
+from transformers import AutoProcessor, AutoModelForCausalLM
+from modeling_git import GitForCausalLMClipEmb
+processor = AutoProcessor.from_pretrained("microsoft/git-large-coco")
+clip_text_model = GitForCausalLMClipEmb.from_pretrained("microsoft/git-large-coco")
+clip_text_model.to(device) # if you get OOM running this script, you can switch this to cpu and lower minibatch_size to 4
+clip_text_model.eval().requires_grad_(False)
+clip_text_seq_dim = 257
+clip_text_emb_dim = 1024
+
+class CLIPConverter(torch.nn.Module):
+    def __init__(self):
+        super(CLIPConverter, self).__init__()
+        self.linear1 = nn.Linear(clip_seq_dim, clip_text_seq_dim)
+        self.linear2 = nn.Linear(clip_emb_dim, clip_text_emb_dim)
+    def forward(self, x):
+        x = x.permute(0,2,1)
+        x = self.linear1(x)
+        x = self.linear2(x.permute(0,2,1))
+        return x
+        
+clip_convert = CLIPConverter()
+state_dict = torch.load(f"{cache_dir}/bigG_to_L_epoch8.pth", map_location='cpu')['model_state_dict']
+clip_convert.load_state_dict(state_dict, strict=True)
+clip_convert.to(device) # if you get OOM running this script, you can switch this to cpu and lower minibatch_size to 4
+del state_dict
+
+
+# In[8]:
+
+
+# prep unCLIP
+config = OmegaConf.load("generative_models/configs/unclip6.yaml")
+config = OmegaConf.to_container(config, resolve=True)
+unclip_params = config["model"]["params"]
+network_config = unclip_params["network_config"]
+denoiser_config = unclip_params["denoiser_config"]
+first_stage_config = unclip_params["first_stage_config"]
+conditioner_config = unclip_params["conditioner_config"]
+sampler_config = unclip_params["sampler_config"]
+scale_factor = unclip_params["scale_factor"]
+disable_first_stage_autocast = unclip_params["disable_first_stage_autocast"]
+offset_noise_level = unclip_params["loss_fn_config"]["params"]["offset_noise_level"]
+
+first_stage_config['target'] = 'sgm.models.autoencoder.AutoencoderKL'
+sampler_config['params']['num_steps'] = 38
+
+diffusion_engine = DiffusionEngine(network_config=network_config,
+                       denoiser_config=denoiser_config,
+                       first_stage_config=first_stage_config,
+                       conditioner_config=conditioner_config,
+                       sampler_config=sampler_config,
+                       scale_factor=scale_factor,
+                       disable_first_stage_autocast=disable_first_stage_autocast)
+# set to inference
+diffusion_engine.eval().requires_grad_(False)
+diffusion_engine.to(device)
+
+ckpt_path = f'{cache_dir}/unclip6_epoch0_step110000.ckpt'
+ckpt = torch.load(ckpt_path, map_location='cpu')
+diffusion_engine.load_state_dict(ckpt['state_dict'])
+
+batch={"jpg": torch.randn(1,3,1,1).to(device), # jpg doesnt get used, it's just a placeholder
+      "original_size_as_tuple": torch.ones(1, 2).to(device) * 768,
+      "crop_coords_top_left": torch.zeros(1, 2).to(device)}
+out = diffusion_engine.conditioner(batch)
+vector_suffix = out["vector"].to(device)
+print("vector_suffix", vector_suffix.shape)
+
+
+# In[9]:
+
+
+# get all reconstructions
+model.to(device)
+model.eval().requires_grad_(False)
+
+# all_images = None
+all_blurryrecons = None
+all_recons = None
+all_predcaptions = []
+all_clipvoxels = None
+
+minibatch_size = 1
+num_samples_per_image = 1
+assert num_samples_per_image == 1
+
+if utils.is_interactive(): plotting=True
+
+with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.float16):
+    for batch in tqdm(range(0,len(np.unique(test_images_idx)),minibatch_size)):
+        uniq_imgs = np.unique(test_images_idx)[batch:batch+minibatch_size]
+        voxel = None
+        for uniq_img in uniq_imgs:
+            locs = np.where(test_images_idx==uniq_img)[0]
+            if len(locs)==1:
+                locs = locs.repeat(3)
+            elif len(locs)==2:
+                locs = locs.repeat(2)[:3]
+            assert len(locs)==3
+            if voxel is None:
+                voxel = test_voxels[None,locs] # 1, num_image_repetitions, num_voxels
+            else:
+                voxel = torch.vstack((voxel, test_voxels[None,locs]))
+        voxel = voxel.to(device)
+        
+        for rep in range(3):
+            voxel_ridge = model.ridge(voxel[:,[rep]],0) # 0th index of subj_list
+            backbone0, clip_voxels0, blurry_image_enc0 = model.backbone(voxel_ridge)
+            if rep==0:
+                clip_voxels = clip_voxels0
+                backbone = backbone0
+                blurry_image_enc = blurry_image_enc0[0]
+            else:
+                clip_voxels += clip_voxels0
+                backbone += backbone0
+                blurry_image_enc += blurry_image_enc0[0]
+        clip_voxels /= 3
+        backbone /= 3
+        blurry_image_enc /= 3
+                
+        # Save retrieval submodule outputs
+        if all_clipvoxels is None:
+            all_clipvoxels = clip_voxels.cpu()
+        else:
+            all_clipvoxels = torch.vstack((all_clipvoxels, clip_voxels.cpu()))
+        
+        # Feed voxels through OpenCLIP-bigG diffusion prior
+        prior_out = model.diffusion_prior.p_sample_loop(backbone.shape, 
+                        text_cond = dict(text_embed = backbone), 
+                        cond_scale = 1., timesteps = 20)
+        
+        pred_caption_emb = clip_convert(prior_out)
+        generated_ids = clip_text_model.generate(pixel_values=pred_caption_emb, max_length=20)
+        generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)
+        all_predcaptions = np.hstack((all_predcaptions, generated_caption))
+        print(generated_caption)
+        
+        # Feed diffusion prior outputs through unCLIP
+        for i in range(len(voxel)):
+            samples = utils.unclip_recon(prior_out[[i]],
+                             diffusion_engine,
+                             vector_suffix,
+                             num_samples=num_samples_per_image)
+            if all_recons is None:
+                all_recons = samples.cpu()
+            else:
+                all_recons = torch.vstack((all_recons, samples.cpu()))
+            if plotting:
+                for s in range(num_samples_per_image):
+                    plt.figure(figsize=(2,2))
+                    plt.imshow(transforms.ToPILImage()(samples[s]))
+                    plt.axis('off')
+                    plt.show()
+
+        if blurry_recon:
+            blurred_image = (autoenc.decode(blurry_image_enc/0.18215).sample/ 2 + 0.5).clamp(0,1)
+            
+            for i in range(len(voxel)):
+                im = torch.Tensor(blurred_image[i])
+                if all_blurryrecons is None:
+                    all_blurryrecons = im[None].cpu()
+                else:
+                    all_blurryrecons = torch.vstack((all_blurryrecons, im[None].cpu()))
+                if plotting:
+                    plt.figure(figsize=(2,2))
+                    plt.imshow(transforms.ToPILImage()(im))
+                    plt.axis('off')
+                    plt.show()
+
+        if plotting: 
+            print(model_name)
+            err # dont actually want to run the whole thing with plotting=True
+
+# resize outputs before saving
+imsize = 256
+all_recons = transforms.Resize((imsize,imsize))(all_recons).float()
+if blurry_recon: 
+    all_blurryrecons = transforms.Resize((imsize,imsize))(all_blurryrecons).float()
+        
+# saving
+print(all_recons.shape)
+# # You can find the all_images file on huggingface: https://huggingface.co/datasets/pscotti/mindeyev2/tree/main/evals
+# torch.save(all_images,"evals/all_images.pt") 
+if blurry_recon:
+    torch.save(all_blurryrecons,f"evals/{model_name}/{model_name}_all_blurryrecons.pt")
+torch.save(all_recons,f"evals/{model_name}/{model_name}_all_recons.pt")
+torch.save(all_predcaptions,f"evals/{model_name}/{model_name}_all_predcaptions.pt")
+torch.save(all_clipvoxels,f"evals/{model_name}/{model_name}_all_clipvoxels.pt")
+print(f"saved {model_name} outputs!")
+
+if not utils.is_interactive():
+    sys.exit(0)
+

MindEyeV2/src/wandb/debug.log

@@ -0,0 +1,30 @@
+2024-12-10 21:57:28,176 INFO    MainThread:842903 [wandb_setup.py:_flush():76] Current SDK version is 0.17.1
+2024-12-10 21:57:28,177 INFO    MainThread:842903 [wandb_setup.py:_flush():76] Configure stats pid to 842903
+2024-12-10 21:57:28,177 INFO    MainThread:842903 [wandb_setup.py:_flush():76] Loading settings from /admin/home-ckadirt/.config/wandb/settings
+2024-12-10 21:57:28,177 INFO    MainThread:842903 [wandb_setup.py:_flush():76] Loading settings from /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/settings
+2024-12-10 21:57:28,177 INFO    MainThread:842903 [wandb_setup.py:_flush():76] Loading settings from environment variables: {}
+2024-12-10 21:57:28,177 INFO    MainThread:842903 [wandb_setup.py:_flush():76] Applying setup settings: {'_disable_service': False}
+2024-12-10 21:57:28,177 INFO    MainThread:842903 [wandb_setup.py:_flush():76] Inferring run settings from compute environment: {'program_relpath': 'src/TrainB5k.py', 'program_abspath': '/weka/proj-fmri/ckadirt/MindEyeV2/src/TrainB5k.py', 'program': '/weka/proj-fmri/ckadirt/MindEyeV2/src/TrainB5k.py'}
+2024-12-10 21:57:28,177 INFO    MainThread:842903 [wandb_setup.py:_flush():76] Applying login settings: {}
+2024-12-10 21:57:28,177 INFO    MainThread:842903 [wandb_init.py:_log_setup():520] Logging user logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20241210_215728-bold5k_nsdm1/logs/debug.log
+2024-12-10 21:57:28,177 INFO    MainThread:842903 [wandb_init.py:_log_setup():521] Logging internal logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20241210_215728-bold5k_nsdm1/logs/debug-internal.log
+2024-12-10 21:57:28,177 INFO    MainThread:842903 [wandb_init.py:init():560] calling init triggers
+2024-12-10 21:57:28,177 INFO    MainThread:842903 [wandb_init.py:init():567] wandb.init called with sweep_config: {}
+config: {'model_name': 'bold5k_nsdm1', 'global_batch_size': '21', 'batch_size': 21, 'num_epochs': 150, 'num_sessions': 15, 'num_params': 2239051016, 'clip_scale': 1.0, 'prior_scale': 30.0, 'blur_scale': 0.5, 'use_image_aug': True, 'max_lr': 0.0003, 'mixup_pct': 0.33, 'num_samples_per_epoch': 4323, 'num_test': 480, 'ckpt_interval': 999, 'ckpt_saving': False, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1}
+2024-12-10 21:57:28,177 INFO    MainThread:842903 [wandb_init.py:init():610] starting backend
+2024-12-10 21:57:28,177 INFO    MainThread:842903 [wandb_init.py:init():614] setting up manager
+2024-12-10 21:57:28,179 INFO    MainThread:842903 [backend.py:_multiprocessing_setup():105] multiprocessing start_methods=fork,spawn,forkserver, using: spawn
+2024-12-10 21:57:28,182 INFO    MainThread:842903 [wandb_init.py:init():622] backend started and connected
+2024-12-10 21:57:28,187 INFO    MainThread:842903 [wandb_init.py:init():711] updated telemetry
+2024-12-10 21:57:28,198 INFO    MainThread:842903 [wandb_init.py:init():744] communicating run to backend with 90.0 second timeout
+2024-12-10 21:57:28,693 INFO    MainThread:842903 [wandb_init.py:init():787] run resumed
+2024-12-10 21:57:28,707 INFO    MainThread:842903 [wandb_run.py:_on_init():2402] communicating current version
+2024-12-10 21:57:28,770 INFO    MainThread:842903 [wandb_run.py:_on_init():2411] got version response upgrade_message: "wandb version 0.19.0 is available!  To upgrade, please run:\n $ pip install wandb --upgrade"
+
+2024-12-10 21:57:28,770 INFO    MainThread:842903 [wandb_init.py:init():795] starting run threads in backend
+2024-12-10 21:57:29,319 INFO    MainThread:842903 [wandb_run.py:_console_start():2380] atexit reg
+2024-12-10 21:57:29,319 INFO    MainThread:842903 [wandb_run.py:_redirect():2235] redirect: wrap_raw
+2024-12-10 21:57:29,319 INFO    MainThread:842903 [wandb_run.py:_redirect():2300] Wrapping output streams.
+2024-12-10 21:57:29,319 INFO    MainThread:842903 [wandb_run.py:_redirect():2325] Redirects installed.
+2024-12-10 21:57:29,322 INFO    MainThread:842903 [wandb_init.py:init():838] run started, returning control to user process
+2024-12-11 00:59:32,641 WARNING MsgRouterThr:842903 [router.py:message_loop():77] message_loop has been closed

MindEyeV2/src/wandb/run-20240618_035407-rr-test/files/config.yaml

@@ -0,0 +1,102 @@
+wandb_version: 1
+
+model_name:
+  desc: null
+  value: rr-test
+global_batch_size:
+  desc: null
+  value: '21'
+batch_size:
+  desc: null
+  value: 3
+num_epochs:
+  desc: null
+  value: 150
+num_sessions:
+  desc: null
+  value: 40
+num_params:
+  desc: null
+  value: 573919937
+clip_scale:
+  desc: null
+  value: 1.0
+prior_scale:
+  desc: null
+  value: 30.0
+blur_scale:
+  desc: null
+  value: 0.5
+use_image_aug:
+  desc: null
+  value: false
+max_lr:
+  desc: null
+  value: 0.0003
+mixup_pct:
+  desc: null
+  value: 0.33
+num_samples_per_epoch:
+  desc: null
+  value: 3750
+num_test:
+  desc: null
+  value: 3000
+ckpt_interval:
+  desc: null
+  value: 999
+ckpt_saving:
+  desc: null
+  value: true
+seed:
+  desc: null
+  value: 42
+distributed:
+  desc: null
+  value: true
+num_devices:
+  desc: null
+  value: 8
+world_size:
+  desc: null
+  value: 1
+train_url:
+  desc: null
+  value: /weka/proj-medarc/shared/mindeyev2_dataset/wds/subj08/train/{0..29}.tar
+test_url:
+  desc: null
+  value: /weka/proj-medarc/shared/mindeyev2_dataset/wds/subj02/new_test/0.tar
+_wandb:
+  desc: null
+  value:
+    python_version: 3.11.9
+    cli_version: 0.17.1
+    framework: huggingface
+    huggingface_version: 4.37.2
+    is_jupyter_run: false
+    is_kaggle_kernel: false
+    start_time: 1718682847
+    t:
+      1:
+      - 1
+      - 9
+      - 11
+      - 41
+      - 49
+      - 55
+      - 63
+      - 71
+      - 79
+      - 83
+      - 103
+      3:
+      - 13
+      - 14
+      - 16
+      - 23
+      4: 3.11.9
+      5: 0.17.1
+      6: 4.37.2
+      8:
+      - 5
+      13: linux-x86_64

MindEyeV2/src/wandb/run-20240618_035407-rr-test/logs/debug-internal.log

@@ -0,0 +1,61 @@
+2024-06-18 03:54:07,351 INFO    StreamThr :2245613 [internal.py:wandb_internal():85] W&B internal server running at pid: 2245613, started at: 2024-06-18 03:54:07.345742
+2024-06-18 03:54:07,352 DEBUG   HandlerThread:2245613 [handler.py:handle_request():158] handle_request: status
+2024-06-18 03:54:07,375 INFO    WriterThread:2245613 [datastore.py:open_for_write():87] open: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_035407-rr-test/run-rr-test.wandb
+2024-06-18 03:54:07,379 DEBUG   SenderThread:2245613 [sender.py:send():379] send: header
+2024-06-18 03:54:07,462 DEBUG   SenderThread:2245613 [sender.py:send():379] send: run
+2024-06-18 03:54:07,474 INFO    SenderThread:2245613 [sender.py:_setup_resume():749] checking resume status for None/mindeye_semantic_cluster/rr-test
+2024-06-18 03:54:07,670 ERROR   SenderThread:2245613 [internal_api.py:execute():344] 401 response executing GraphQL.
+2024-06-18 03:54:07,670 ERROR   SenderThread:2245613 [internal_api.py:execute():345] {"errors":[{"message":"user is not logged in","path":["upsertBucket"],"extensions":{"code":"PERMISSION_ERROR"}}],"data":{"upsertBucket":null}}
+2024-06-18 03:54:07,670 ERROR   SenderThread:2245613 [sender.py:send_run():1009] The API key you provided is either invalid or missing.  If the `WANDB_API_KEY` environment variable is set, make sure it is correct. Otherwise, to resolve this issue, you may try running the 'wandb login --relogin' command. If you are using a local server, make sure that you're using the correct hostname. If you're not sure, you can try logging in again using the 'wandb login --relogin --host [hostname]' command.(Error 401: Unauthorized)
+Traceback (most recent call last):
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/sdk/lib/retry.py", line 131, in __call__
+    result = self._call_fn(*args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/sdk/internal/internal_api.py", line 340, in execute
+    return self.client.execute(*args, **kwargs)  # type: ignore
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/vendor/gql-0.2.0/wandb_gql/client.py", line 52, in execute
+    result = self._get_result(document, *args, **kwargs)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/vendor/gql-0.2.0/wandb_gql/client.py", line 60, in _get_result
+    return self.transport.execute(document, *args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/sdk/lib/gql_request.py", line 59, in execute
+    request.raise_for_status()
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/requests/models.py", line 1024, in raise_for_status
+    raise HTTPError(http_error_msg, response=self)
+requests.exceptions.HTTPError: 401 Client Error: Unauthorized for url: https://api.wandb.ai/graphql
+
+During handling of the above exception, another exception occurred:
+
+Traceback (most recent call last):
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/sdk/internal/sender.py", line 1007, in send_run
+    server_run = self._init_run(run, config_value_dict)
+                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/sdk/internal/sender.py", line 1079, in _init_run
+    server_run, inserted, server_messages = self._api.upsert_run(
+                                            ^^^^^^^^^^^^^^^^^^^^^
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/apis/normalize.py", line 73, in wrapper
+    raise err
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/apis/normalize.py", line 41, in wrapper
+    return func(*args, **kwargs)
+           ^^^^^^^^^^^^^^^^^^^^^
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/sdk/internal/internal_api.py", line 2188, in upsert_run
+    response = self.gql(
+               ^^^^^^^^^
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/sdk/internal/internal_api.py", line 312, in gql
+    ret = self._retry_gql(
+          ^^^^^^^^^^^^^^^^
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/sdk/lib/retry.py", line 147, in __call__
+    retry_timedelta_triggered = check_retry_fn(e)
+                                ^^^^^^^^^^^^^^^^^
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/util.py", line 965, in check_retry_fn
+    return fallback_retry_fn(e)
+           ^^^^^^^^^^^^^^^^^^^^
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/util.py", line 896, in no_retry_auth
+    raise AuthenticationError(
+wandb.errors.AuthenticationError: The API key you provided is either invalid or missing.  If the `WANDB_API_KEY` environment variable is set, make sure it is correct. Otherwise, to resolve this issue, you may try running the 'wandb login --relogin' command. If you are using a local server, make sure that you're using the correct hostname. If you're not sure, you can try logging in again using the 'wandb login --relogin --host [hostname]' command.(Error 401: Unauthorized)
+2024-06-18 03:54:07,847 DEBUG   HandlerThread:2245613 [handler.py:handle_request():158] handle_request: shutdown
+2024-06-18 03:54:07,847 INFO    HandlerThread:2245613 [handler.py:finish():882] shutting down handler
+2024-06-18 03:54:08,462 INFO    WriterThread:2245613 [datastore.py:close():296] close: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_035407-rr-test/run-rr-test.wandb
+2024-06-18 03:54:08,695 INFO    SenderThread:2245613 [sender.py:finish():1608] shutting down sender

MindEyeV2/src/wandb/run-20240618_035407-rr-test/logs/debug.log

@@ -0,0 +1,29 @@
+2024-06-18 03:54:07,333 INFO    MainThread:2245259 [wandb_setup.py:_flush():76] Current SDK version is 0.17.1
+2024-06-18 03:54:07,334 INFO    MainThread:2245259 [wandb_setup.py:_flush():76] Configure stats pid to 2245259
+2024-06-18 03:54:07,334 INFO    MainThread:2245259 [wandb_setup.py:_flush():76] Loading settings from /admin/home-ckadirt/.config/wandb/settings
+2024-06-18 03:54:07,334 INFO    MainThread:2245259 [wandb_setup.py:_flush():76] Loading settings from /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/settings
+2024-06-18 03:54:07,334 INFO    MainThread:2245259 [wandb_setup.py:_flush():76] Loading settings from environment variables: {}
+2024-06-18 03:54:07,334 INFO    MainThread:2245259 [wandb_setup.py:_flush():76] Applying setup settings: {'_disable_service': False}
+2024-06-18 03:54:07,334 INFO    MainThread:2245259 [wandb_setup.py:_flush():76] Inferring run settings from compute environment: {'program_relpath': 'src/TrainCluster.py', 'program_abspath': '/weka/proj-fmri/ckadirt/MindEyeV2/src/TrainCluster.py', 'program': '/weka/proj-fmri/ckadirt/MindEyeV2/src/TrainCluster.py'}
+2024-06-18 03:54:07,334 INFO    MainThread:2245259 [wandb_setup.py:_flush():76] Applying login settings: {}
+2024-06-18 03:54:07,334 INFO    MainThread:2245259 [wandb_init.py:_log_setup():520] Logging user logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_035407-rr-test/logs/debug.log
+2024-06-18 03:54:07,335 INFO    MainThread:2245259 [wandb_init.py:_log_setup():521] Logging internal logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_035407-rr-test/logs/debug-internal.log
+2024-06-18 03:54:07,335 INFO    MainThread:2245259 [wandb_init.py:init():560] calling init triggers
+2024-06-18 03:54:07,335 INFO    MainThread:2245259 [wandb_init.py:init():567] wandb.init called with sweep_config: {}
+config: {'model_name': 'rr-test', 'global_batch_size': '21', 'batch_size': 3, 'num_epochs': 150, 'num_sessions': 40, 'num_params': 573919937, 'clip_scale': 1.0, 'prior_scale': 30.0, 'blur_scale': 0.5, 'use_image_aug': False, 'max_lr': 0.0003, 'mixup_pct': 0.33, 'num_samples_per_epoch': 3750, 'num_test': 3000, 'ckpt_interval': 999, 'ckpt_saving': True, 'seed': 42, 'distributed': True, 'num_devices': 8, 'world_size': 1, 'train_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj08/train/{0..29}.tar', 'test_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj02/new_test/0.tar'}
+2024-06-18 03:54:07,335 INFO    MainThread:2245259 [wandb_init.py:init():610] starting backend
+2024-06-18 03:54:07,335 INFO    MainThread:2245259 [wandb_init.py:init():614] setting up manager
+2024-06-18 03:54:07,344 INFO    MainThread:2245259 [backend.py:_multiprocessing_setup():105] multiprocessing start_methods=fork,spawn,forkserver, using: spawn
+2024-06-18 03:54:07,352 INFO    MainThread:2245259 [wandb_init.py:init():622] backend started and connected
+2024-06-18 03:54:07,374 INFO    MainThread:2245259 [wandb_init.py:init():711] updated telemetry
+2024-06-18 03:54:07,461 INFO    MainThread:2245259 [wandb_init.py:init():744] communicating run to backend with 90.0 second timeout
+2024-06-18 03:54:07,695 ERROR   MainThread:2245259 [wandb_init.py:init():770] encountered error: The API key you provided is either invalid or missing.  If the `WANDB_API_KEY` environment variable is set, make sure it is correct. Otherwise, to resolve this issue, you may try running the 'wandb login --relogin' command. If you are using a local server, make sure that you're using the correct hostname. If you're not sure, you can try logging in again using the 'wandb login --relogin --host [hostname]' command.(Error 401: Unauthorized)
+2024-06-18 03:54:07,695 ERROR   MainThread:2245259 [wandb_init.py:init():1181] error in wandb.init()
+Traceback (most recent call last):
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/sdk/wandb_init.py", line 1171, in init
+    return wi.init()
+           ^^^^^^^^^
+  File "/admin/home-ckadirt/fmri/lib/python3.11/site-packages/wandb/sdk/wandb_init.py", line 776, in init
+    raise error
+wandb.errors.AuthenticationError: The API key you provided is either invalid or missing.  If the `WANDB_API_KEY` environment variable is set, make sure it is correct. Otherwise, to resolve this issue, you may try running the 'wandb login --relogin' command. If you are using a local server, make sure that you're using the correct hostname. If you're not sure, you can try logging in again using the 'wandb login --relogin --host [hostname]' command.(Error 401: Unauthorized)
+2024-06-18 03:54:08,847 WARNING MsgRouterThr:2245259 [router.py:message_loop():77] message_loop has been closed

MindEyeV2/src/wandb/run-20240618_121613-rr-test/files/code/src/TrainCluster.py

@@ -0,0 +1,1069 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# # Import packages & functions
+
+# In[1]:
+
+
+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 matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from accelerate import Accelerator
+
+# SDXL unCLIP requires code from https://github.com/Stability-AI/generative-models/tree/main
+sys.path.append('generative_models/')
+import sgm
+from generative_models.sgm.modules.encoders.modules import FrozenOpenCLIPImageEmbedder # bigG embedder
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+
+
+# In[2]:
+
+
+### 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)  
+
+data_type = torch.float16 # change depending on your mixed_precision
+num_devices = torch.cuda.device_count()
+if num_devices==0: num_devices = 1
+
+# First use "accelerate config" in terminal and setup using deepspeed stage 2 with CPU offloading!
+accelerator = Accelerator(split_batches=False, mixed_precision="fp16")
+if utils.is_interactive(): # set batch size here if using interactive notebook instead of submitting job
+    global_batch_size = batch_size = 8
+else:
+    global_batch_size = os.environ["GLOBAL_BATCH_SIZE"]
+    batch_size = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices
+
+
+# In[3]:
+
+
+print("PID of this process =",os.getpid())
+device = accelerator.device
+print("device:",device)
+world_size = accelerator.state.num_processes
+distributed = not accelerator.state.distributed_type == 'NO'
+num_devices = torch.cuda.device_count()
+if num_devices==0 or not distributed: num_devices = 1
+num_workers = num_devices
+print(accelerator.state)
+
+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[4]:
+
+
+# if running this interactively, can specify jupyter_args here for argparser to use
+if utils.is_interactive():
+    model_name = "testing"
+    print("model_name:", model_name)
+    
+    # global_batch_size and batch_size should already be defined in the 2nd cell block
+    jupyter_args = f"--data_path=/weka/proj-medarc/shared/mindeyev2_dataset \
+                    --cache_dir=/weka/proj-medarc/shared/cache \
+                    --model_name={model_name} \
+                    --no-multi_subject --subj=1 --batch_size={batch_size} --num_sessions=40 \
+                    --hidden_dim=1024 --clip_scale=1. \
+                    --no-blurry_recon --blur_scale=.5  \
+                    --use_prior --prior_scale=30 \
+                    --n_blocks=4 --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --no-use_image_aug \
+                    --ckpt_interval=999 --no-ckpt_saving --no-wandb_log"
+    # --multisubject_ckpt=../train_logs/multisubject_subj01_1024_24bs_nolow
+
+    print(jupyter_args)
+    jupyter_args = jupyter_args.split()
+    
+    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[5]:
+
+
+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=os.getcwd(),
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--cache_dir", type=str, default=os.getcwd(),
+    help="Path to where misc. files downloaded from huggingface are stored. Defaults to current src directory.",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,3,4,5,6,7,8],
+    help="Validate on which subject?",
+)
+parser.add_argument(
+    "--multisubject_ckpt", type=str, default=None,
+    help="Path to pre-trained multisubject model to finetune a single subject from. multisubject must be False.",
+)
+parser.add_argument(
+    "--num_sessions", type=int, default=1,
+    help="Number of training sessions to include",
+)
+parser.add_argument(
+    "--use_prior",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to train diffusion prior (True) or just rely on retrieval part of the pipeline (False)",
+)
+parser.add_argument(
+    "--batch_size", type=int, default=16,
+    help="Batch size can be increased by 10x if only training retreival submodule and not diffusion prior",
+)
+parser.add_argument(
+    "--wandb_log",action=argparse.BooleanOptionalAction,default=False,
+    help="whether to log to wandb",
+)
+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(
+    "--blur_scale",type=float,default=.5,
+    help="multiply loss from blurry recons by this number",
+)
+parser.add_argument(
+    "--clip_scale",type=float,default=1.,
+    help="multiply contrastive loss by this number",
+)
+parser.add_argument(
+    "--prior_scale",type=float,default=30,
+    help="multiply diffusion prior loss by this",
+)
+parser.add_argument(
+    "--use_image_aug",action=argparse.BooleanOptionalAction,default=False,
+    help="whether to use image augmentation",
+)
+parser.add_argument(
+    "--num_epochs",type=int,default=150,
+    help="number of epochs of training",
+)
+parser.add_argument(
+    "--multi_subject",action=argparse.BooleanOptionalAction,default=False,
+)
+parser.add_argument(
+    "--new_test",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--n_blocks",type=int,default=4,
+)
+parser.add_argument(
+    "--hidden_dim",type=int,default=1024,
+)
+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)
+    
+# seed all random functions
+utils.seed_everything(seed)
+
+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 or blurry_recon:
+    import kornia
+    from kornia.augmentation.container import AugmentationSequential
+if use_image_aug:
+    img_augment = AugmentationSequential(
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1, p=0.3),
+        same_on_batch=False,
+        data_keys=["input"],
+    )
+    
+if multi_subject:
+    subj_list = np.arange(1,9)
+    subj_list = subj_list[subj_list != subj]
+else:
+    subj_list = [subj]
+
+print("subj_list", subj_list, "num_sessions", num_sessions)
+
+
+# # Prep data, models, and dataloaders
+
+# ### Creating wds dataloader, preload betas and all 73k possible images
+
+# In[6]:
+
+
+def my_split_by_node(urls): return urls
+num_voxels_list = []
+
+if multi_subject:
+    nsessions_allsubj=np.array([40, 40, 32, 30, 40, 32, 40, 30])
+    num_samples_per_epoch = (750*40) // num_devices 
+else:
+    num_samples_per_epoch = (750*num_sessions) // num_devices 
+
+print("dividing batch size by subj_list, which will then be concatenated across subj during training...") 
+batch_size = batch_size // len(subj_list)
+
+num_iterations_per_epoch = num_samples_per_epoch // (batch_size*len(subj_list))
+
+print("batch_size =", batch_size, "num_iterations_per_epoch =",num_iterations_per_epoch, "num_samples_per_epoch =",num_samples_per_epoch)
+
+
+# In[7]:
+
+
+train_data = {}
+train_dl = {}
+num_voxels = {}
+voxels = {}
+for s in subj_list:
+    print(f"Training with {num_sessions} sessions")
+    if multi_subject:
+        train_url = f"{data_path}/wds/subj0{s}/train/" + "{0.." + f"{nsessions_allsubj[s-1]-1}" + "}.tar"
+    else:
+        train_url = f"{data_path}/wds/subj0{s}/train/" + "{0.." + f"{num_sessions-1}" + "}.tar"
+    print(train_url)
+    
+    train_data[f'subj0{s}'] = wds.WebDataset(train_url,resampled=True,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[f'subj0{s}'] = torch.utils.data.DataLoader(train_data[f'subj0{s}'], batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+    f = h5py.File(f'{data_path}/betas_all_subj0{s}_fp32_renorm.hdf5', 'r')
+    betas = f['betas'][:]
+    betas = torch.Tensor(betas).to("cpu").to(data_type)
+    num_voxels_list.append(betas[0].shape[-1])
+    num_voxels[f'subj0{s}'] = betas[0].shape[-1]
+    voxels[f'subj0{s}'] = betas
+    print(f"num_voxels for subj0{s}: {num_voxels[f'subj0{s}']}")
+
+print("Loaded all subj train dls and betas!\n")
+
+# Validate only on one subject
+if multi_subject: 
+    subj = subj_list[0] # cant validate on the actual held out person so picking first in subj_list
+if not new_test: # using old test set from before full dataset released (used in original MindEye paper)
+    if subj==3:
+        num_test=2113
+    elif subj==4:
+        num_test=1985
+    elif subj==6:
+        num_test=2113
+    elif subj==8:
+        num_test=1985
+    else:
+        num_test=2770
+    test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+elif new_test: # using larger test set from after full dataset released
+    if subj==3:
+        num_test=2371
+    elif subj==4:
+        num_test=2188
+    elif subj==6:
+        num_test=2371
+    elif subj==8:
+        num_test=2188
+    else:
+        num_test=3000
+    test_url = f"{data_path}/wds/subj0{subj}/new_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=num_test, shuffle=False, drop_last=True, pin_memory=True)
+print(f"Loaded test dl for subj{subj}!\n")
+
+
+# In[8]:
+
+
+# Load 73k NSD images
+f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+images = f['images']
+print("Loaded all 73k possible NSD images to cpu!", images.shape)
+
+
+# ## Load models
+
+# ### CLIP image embeddings  model
+
+# In[9]:
+
+
+clip_img_embedder = FrozenOpenCLIPImageEmbedder(
+    arch="ViT-bigG-14",
+    version="laion2b_s39b_b160k",
+    output_tokens=True,
+    only_tokens=True,
+)
+clip_img_embedder.to(device)
+
+clip_seq_dim = 256
+clip_emb_dim = 1664
+
+
+# ### SD VAE
+
+# In[10]:
+
+
+if blurry_recon:
+    from diffusers import AutoencoderKL    
+    autoenc = AutoencoderKL(
+        down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D'],
+        up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D'],
+        block_out_channels=[128, 256, 512, 512],
+        layers_per_block=2,
+        sample_size=256,
+    )
+    ckpt = torch.load(f'{cache_dir}/sd_image_var_autoenc.pth')
+    autoenc.load_state_dict(ckpt)
+    
+    autoenc.eval()
+    autoenc.requires_grad_(False)
+    autoenc.to(device)
+    utils.count_params(autoenc)
+    
+    from autoencoder.convnext import ConvnextXL
+    cnx = ConvnextXL(f'{cache_dir}/convnext_xlarge_alpha0.75_fullckpt.pth')
+    cnx.requires_grad_(False)
+    cnx.eval()
+    cnx.to(device)
+    
+    mean = torch.tensor([0.485, 0.456, 0.406]).to(device).reshape(1,3,1,1)
+    std = torch.tensor([0.228, 0.224, 0.225]).to(device).reshape(1,3,1,1)
+    
+    blur_augs = AugmentationSequential(
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.8),
+        kornia.augmentation.RandomGrayscale(p=0.1),
+        kornia.augmentation.RandomSolarize(p=0.1),
+        kornia.augmentation.RandomResizedCrop((224,224), scale=(.9,.9), ratio=(1,1), p=1.0),
+        data_keys=["input"],
+    )
+
+
+# ### MindEye modules
+
+# In[11]:
+
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+model
+
+
+# In[12]:
+
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer to enable regularization
+    def __init__(self, input_sizes, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linears = torch.nn.ModuleList([
+                torch.nn.Linear(input_size, out_features) for input_size in input_sizes
+            ])
+    def forward(self, x, subj_idx):
+        out = self.linears[subj_idx](x[:,0]).unsqueeze(1)
+        return out
+        
+class IndividRidgeRegression(torch.nn.Module):
+    def __init__(self, input_size, out_features):
+        super(IndividRidgeRegression, self).__init__()
+        self.out_features = out_features
+        self.linear = torch.nn.Linear(input_size, out_features)
+    def forward(self, x):
+        out = self.linear(x)
+        return out
+    
+model.ridge = RidgeRegression(num_voxels_list, out_features=hidden_dim)
+utils.count_params(model.ridge)
+utils.count_params(model)
+
+# test on subject 1 with fake data
+b = torch.randn((2,1,num_voxels_list[0]))
+print(b.shape, model.ridge(b,0).shape)
+
+
+# In[13]:
+
+
+from models import BrainNetwork
+model.backbone = BrainNetwork(h=hidden_dim, in_dim=hidden_dim, seq_len=1, n_blocks=n_blocks,
+                          clip_size=clip_emb_dim, out_dim=clip_emb_dim*clip_seq_dim, 
+                          blurry_recon=blurry_recon, clip_scale=clip_scale)
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+# test that the model works on some fake data
+b = torch.randn((2,1,hidden_dim))
+print("b.shape",b.shape)
+
+backbone_, clip_, blur_ = model.backbone(b)
+print(backbone_.shape, clip_.shape, blur_[0].shape, blur_[1].shape)
+
+
+# ### Load semantic clusters
+
+# In[14]:
+
+
+path_semantic_names = "/weka/proj-medarc/shared/mindeyev2_dataset/semantic_cluster_names.npy"
+path_semantic_cluster = "/weka/proj-medarc/shared/mindeyev2_dataset/COCO_73k_semantic_cluster.npy"
+semantic_cluster_names = np.load(path_semantic_names)
+semantic_cluster = np.load(path_semantic_cluster)
+possible_semantic_clusters = np.unique(semantic_cluster)
+
+# one-hot encode semantic clusters
+# move possible_semantic_clusters to numbers and create a dictionary
+semantic_cluster_dict = {cluster: i for i, cluster in enumerate(possible_semantic_clusters)}
+semantic_cluster_onehot = torch.zeros((len(semantic_cluster), len(possible_semantic_clusters)))
+for i, cluster in enumerate(semantic_cluster):
+    semantic_cluster_onehot[i, semantic_cluster_dict[cluster]] = 1
+
+
+print("semantic_cluster_onehot.shape", semantic_cluster_onehot.shape)
+
+num_seman_clusters = len(np.unique(semantic_cluster))
+print("num_seman_clusters", num_seman_clusters)
+
+
+# ### Adding the ridge regression to the class
+
+# In[15]:
+
+
+# if use_prior:
+#     from models import *
+
+#     # setup diffusion prior network
+#     out_dim = clip_emb_dim
+#     depth = 6
+#     dim_head = 52
+#     heads = clip_emb_dim//52 # heads * dim_head = clip_emb_dim
+#     timesteps = 100
+
+#     prior_network = PriorNetwork(
+#             dim=out_dim,
+#             depth=depth,
+#             dim_head=dim_head,
+#             heads=heads,
+#             causal=False,
+#             num_tokens = clip_seq_dim,
+#             learned_query_mode="pos_emb"
+#         )
+
+#     model.diffusion_prior = BrainDiffusionPrior(
+#         net=prior_network,
+#         image_embed_dim=out_dim,
+#         condition_on_text_encodings=False,
+#         timesteps=timesteps,
+#         cond_drop_prob=0.2,
+#         image_embed_scale=None,
+#     )
+    
+#     utils.count_params(model.diffusion_prior)
+#     utils.count_params(model)
+
+model.RRClassifier = IndividRidgeRegression(clip_emb_dim*clip_seq_dim, out_features=num_seman_clusters)
+utils.count_params(model.RRClassifier)
+utils.count_params(model)
+
+
+# ### Setup optimizer / lr / ckpt saving
+
+# In[16]:
+
+
+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},
+]
+# if use_prior:
+#     opt_grouped_parameters.extend([
+#         {'params': [p for n, p in model.diffusion_prior.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},
+#         {'params': [p for n, p in model.diffusion_prior.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
+#     ])
+opt_grouped_parameters.extend([
+    {'params': [p for n, p in model.RRClassifier.named_parameters()], 'weight_decay': 1e-2},
+])
+
+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_iterations_per_epoch)),
+        last_epoch=-1
+    )
+elif lr_scheduler_type == 'cycle':
+    total_steps=int(np.floor(num_epochs*num_iterations_per_epoch))
+    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'
+    if accelerator.is_main_process:
+        unwrapped_model = accelerator.unwrap_model(model)
+        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)
+    print(f"\n---saved {outdir}/{tag} ckpt!---\n")
+
+def load_ckpt(tag,load_lr=True,load_optimizer=True,load_epoch=True,strict=True,outdir=outdir,multisubj_loading=False): 
+    print(f"\n---loading {outdir}/{tag}.pth ckpt---\n")
+    checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+    state_dict = checkpoint['model_state_dict']
+    if multisubj_loading: # remove incompatible ridge layer that will otherwise error
+        state_dict.pop('ridge.linears.0.weight',None)
+    model.load_state_dict(state_dict, strict=strict)
+    if load_epoch:
+        globals()["epoch"] = checkpoint['epoch']
+        print("Epoch",epoch)
+    if load_optimizer:
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    if load_lr:
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    del checkpoint
+
+print("\nDone with model preparations!")
+num_params = utils.count_params(model)
+
+
+# # Weights and Biases
+
+# In[17]:
+
+
+if local_rank==0 and wandb_log: # only use main process for wandb logging
+    import wandb
+    wandb_project = 'mindeye_semantic_cluster'
+    print(f"wandb {wandb_project} run {model_name}")
+    # need to configure wandb beforehand in terminal with "wandb init"!
+    wandb_config = {
+      "model_name": model_name,
+      "global_batch_size": global_batch_size,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "num_sessions": num_sessions,
+      "num_params": num_params,
+      "clip_scale": clip_scale,
+      "prior_scale": prior_scale,
+      "blur_scale": blur_scale,
+      "use_image_aug": use_image_aug,
+      "max_lr": max_lr,
+      "mixup_pct": mixup_pct,
+      "num_samples_per_epoch": num_samples_per_epoch,
+      "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)
+    print("wandb_id:",model_name)
+    wandb.init(
+        id=model_name,
+        project=wandb_project,
+        name=model_name,
+        config=wandb_config,
+        resume="allow",
+    )
+else:
+    wandb_log = False
+
+
+# # Main
+
+# In[18]:
+
+
+epoch = 0
+losses, test_losses, lrs = [], [], []
+best_test_loss = 1e9
+torch.cuda.empty_cache()
+
+
+# In[19]:
+
+
+# load multisubject stage1 ckpt if set
+if multisubject_ckpt is not None:
+    load_ckpt("last",outdir=multisubject_ckpt,load_lr=False,load_optimizer=False,load_epoch=False,strict=False,multisubj_loading=True)
+
+
+# In[20]:
+
+
+train_dls = [train_dl[f'subj0{s}'] for s in subj_list]
+
+model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot = accelerator.prepare(model, optimizer, *train_dls, lr_scheduler, semantic_cluster_onehot)
+# leaving out test_dl since we will only have local_rank 0 device do evals
+
+
+# In[21]:
+
+
+print(num_iterations_per_epoch)
+
+
+# In[26]:
+
+
+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()
+soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))
+
+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.
+    
+    recon_cossim = 0.
+    test_recon_cossim = 0.
+    recon_mse = 0.
+    test_recon_mse = 0.
+
+    loss_clip_total = 0.
+    loss_blurry_total = 0.
+    loss_blurry_cont_total = 0.
+    test_loss_clip_total = 0.
+    
+    loss_prior_total = 0.
+    test_loss_prior_total = 0.
+    
+    loss_RR_total = 0.
+    test_loss_RR_total = 0.
+
+    blurry_pixcorr = 0.
+    test_blurry_pixcorr = 0. # needs >.456 to beat low-level subj01 results in mindeye v1
+
+    # pre-load all batches for this epoch (it's MUCH faster to pre-load in bulk than to separate loading per batch)
+    voxel_iters = {} # empty dict because diff subjects have differing # of voxels
+    image_iters = torch.zeros(num_iterations_per_epoch, batch_size*len(subj_list), 3, 224, 224).float()
+    annot_iters = {}
+    perm_iters, betas_iters, select_iters = {}, {}, {}
+    images_indexes = {}
+    for s, train_dl in enumerate(train_dls):
+        with torch.cuda.amp.autocast(dtype=data_type):
+            iter = -1
+            for behav0, past_behav0, future_behav0, old_behav0 in train_dl: 
+                # Load images to cpu from hdf5 (requires sorted indexing)
+                image_idx = behav0[:,0,0].cpu().long().numpy()
+
+                image0, image_sorted_idx = np.unique(image_idx, return_index=True)                
+                if len(image0) != len(image_idx): # hdf5 cant handle duplicate indexing
+                    continue
+                iter += 1
+                image0 = torch.tensor(images[image0], dtype=data_type)
+                image_iters[iter,s*batch_size:s*batch_size+batch_size] = image0
+                images_indexes[f"subj0{s}_iter{iter}"] = image_sorted_idx
+                
+                # Load voxels for current batch, matching above indexing
+                voxel_idx = behav0[:,0,5].cpu().long().numpy()
+                voxel_sorted_idx = voxel_idx[image_sorted_idx]
+                voxel0 = voxels[f'subj0{subj_list[s]}'][voxel_sorted_idx]
+                voxel0 = torch.Tensor(voxel0).unsqueeze(1)
+
+                if epoch < int(mixup_pct * num_epochs):
+                    voxel0, perm, betas, select = utils.mixco(voxel0)
+                    perm_iters[f"subj0{subj_list[s]}_iter{iter}"] = perm
+                    betas_iters[f"subj0{subj_list[s]}_iter{iter}"] = betas
+                    select_iters[f"subj0{subj_list[s]}_iter{iter}"] = select
+
+                voxel_iters[f"subj0{subj_list[s]}_iter{iter}"] = voxel0
+
+                if iter >= num_iterations_per_epoch-1:
+                    break
+
+    # you now have voxel_iters and image_iters with num_iterations_per_epoch batches each
+    for train_i in range(num_iterations_per_epoch):
+        with torch.cuda.amp.autocast(dtype=data_type):
+            optimizer.zero_grad()
+            loss=0.
+
+            voxel_list = [voxel_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list]
+            image = image_iters[train_i].detach()
+            image = image.to(device)
+
+            if use_image_aug: 
+                image = img_augment(image)
+
+            clip_target = clip_img_embedder(image)
+            assert not torch.any(torch.isnan(clip_target))
+
+            if epoch < int(mixup_pct * num_epochs):
+                perm_list = [perm_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list]
+                perm = torch.cat(perm_list, dim=0)
+                betas_list = [betas_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list]
+                betas = torch.cat(betas_list, dim=0)
+                select_list = [select_iters[f"subj0{s}_iter{train_i}"].detach().to(device) for s in subj_list]
+                select = torch.cat(select_list, dim=0)
+
+            voxel_ridge_list = [model.ridge(voxel_list[si],si) for si,s in enumerate(subj_list)]
+            voxel_ridge = torch.cat(voxel_ridge_list, dim=0)
+
+            backbone, clip_voxels, blurry_image_enc_ = model.backbone(voxel_ridge)
+
+            if clip_scale>0:
+                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 use_prior:
+            #     loss_prior, prior_out = model.diffusion_prior(text_embed=backbone, image_embed=clip_target)
+            #     loss_prior_total += loss_prior.item()
+            #     loss_prior *= prior_scale
+            #     loss += loss_prior
+
+            #     recon_cossim += nn.functional.cosine_similarity(prior_out, clip_target).mean().item()
+            #     recon_mse += mse(prior_out, clip_target).item()
+
+            logits = model.RRClassifier(backbone.flatten(1))
+            loss_RR = nn.functional.cross_entropy(logits, torch.argmax(semantic_cluster_onehot[images_indexes[f"subj0{s}_iter{train_i}"]], dim=1).to(logits.device))
+            #print("backbone.shape",backbone.shape, "clip_voxels.shape",clip_voxels.shape, "blurry_image_enc_[0].shape",blurry_image_enc_[0].shape, "blurry_image_enc_[1].shape",blurry_image_enc_[1].shape)
+            #something 
+
+            loss_RR_total += loss_RR.item()
+            loss += loss_RR
+
+            if clip_scale>0:
+                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:     
+                image_enc_pred, transformer_feats = blurry_image_enc_
+
+                image_enc = autoenc.encode(2*image-1).latent_dist.mode() * 0.18215
+                loss_blurry = l1(image_enc_pred, image_enc)
+                loss_blurry_total += loss_blurry.item()
+
+                if epoch < int(mixup_pct * num_epochs):
+                    image_enc_shuf = image_enc[perm]
+                    betas_shape = [-1] + [1]*(len(image_enc.shape)-1)
+                    image_enc[select] = image_enc[select] * betas[select].reshape(*betas_shape) + \
+                        image_enc_shuf[select] * (1 - betas[select]).reshape(*betas_shape)
+
+                image_norm = (image - mean)/std
+                image_aug = (blur_augs(image) - mean)/std
+                _, cnx_embeds = cnx(image_norm)
+                _, cnx_aug_embeds = cnx(image_aug)
+
+                cont_loss = utils.soft_cont_loss(
+                    nn.functional.normalize(transformer_feats.reshape(-1, transformer_feats.shape[-1]), dim=-1),
+                    nn.functional.normalize(cnx_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),
+                    nn.functional.normalize(cnx_aug_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),
+                    temp=0.2)
+                loss_blurry_cont_total += cont_loss.item()
+
+                loss += (loss_blurry + 0.1*cont_loss) * blur_scale #/.18215
+
+            if clip_scale>0:
+                # forward and backward top 1 accuracy        
+                labels = torch.arange(len(clip_voxels_norm)).to(clip_voxels_norm.device) 
+                fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1).item()
+                bwd_percent_correct += utils.topk(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(image)), size=len(image)//5, replace=False)
+                    blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample/ 2 + 0.5).clamp(0,1)
+                    pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)
+                    blurry_pixcorr += pixcorr.item()
+
+            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[f'subj0{subj}'][behav[:,0,5].cpu().long()].unsqueeze(1)
+                    
+                    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 len(locs)==1:
+                            locs = locs.repeat(3)
+                        elif len(locs)==2:
+                            locs = locs.repeat(2)[:3]
+                        assert len(locs)==3
+                        if test_image is None:
+                            test_image = torch.Tensor(images[im][None])
+                            test_voxel = voxel[locs][None]
+                        else:
+                            test_image = torch.vstack((test_image, torch.Tensor(images[im][None])))
+                            test_voxel = torch.vstack((test_voxel, voxel[locs][None]))
+
+                loss=0.
+                            
+                test_indices = torch.arange(len(test_voxel))[:300]
+                voxel = test_voxel[test_indices].to(device)
+                image = test_image[test_indices].to(device)
+                assert len(image) == 300
+
+                clip_target = clip_img_embedder(image.float())
+
+                for rep in range(3):
+                    voxel_ridge = model.ridge(voxel[:,rep],0) # 0th index of subj_list
+                    backbone0, clip_voxels0, blurry_image_enc_ = model.backbone(voxel_ridge)
+
+                    logits0 = model.RRClassifier(backbone0.flatten(1))
+
+                    if rep==0:
+                        clip_voxels = clip_voxels0
+                        backbone = backbone0
+                        logits = logits0
+                    else:
+                        clip_voxels += clip_voxels0
+                        backbone += backbone0
+                        logits += logits0
+                clip_voxels /= 3
+                backbone /= 3
+                logits /= 3
+
+                print(logits.shape, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).shape)
+                RR_loss = nn.functional.cross_entropy(logits, torch.argmax(semantic_cluster_onehot[test_indices], dim=1).to(logits.device))
+                test_loss_RR_total += RR_loss.item()
+                loss += RR_loss
+
+                if clip_scale>0:
+                    clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+                    clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+                
+                # for some evals, only doing a subset of the samples per batch because of computational cost
+                random_samps = np.random.choice(np.arange(len(image)), size=len(image)//5, replace=False)
+                
+                # if use_prior:
+                #     loss_prior, contaminated_prior_out = model.diffusion_prior(text_embed=backbone[random_samps], image_embed=clip_target[random_samps])
+                #     test_loss_prior_total += loss_prior.item()
+                #     loss_prior *= prior_scale
+                #     loss += loss_prior
+                    
+                if clip_scale>0:
+                    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:
+                    image_enc_pred, _ = blurry_image_enc_
+                    blurry_recon_images = (autoenc.decode(image_enc_pred[random_samps]/0.18215).sample / 2 + 0.5).clamp(0,1)
+                    pixcorr = utils.pixcorr(image[random_samps], blurry_recon_images)
+                    test_blurry_pixcorr += pixcorr.item()
+
+                if clip_scale>0:
+                    # forward and backward top 1 accuracy        
+                    labels = torch.arange(len(clip_voxels_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())
+
+            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),
+                "train/loss_blurry_cont_total": loss_blurry_cont_total / (train_i + 1),
+                "test/loss_clip_total": test_loss_clip_total / (test_i + 1),
+                "train/blurry_pixcorr": blurry_pixcorr / (train_i + 1),
+                "test/blurry_pixcorr": test_blurry_pixcorr / (test_i + 1),
+                "train/recon_cossim": recon_cossim / (train_i + 1),
+                "test/recon_cossim": test_recon_cossim / (test_i + 1),
+                "train/recon_mse": recon_mse / (train_i + 1),
+                "test/recon_mse": test_recon_mse / (test_i + 1),
+                "train/loss_prior": loss_prior_total / (train_i + 1),
+                "test/loss_prior": test_loss_prior_total / (test_i + 1),
+                "train/loss_RR": loss_RR_total / (train_i + 1),
+                "test/loss_RR": test_loss_RR_total / (test_i + 1),
+                }
+
+            # if finished training, save jpg recons if they exist
+            if (epoch == num_epochs-1) or (epoch % ckpt_interval == 0):
+                if blurry_recon:    
+                    image_enc = autoenc.encode(2*image[:4]-1).latent_dist.mode() * 0.18215
+                    # 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(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(image_enc_pred[[j]]/0.18215).sample / 2 + 0.5).clamp(0,1)))
+                        axes[jj].axis('off')
+
+                    if wandb_log:
+                        logs[f"test/blur_recons"] = wandb.Image(fig, caption=f"epoch{epoch:03d}")
+                        plt.close()
+                    else:
+                        plt.show()
+
+            progress_bar.set_postfix(**logs)
+
+            if wandb_log: wandb.log(logs)
+            
+    # Save model checkpoint and reconstruct
+    if (ckpt_saving) and (epoch % ckpt_interval == 0):
+        save_ckpt(f'last')
+
+    # wait for other GPUs to catch up if needed
+    accelerator.wait_for_everyone()
+    torch.cuda.empty_cache()
+
+print("\n===Finished!===\n")
+if ckpt_saving:
+    save_ckpt(f'last')
+
+
+# In[ ]:
+
+
+plt.plot(losses)
+plt.show()
+plt.plot(test_losses)
+plt.show()
+
+
+# In[29]:
+
+
+import wandb
+wandb.login()
+

MindEyeV2/src/wandb/run-20240618_121613-rr-test/files/config.yaml

@@ -0,0 +1,103 @@
+wandb_version: 1
+
+model_name:
+  desc: null
+  value: rr-test
+global_batch_size:
+  desc: null
+  value: '21'
+batch_size:
+  desc: null
+  value: 3
+num_epochs:
+  desc: null
+  value: 150
+num_sessions:
+  desc: null
+  value: 40
+num_params:
+  desc: null
+  value: 573919937
+clip_scale:
+  desc: null
+  value: 1.0
+prior_scale:
+  desc: null
+  value: 30.0
+blur_scale:
+  desc: null
+  value: 0.5
+use_image_aug:
+  desc: null
+  value: false
+max_lr:
+  desc: null
+  value: 0.0003
+mixup_pct:
+  desc: null
+  value: 0.33
+num_samples_per_epoch:
+  desc: null
+  value: 3750
+num_test:
+  desc: null
+  value: 3000
+ckpt_interval:
+  desc: null
+  value: 999
+ckpt_saving:
+  desc: null
+  value: true
+seed:
+  desc: null
+  value: 42
+distributed:
+  desc: null
+  value: true
+num_devices:
+  desc: null
+  value: 8
+world_size:
+  desc: null
+  value: 1
+train_url:
+  desc: null
+  value: /weka/proj-medarc/shared/mindeyev2_dataset/wds/subj08/train/{0..29}.tar
+test_url:
+  desc: null
+  value: /weka/proj-medarc/shared/mindeyev2_dataset/wds/subj02/new_test/0.tar
+_wandb:
+  desc: null
+  value:
+    code_path: code/src/TrainCluster.py
+    python_version: 3.11.9
+    cli_version: 0.17.1
+    framework: huggingface
+    huggingface_version: 4.37.2
+    is_jupyter_run: false
+    is_kaggle_kernel: false
+    start_time: 1718712973
+    t:
+      1:
+      - 1
+      - 9
+      - 11
+      - 41
+      - 49
+      - 55
+      - 63
+      - 71
+      - 79
+      - 83
+      - 103
+      3:
+      - 13
+      - 14
+      - 16
+      - 23
+      4: 3.11.9
+      5: 0.17.1
+      6: 4.37.2
+      8:
+      - 5
+      13: linux-x86_64

MindEyeV2/src/wandb/run-20240618_121613-rr-test/files/diff.patch

@@ -0,0 +1,230 @@
+diff --git a/src/Train.ipynb b/src/Train.ipynb
+index 91922c8..396549d 100644
+--- a/src/Train.ipynb
++++ b/src/Train.ipynb
+@@ -99,7 +99,7 @@
+      "name": "stdout",
+      "output_type": "stream",
+      "text": [
+-      "PID of this process = 3001293\n",
++      "PID of this process = 3520819\n",
+       "device: cuda\n",
+       "Distributed environment: DistributedType.NO\n",
+       "Num processes: 1\n",
+@@ -1362,9 +1362,9 @@
+  ],
+  "metadata": {
+   "kernelspec": {
+-   "display_name": "mindeye",
++   "display_name": "Python 3 (ipykernel)",
+    "language": "python",
+-   "name": "mindeye"
++   "name": "python3"
+   },
+   "language_info": {
+    "codemirror_mode": {
+diff --git a/src/accel.slurm b/src/accel.slurm
+index 1c1dfbc..5623c78 100644
+--- a/src/accel.slurm
++++ b/src/accel.slurm
+@@ -1,7 +1,7 @@
+ #!/bin/bash
+-#SBATCH --account=topfmri
+-#SBATCH --partition=a40x
+-#SBATCH --job-name=eye
++#SBATCH --account=fmri
++#SBATCH --partition=p5
++#SBATCH --job-name=eyeRR
+ #SBATCH --nodes=1           
+ #SBATCH --gres=gpu:1
+ #SBATCH --time=48:00:00          # total run time limit (HH:MM:SS)
+@@ -12,11 +12,11 @@
+ #SBATCH --exclusive
+ 
+ # Make sure you activate your fmri environment created from src/setup.sh
+-cd /weka/proj-fmri/paulscotti/MindEyeV2/src
+-source fmri/bin/activate
++cd /weka/proj-fmri/ckadirt/MindEyeV2/src
++source /admin/home-ckadirt/fmri/bin/activate
+ 
+ # The following line converts your jupyter notebook into a python script runnable with Slurm
+-jupyter nbconvert Train.ipynb --to python
++jupyter nbconvert TrainCluster.ipynb --to python
+ 
+ export NUM_GPUS=1  # Set to equal gres=gpu:#!
+ export BATCH_SIZE=21 # 21 for multisubject / 24 for singlesubject (orig. paper used 42 for multisubject / 24 for singlesubject)
+@@ -32,9 +32,9 @@ echo MASTER_PORT=${MASTER_PORT}
+ echo WORLD_SIZE=${COUNT_NODE}
+ 
+ # multisubject pretraining
+-model_name="multisubject_excludingsubj01_40sess"
++model_name="rr-test"
+ echo model_name=${model_name}
+-accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --mixed_precision=fp16 Train.py --data_path=/weka/proj-fmri/shared/mindeyev2_dataset --cache_dir=/weka/proj-fmri/shared/cache --model_name=${model_name} --multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --use_prior --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log
++accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --mixed_precision=fp16 TrainCluster.py --data_path=/weka/proj-medarc/shared/mindeyev2_dataset --cache_dir=/weka/proj-medarc/shared/cache --model_name=${model_name} --multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log
+ 
+ # singlesubject finetuning
+ #model_name="finetuned_subj01_40sess"
+diff --git a/src/recon_inference.ipynb b/src/recon_inference.ipynb
+index 10d7e27..0593e7e 100644
+--- a/src/recon_inference.ipynb
++++ b/src/recon_inference.ipynb
+@@ -12,7 +12,6 @@
+      "name": "stdout",
+      "output_type": "stream",
+      "text": [
+-      "LOCAL RANK  0\n",
+       "device: cuda\n"
+      ]
+     }
+@@ -60,7 +59,7 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 33,
++   "execution_count": 2,
+    "id": "e52985b1-95ff-487b-8b2d-cc1ad1c190b8",
+    "metadata": {
+     "tags": []
+@@ -71,9 +70,7 @@
+      "output_type": "stream",
+      "text": [
+       "model_name: final_subj01_pretrained_40sess_24bs\n",
+-      "--data_path=/weka/proj-medarc/shared/mindeyev2_dataset                     --cache_dir=/weka/proj-medarc/shared/mindeyev2_dataset                     --model_name=final_subj01_pretrained_40sess_24bs --subj=1                     --hidden_dim=4096 --n_blocks=4 --new_test\n",
+-      "The autoreload extension is already loaded. To reload it, use:\n",
+-      "  %reload_ext autoreload\n"
++      "--data_path=/weka/proj-medarc/shared/mindeyev2_dataset                     --cache_dir=/weka/proj-medarc/shared/mindeyev2_dataset                     --model_name=final_subj01_pretrained_40sess_24bs --subj=1                     --hidden_dim=4096 --n_blocks=4 --new_test\n"
+      ]
+     }
+    ],
+@@ -99,7 +96,7 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 34,
++   "execution_count": 3,
+    "id": "49e5dae4-606d-4dc6-b420-df9e4c14737e",
+    "metadata": {
+     "tags": []
+@@ -157,7 +154,7 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 35,
++   "execution_count": 4,
+    "id": "64672583-9f00-46f5-8d4e-00e4c7068a1d",
+    "metadata": {
+     "tags": []
+@@ -221,7 +218,7 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 36,
++   "execution_count": 5,
+    "id": "a3cbeea8-e95b-48d9-9bc2-91af260c93d1",
+    "metadata": {},
+    "outputs": [
+@@ -254,39 +251,10 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 38,
++   "execution_count": null,
+    "id": "3afc4858-b6a6-4a52-9303-b4a50ea5cc0b",
+    "metadata": {},
+-   "outputs": [
+-    {
+-     "name": "stdout",
+-     "output_type": "stream",
+-     "text": [
+-      "param counts:\n",
+-      "83,653,863 total\n",
+-      "0 trainable\n",
+-      "param counts:\n",
+-      "64,409,600 total\n",
+-      "64,409,600 trainable\n",
+-      "param counts:\n",
+-      "1,903,020,028 total\n",
+-      "1,903,020,028 trainable\n",
+-      "param counts:\n",
+-      "1,967,429,628 total\n",
+-      "1,967,429,628 trainable\n",
+-      "param counts:\n",
+-      "259,865,216 total\n",
+-      "259,865,200 trainable\n",
+-      "param counts:\n",
+-      "2,227,294,844 total\n",
+-      "2,227,294,828 trainable\n",
+-      "\n",
+-      "---loading /weka/proj-fmri/paulscotti/MindEye2_git/train_logs/final_subj01_pretrained_40sess_24bs/last.pth ckpt---\n",
+-      "\n",
+-      "ckpt loaded!\n"
+-     ]
+-    }
+-   ],
++   "outputs": [],
+    "source": [
+     "clip_img_embedder = FrozenOpenCLIPImageEmbedder(\n",
+     "    arch=\"ViT-bigG-14\",\n",
+@@ -393,7 +361,7 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 30,
++   "execution_count": null,
+    "id": "295824db-ab3d-450c-90fb-f656e48994ba",
+    "metadata": {
+     "tags": []
+@@ -430,38 +398,10 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 31,
++   "execution_count": null,
+    "id": "f726f617-39f5-49e2-8d0c-d11d27d01c30",
+    "metadata": {},
+-   "outputs": [
+-    {
+-     "name": "stderr",
+-     "output_type": "stream",
+-     "text": [
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now.\n"
+-     ]
+-    },
+-    {
+-     "name": "stdout",
+-     "output_type": "stream",
+-     "text": [
+-      "Initialized embedder #0: FrozenOpenCLIPImageEmbedder with 1909889025 params. Trainable: False\n",
+-      "Initialized embedder #1: ConcatTimestepEmbedderND with 0 params. Trainable: False\n",
+-      "Initialized embedder #2: ConcatTimestepEmbedderND with 0 params. Trainable: False\n",
+-      "vector_suffix torch.Size([1, 1024])\n"
+-     ]
+-    }
+-   ],
++   "outputs": [],
+    "source": [
+     "# prep unCLIP\n",
+     "config = OmegaConf.load(\"generative_models/configs/unclip6.yaml\")\n",
+@@ -715,9 +655,9 @@
+  ],
+  "metadata": {
+   "kernelspec": {
+-   "display_name": "mindeye",
++   "display_name": "Python 3 (ipykernel)",
+    "language": "python",
+-   "name": "mindeye"
++   "name": "python3"
+   },
+   "language_info": {
+    "codemirror_mode": {

MindEyeV2/src/wandb/run-20240618_121613-rr-test/files/requirements.txt

@@ -0,0 +1,214 @@
+Babel==2.15.0
+CoCa-pytorch==0.1.0
+GitPython==3.1.43
+Jinja2==3.1.4
+MarkupSafe==2.1.5
+PyYAML==6.0.1
+Pygments==2.18.0
+QtPy==2.4.1
+Send2Trash==1.8.3
+absl-py==2.1.0
+accelerate==0.24.1
+aiohttp==3.9.5
+aiosignal==1.3.1
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+anyio==4.4.0
+argon2-cffi-bindings==21.2.0
+argon2-cffi==23.1.0
+arrow==1.3.0
+asttokens==2.4.1
+async-lru==2.0.4
+attrs==23.2.0
+beartype==0.18.5
+beautifulsoup4==4.12.3
+bleach==6.1.0
+braceexpand==0.1.7
+certifi==2024.6.2
+cffi==1.16.0
+charset-normalizer==3.3.2
+click==8.1.7
+clip-anytorch==2.6.0
+clip==1.0
+comm==0.2.2
+contourpy==1.2.1
+cycler==0.12.1
+dalle2-pytorch==1.15.6
+datasets==2.20.0
+debugpy==1.8.1
+decorator==5.1.1
+deepspeed==0.13.1
+defusedxml==0.7.1
+diffusers==0.23.0
+dill==0.3.8
+docker-pycreds==0.4.0
+einops==0.8.0
+einx==0.3.0
+ema-pytorch==0.5.0
+embedding-reader==1.7.0
+evaluate==0.4.1
+executing==2.0.1
+fastjsonschema==2.20.0
+filelock==3.15.1
+fonttools==4.53.0
+fqdn==1.5.1
+frozendict==2.4.4
+frozenlist==1.4.1
+fsspec==2024.5.0
+ftfy==6.2.0
+gitdb==4.0.11
+h11==0.14.0
+h5py==3.10.0
+hjson==3.1.0
+httpcore==1.0.5
+httpx==0.27.0
+huggingface-hub==0.23.4
+idna==3.7
+imageio==2.34.1
+importlib_metadata==7.1.0
+ipykernel==6.29.4
+ipython==8.25.0
+ipywidgets==8.1.3
+isoduration==20.11.0
+jedi==0.19.1
+joblib==1.4.2
+json5==0.9.25
+jsonpointer==3.0.0
+jsonschema-specifications==2023.12.1
+jsonschema==4.22.0
+jupyter-console==6.6.3
+jupyter-events==0.10.0
+jupyter-lsp==2.2.5
+jupyter==1.0.0
+jupyter_client==8.6.2
+jupyter_core==5.7.2
+jupyter_server==2.14.1
+jupyter_server_terminals==0.5.3
+jupyterlab==4.2.2
+jupyterlab_nvdashboard==0.11.0
+jupyterlab_pygments==0.3.0
+jupyterlab_server==2.27.2
+jupyterlab_widgets==3.0.11
+kiwisolver==1.4.5
+kornia==0.7.1
+lazy_loader==0.4
+lightning-utilities==0.11.2
+matplotlib-inline==0.1.7
+matplotlib==3.8.2
+mistune==3.0.2
+mpmath==1.3.0
+multidict==6.0.5
+multiprocess==0.70.16
+nbclient==0.10.0
+nbconvert==7.16.4
+nbformat==5.10.4
+nest-asyncio==1.6.0
+networkx==3.3
+ninja==1.11.1.1
+nltk==3.8.1
+notebook==7.2.1
+notebook_shim==0.2.4
+numpy==1.26.4
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.18.1
+nvidia-nvjitlink-cu12==12.5.40
+nvidia-nvtx-cu12==12.1.105
+omegaconf==2.3.0
+open-clip-torch==2.24.0
+overrides==7.7.0
+packaging==24.1
+pandas==2.2.0
+pandocfilters==1.5.1
+parso==0.8.4
+pexpect==4.9.0
+pillow==10.3.0
+pip==24.0
+platformdirs==4.2.2
+prometheus_client==0.20.0
+prompt_toolkit==3.0.47
+protobuf==5.27.1
+psutil==5.9.8
+ptyprocess==0.7.0
+pure-eval==0.2.2
+py-cpuinfo==9.0.0
+pyarrow-hotfix==0.6
+pyarrow==15.0.2
+pycparser==2.22
+pydantic==2.7.4
+pydantic_core==2.18.4
+pynvml==11.5.0
+pyparsing==3.1.2
+python-dateutil==2.9.0.post0
+python-json-logger==2.0.7
+pytorch-lightning==2.0.1
+pytorch-warmup==0.1.1
+pytz==2024.1
+pyzmq==26.0.3
+qtconsole==5.5.2
+referencing==0.35.1
+regex==2024.5.15
+requests==2.32.3
+resize-right==0.0.2
+responses==0.18.0
+rfc3339-validator==0.1.4
+rfc3986-validator==0.1.1
+rotary-embedding-torch==0.6.2
+rouge_score==0.1.2
+rpds-py==0.18.1
+safetensors==0.4.3
+scikit-image==0.22.0
+scikit-learn==1.5.0
+scipy==1.13.1
+sentence-transformers==2.5.1
+sentencepiece==0.2.0
+sentry-sdk==2.5.1
+setproctitle==1.3.3
+setuptools==65.5.0
+six==1.16.0
+smmap==5.0.1
+sniffio==1.3.1
+soupsieve==2.5
+stack-data==0.6.3
+sympy==1.12.1
+terminado==0.18.1
+threadpoolctl==3.5.0
+tifffile==2024.5.22
+timm==1.0.3
+tinycss2==1.3.0
+tokenizers==0.15.2
+torch-fidelity==0.3.0
+torch==2.1.0
+torchmetrics==1.3.0.post0
+torchvision==0.16.0
+tornado==6.4.1
+tqdm==4.66.4
+traitlets==5.14.3
+transformers==4.37.2
+triton==2.1.0
+types-python-dateutil==2.9.0.20240316
+typing_extensions==4.12.2
+tzdata==2024.1
+umap==0.1.1
+uri-template==1.3.0
+urllib3==2.2.1
+vector-quantize-pytorch==1.14.24
+wandb==0.17.1
+wcwidth==0.2.13
+webcolors==24.6.0
+webdataset==0.2.73
+webencodings==0.5.1
+websocket-client==1.8.0
+widgetsnbextension==4.0.11
+x-clip==0.14.4
+xformers==0.0.22.post7
+xxhash==3.4.1
+yarl==1.9.4
+zipp==3.19.2

MindEyeV2/src/wandb/run-20240618_121613-rr-test/files/wandb-metadata.json

@@ -0,0 +1,1057 @@
+{
+    "os": "Linux-5.15.0-1058-aws-x86_64-with-glibc2.31",
+    "python": "3.11.9",
+    "heartbeatAt": "2024-06-18T12:16:14.300329",
+    "startedAt": "2024-06-18T12:16:13.354272",
+    "docker": null,
+    "cuda": null,
+    "args": [
+        "--data_path=/weka/proj-medarc/shared/mindeyev2_dataset",
+        "--cache_dir=/weka/proj-medarc/shared/cache",
+        "--model_name=rr-test",
+        "--multi_subject",
+        "--subj=1",
+        "--batch_size=21",
+        "--max_lr=3e-4",
+        "--mixup_pct=.33",
+        "--num_epochs=150",
+        "--prior_scale=30",
+        "--clip_scale=1",
+        "--no-blurry_recon",
+        "--blur_scale=.5",
+        "--no-use_image_aug",
+        "--n_blocks=4",
+        "--hidden_dim=1024",
+        "--num_sessions=40",
+        "--ckpt_interval=999",
+        "--ckpt_saving",
+        "--wandb_log"
+    ],
+    "state": "running",
+    "program": "/weka/proj-fmri/ckadirt/MindEyeV2/src/TrainCluster.py",
+    "codePathLocal": "TrainCluster.py",
+    "codePath": "src/TrainCluster.py",
+    "git": {
+        "remote": "https://github.com/MedARC-AI/MindEyeV2.git",
+        "commit": "a3588f0fadc81bdb0cc3e15993c01f8dd64dc7fe"
+    },
+    "email": null,
+    "root": "/weka/proj-fmri/ckadirt/MindEyeV2",
+    "host": "ip-10-0-154-61",
+    "username": "ckadirt",
+    "executable": "/admin/home-ckadirt/fmri/bin/python3.11",
+    "cpu_count": 96,
+    "cpu_count_logical": 192,
+    "cpu_freq": {
+        "current": 2694.2966979166763,
+        "min": 0.0,
+        "max": 0.0
+    },
+    "cpu_freq_per_core": [
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3112.839,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.447,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.107,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3596.856,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.522,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.648,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3598.753,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3598.438,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.905,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3362.07,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3374.828,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        }
+    ],
+    "disk": {
+        "/": {
+            "total": 232.4169158935547,
+            "used": 170.42882919311523
+        }
+    },
+    "gpu": "NVIDIA H100 80GB HBM3",
+    "gpu_count": 8,
+    "gpu_devices": [
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        }
+    ],
+    "memory": {
+        "total": 1999.962516784668
+    }
+}

MindEyeV2/src/wandb/run-20240618_121613-rr-test/logs/debug-internal.log

@@ -0,0 +1,52 @@
+2024-06-18 12:16:13,467 INFO    StreamThr :828956 [internal.py:wandb_internal():85] W&B internal server running at pid: 828956, started at: 2024-06-18 12:16:13.463596
+2024-06-18 12:16:13,468 DEBUG   HandlerThread:828956 [handler.py:handle_request():158] handle_request: status
+2024-06-18 12:16:13,492 INFO    WriterThread:828956 [datastore.py:open_for_write():87] open: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_121613-rr-test/run-rr-test.wandb
+2024-06-18 12:16:13,496 DEBUG   SenderThread:828956 [sender.py:send():379] send: header
+2024-06-18 12:16:13,562 DEBUG   SenderThread:828956 [sender.py:send():379] send: run
+2024-06-18 12:16:13,574 INFO    SenderThread:828956 [sender.py:_setup_resume():749] checking resume status for None/mindeye_semantic_cluster/rr-test
+2024-06-18 12:16:14,023 INFO    SenderThread:828956 [dir_watcher.py:__init__():211] watching files in: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_121613-rr-test/files
+2024-06-18 12:16:14,023 INFO    SenderThread:828956 [sender.py:_start_run_threads():1188] run started: rr-test with start time 1718712973.469764
+2024-06-18 12:16:14,048 DEBUG   HandlerThread:828956 [handler.py:handle_request():158] handle_request: check_version
+2024-06-18 12:16:14,048 DEBUG   SenderThread:828956 [sender.py:send_request():406] send_request: check_version
+2024-06-18 12:16:14,148 DEBUG   HandlerThread:828956 [handler.py:handle_request():158] handle_request: run_start
+2024-06-18 12:16:14,236 DEBUG   HandlerThread:828956 [system_info.py:__init__():26] System info init
+2024-06-18 12:16:14,236 DEBUG   HandlerThread:828956 [system_info.py:__init__():41] System info init done
+2024-06-18 12:16:14,236 INFO    HandlerThread:828956 [system_monitor.py:start():194] Starting system monitor
+2024-06-18 12:16:14,237 INFO    SystemMonitor:828956 [system_monitor.py:_start():158] Starting system asset monitoring threads
+2024-06-18 12:16:14,237 INFO    HandlerThread:828956 [system_monitor.py:probe():214] Collecting system info
+2024-06-18 12:16:14,237 INFO    SystemMonitor:828956 [interfaces.py:start():188] Started cpu monitoring
+2024-06-18 12:16:14,238 INFO    SystemMonitor:828956 [interfaces.py:start():188] Started disk monitoring
+2024-06-18 12:16:14,238 INFO    SystemMonitor:828956 [interfaces.py:start():188] Started gpu monitoring
+2024-06-18 12:16:14,239 INFO    SystemMonitor:828956 [interfaces.py:start():188] Started memory monitoring
+2024-06-18 12:16:14,239 INFO    SystemMonitor:828956 [interfaces.py:start():188] Started network monitoring
+2024-06-18 12:16:14,300 DEBUG   HandlerThread:828956 [system_info.py:probe():152] Probing system
+2024-06-18 12:16:14,308 DEBUG   HandlerThread:828956 [system_info.py:_probe_git():137] Probing git
+2024-06-18 12:16:14,331 DEBUG   HandlerThread:828956 [system_info.py:_probe_git():145] Probing git done
+2024-06-18 12:16:14,331 DEBUG   HandlerThread:828956 [system_info.py:probe():200] Probing system done
+2024-06-18 12:16:14,332 DEBUG   HandlerThread:828956 [system_monitor.py:probe():223] {'os': 'Linux-5.15.0-1058-aws-x86_64-with-glibc2.31', 'python': '3.11.9', 'heartbeatAt': '2024-06-18T12:16:14.300329', 'startedAt': '2024-06-18T12:16:13.354272', 'docker': None, 'cuda': None, 'args': ('--data_path=/weka/proj-medarc/shared/mindeyev2_dataset', '--cache_dir=/weka/proj-medarc/shared/cache', '--model_name=rr-test', '--multi_subject', '--subj=1', '--batch_size=21', '--max_lr=3e-4', '--mixup_pct=.33', '--num_epochs=150', '--prior_scale=30', '--clip_scale=1', '--no-blurry_recon', '--blur_scale=.5', '--no-use_image_aug', '--n_blocks=4', '--hidden_dim=1024', '--num_sessions=40', '--ckpt_interval=999', '--ckpt_saving', '--wandb_log'), 'state': 'running', 'program': '/weka/proj-fmri/ckadirt/MindEyeV2/src/TrainCluster.py', 'codePathLocal': 'TrainCluster.py', 'codePath': 'src/TrainCluster.py', 'git': {'remote': 'https://github.com/MedARC-AI/MindEyeV2.git', 'commit': 'a3588f0fadc81bdb0cc3e15993c01f8dd64dc7fe'}, 'email': None, 'root': '/weka/proj-fmri/ckadirt/MindEyeV2', 'host': 'ip-10-0-154-61', 'username': 'ckadirt', 'executable': '/admin/home-ckadirt/fmri/bin/python3.11', 'cpu_count': 96, 'cpu_count_logical': 192, 'cpu_freq': {'current': 2694.2966979166763, 'min': 0.0, 'max': 0.0}, 'cpu_freq_per_core': [{'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 3112.839, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 3597.447, 'min': 0.0, 'max': 0.0}, {'current': 3597.107, 'min': 0.0, 'max': 0.0}, {'current': 3596.856, 'min': 0.0, 'max': 0.0}, {'current': 3597.522, 'min': 0.0, 'max': 0.0}, {'current': 3597.648, 'min': 0.0, 'max': 0.0}, {'current': 3598.753, 'min': 0.0, 'max': 0.0}, {'current': 3598.438, 'min': 0.0, 'max': 0.0}, {'current': 3597.905, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 3362.07, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 3374.828, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}], 'disk': {'/': {'total': 232.4169158935547, 'used': 170.42882919311523}}, 'gpu': 'NVIDIA H100 80GB HBM3', 'gpu_count': 8, 'gpu_devices': [{'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}, {'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}, {'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}, {'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}, {'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}, {'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}, {'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}, {'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}], 'memory': {'total': 1999.962516784668}}
+2024-06-18 12:16:14,332 INFO    HandlerThread:828956 [system_monitor.py:probe():224] Finished collecting system info
+2024-06-18 12:16:14,332 INFO    HandlerThread:828956 [system_monitor.py:probe():227] Publishing system info
+2024-06-18 12:16:14,332 DEBUG   HandlerThread:828956 [system_info.py:_save_code():44] Saving code
+2024-06-18 12:16:14,365 DEBUG   HandlerThread:828956 [system_info.py:_save_code():67] Saving code done
+2024-06-18 12:16:14,365 DEBUG   HandlerThread:828956 [system_info.py:_save_patches():84] Saving git patches
+2024-06-18 12:16:15,014 DEBUG   HandlerThread:828956 [system_info.py:_save_patches():126] Saving git patches done
+2024-06-18 12:16:15,021 INFO    HandlerThread:828956 [system_monitor.py:probe():229] Finished publishing system info
+2024-06-18 12:16:15,033 INFO    Thread-12 :828956 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_121613-rr-test/files/diff.patch
+2024-06-18 12:16:15,033 INFO    Thread-12 :828956 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_121613-rr-test/files/code/src/TrainCluster.py
+2024-06-18 12:16:15,033 INFO    Thread-12 :828956 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_121613-rr-test/files/wandb-metadata.json
+2024-06-18 12:16:15,033 INFO    Thread-12 :828956 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_121613-rr-test/files/code/src
+2024-06-18 12:16:15,033 INFO    Thread-12 :828956 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_121613-rr-test/files/code
+2024-06-18 12:16:15,039 DEBUG   SenderThread:828956 [sender.py:send():379] send: files
+2024-06-18 12:16:15,040 INFO    SenderThread:828956 [sender.py:_save_file():1454] saving file wandb-metadata.json with policy now
+2024-06-18 12:16:15,040 INFO    SenderThread:828956 [sender.py:_save_file():1454] saving file code/src/TrainCluster.py with policy now
+2024-06-18 12:16:15,040 INFO    SenderThread:828956 [sender.py:_save_file():1454] saving file diff.patch with policy now
+2024-06-18 12:16:15,566 INFO    wandb-upload_0:828956 [upload_job.py:push():130] Uploaded file /tmp/tmp9o0_c5jjwandb/05oxlsji-wandb-metadata.json
+2024-06-18 12:16:15,628 DEBUG   HandlerThread:828956 [handler.py:handle_request():158] handle_request: python_packages
+2024-06-18 12:16:15,629 DEBUG   SenderThread:828956 [sender.py:send_request():406] send_request: python_packages
+2024-06-18 12:16:15,629 DEBUG   HandlerThread:828956 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 12:16:15,630 DEBUG   HandlerThread:828956 [handler.py:handle_request():158] handle_request: internal_messages
+2024-06-18 12:16:15,635 DEBUG   SenderThread:828956 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 12:16:15,811 DEBUG   SenderThread:828956 [sender.py:send():379] send: telemetry
+2024-06-18 12:16:15,866 INFO    wandb-upload_2:828956 [upload_job.py:push():130] Uploaded file /tmp/tmp9o0_c5jjwandb/xypu7qg2-diff.patch
+2024-06-18 12:16:15,886 INFO    wandb-upload_1:828956 [upload_job.py:push():130] Uploaded file /tmp/tmp9o0_c5jjwandb/z7aiv2ba-code/src/TrainCluster.py
+2024-06-18 12:16:16,048 INFO    Thread-12 :828956 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_121613-rr-test/files/requirements.txt

MindEyeV2/src/wandb/run-20240618_121613-rr-test/logs/debug.log

@@ -0,0 +1,28 @@
+2024-06-18 12:16:13,451 INFO    MainThread:828664 [wandb_setup.py:_flush():76] Current SDK version is 0.17.1
+2024-06-18 12:16:13,451 INFO    MainThread:828664 [wandb_setup.py:_flush():76] Configure stats pid to 828664
+2024-06-18 12:16:13,451 INFO    MainThread:828664 [wandb_setup.py:_flush():76] Loading settings from /admin/home-ckadirt/.config/wandb/settings
+2024-06-18 12:16:13,451 INFO    MainThread:828664 [wandb_setup.py:_flush():76] Loading settings from /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/settings
+2024-06-18 12:16:13,451 INFO    MainThread:828664 [wandb_setup.py:_flush():76] Loading settings from environment variables: {}
+2024-06-18 12:16:13,451 INFO    MainThread:828664 [wandb_setup.py:_flush():76] Applying setup settings: {'_disable_service': False}
+2024-06-18 12:16:13,451 INFO    MainThread:828664 [wandb_setup.py:_flush():76] Inferring run settings from compute environment: {'program_relpath': 'src/TrainCluster.py', 'program_abspath': '/weka/proj-fmri/ckadirt/MindEyeV2/src/TrainCluster.py', 'program': '/weka/proj-fmri/ckadirt/MindEyeV2/src/TrainCluster.py'}
+2024-06-18 12:16:13,451 INFO    MainThread:828664 [wandb_setup.py:_flush():76] Applying login settings: {}
+2024-06-18 12:16:13,452 INFO    MainThread:828664 [wandb_init.py:_log_setup():520] Logging user logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_121613-rr-test/logs/debug.log
+2024-06-18 12:16:13,453 INFO    MainThread:828664 [wandb_init.py:_log_setup():521] Logging internal logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_121613-rr-test/logs/debug-internal.log
+2024-06-18 12:16:13,453 INFO    MainThread:828664 [wandb_init.py:init():560] calling init triggers
+2024-06-18 12:16:13,453 INFO    MainThread:828664 [wandb_init.py:init():567] wandb.init called with sweep_config: {}
+config: {'model_name': 'rr-test', 'global_batch_size': '21', 'batch_size': 3, 'num_epochs': 150, 'num_sessions': 40, 'num_params': 573919937, 'clip_scale': 1.0, 'prior_scale': 30.0, 'blur_scale': 0.5, 'use_image_aug': False, 'max_lr': 0.0003, 'mixup_pct': 0.33, 'num_samples_per_epoch': 3750, 'num_test': 3000, 'ckpt_interval': 999, 'ckpt_saving': True, 'seed': 42, 'distributed': True, 'num_devices': 8, 'world_size': 1, 'train_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj08/train/{0..29}.tar', 'test_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj02/new_test/0.tar'}
+2024-06-18 12:16:13,453 INFO    MainThread:828664 [wandb_init.py:init():610] starting backend
+2024-06-18 12:16:13,453 INFO    MainThread:828664 [wandb_init.py:init():614] setting up manager
+2024-06-18 12:16:13,462 INFO    MainThread:828664 [backend.py:_multiprocessing_setup():105] multiprocessing start_methods=fork,spawn,forkserver, using: spawn
+2024-06-18 12:16:13,468 INFO    MainThread:828664 [wandb_init.py:init():622] backend started and connected
+2024-06-18 12:16:13,492 INFO    MainThread:828664 [wandb_init.py:init():711] updated telemetry
+2024-06-18 12:16:13,562 INFO    MainThread:828664 [wandb_init.py:init():744] communicating run to backend with 90.0 second timeout
+2024-06-18 12:16:14,047 INFO    MainThread:828664 [wandb_run.py:_on_init():2402] communicating current version
+2024-06-18 12:16:14,091 INFO    MainThread:828664 [wandb_run.py:_on_init():2411] got version response upgrade_message: "wandb version 0.17.2 is available!  To upgrade, please run:\n $ pip install wandb --upgrade"
+
+2024-06-18 12:16:14,092 INFO    MainThread:828664 [wandb_init.py:init():795] starting run threads in backend
+2024-06-18 12:16:15,629 INFO    MainThread:828664 [wandb_run.py:_console_start():2380] atexit reg
+2024-06-18 12:16:15,629 INFO    MainThread:828664 [wandb_run.py:_redirect():2235] redirect: wrap_raw
+2024-06-18 12:16:15,630 INFO    MainThread:828664 [wandb_run.py:_redirect():2300] Wrapping output streams.
+2024-06-18 12:16:15,630 INFO    MainThread:828664 [wandb_run.py:_redirect():2325] Redirects installed.
+2024-06-18 12:16:15,638 INFO    MainThread:828664 [wandb_init.py:init():838] run started, returning control to user process

MindEyeV2/src/wandb/run-20240618_123614-testing/files/config.yaml

@@ -0,0 +1,115 @@
+wandb_version: 1
+
+model_name:
+  desc: null
+  value: testing
+global_batch_size:
+  desc: null
+  value: 8
+batch_size:
+  desc: null
+  value: 8
+num_epochs:
+  desc: null
+  value: 150
+num_sessions:
+  desc: null
+  value: 40
+num_params:
+  desc: null
+  value: 486928065
+clip_scale:
+  desc: null
+  value: 1.0
+prior_scale:
+  desc: null
+  value: 30.0
+blur_scale:
+  desc: null
+  value: 0.5
+use_image_aug:
+  desc: null
+  value: false
+max_lr:
+  desc: null
+  value: 0.0003
+mixup_pct:
+  desc: null
+  value: 0.33
+num_samples_per_epoch:
+  desc: null
+  value: 30000
+num_test:
+  desc: null
+  value: 3000
+ckpt_interval:
+  desc: null
+  value: 999
+ckpt_saving:
+  desc: null
+  value: false
+seed:
+  desc: null
+  value: 42
+distributed:
+  desc: null
+  value: false
+num_devices:
+  desc: null
+  value: 1
+world_size:
+  desc: null
+  value: 1
+train_url:
+  desc: null
+  value: /weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar
+test_url:
+  desc: null
+  value: /weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar
+_wandb:
+  desc: null
+  value:
+    python_version: 3.11.9
+    cli_version: 0.17.1
+    framework: huggingface
+    huggingface_version: 4.37.2
+    is_jupyter_run: true
+    is_kaggle_kernel: false
+    start_time: 1718714174
+    t:
+      1:
+      - 1
+      - 9
+      - 11
+      - 41
+      - 49
+      - 55
+      - 63
+      - 71
+      - 79
+      - 83
+      - 103
+      2:
+      - 1
+      - 9
+      - 11
+      - 41
+      - 49
+      - 55
+      - 63
+      - 71
+      - 79
+      - 83
+      - 103
+      3:
+      - 13
+      - 14
+      - 16
+      - 23
+      4: 3.11.9
+      5: 0.17.1
+      6: 4.37.2
+      8:
+      - 1
+      - 5
+      13: linux-x86_64

MindEyeV2/src/wandb/run-20240618_123614-testing/files/diff.patch

@@ -0,0 +1,239 @@
+diff --git a/src/Train.ipynb b/src/Train.ipynb
+index 91922c8..396549d 100644
+--- a/src/Train.ipynb
++++ b/src/Train.ipynb
+@@ -99,7 +99,7 @@
+      "name": "stdout",
+      "output_type": "stream",
+      "text": [
+-      "PID of this process = 3001293\n",
++      "PID of this process = 3520819\n",
+       "device: cuda\n",
+       "Distributed environment: DistributedType.NO\n",
+       "Num processes: 1\n",
+@@ -1362,9 +1362,9 @@
+  ],
+  "metadata": {
+   "kernelspec": {
+-   "display_name": "mindeye",
++   "display_name": "Python 3 (ipykernel)",
+    "language": "python",
+-   "name": "mindeye"
++   "name": "python3"
+   },
+   "language_info": {
+    "codemirror_mode": {
+diff --git a/src/accel.slurm b/src/accel.slurm
+index 1c1dfbc..50097be 100644
+--- a/src/accel.slurm
++++ b/src/accel.slurm
+@@ -1,7 +1,7 @@
+ #!/bin/bash
+-#SBATCH --account=topfmri
+-#SBATCH --partition=a40x
+-#SBATCH --job-name=eye
++#SBATCH --account=fmri
++#SBATCH --partition=p5
++#SBATCH --job-name=eyeRR
+ #SBATCH --nodes=1           
+ #SBATCH --gres=gpu:1
+ #SBATCH --time=48:00:00          # total run time limit (HH:MM:SS)
+@@ -12,11 +12,11 @@
+ #SBATCH --exclusive
+ 
+ # Make sure you activate your fmri environment created from src/setup.sh
+-cd /weka/proj-fmri/paulscotti/MindEyeV2/src
+-source fmri/bin/activate
++cd /weka/proj-fmri/ckadirt/MindEyeV2/src
++source /admin/home-ckadirt/fmri/bin/activate
+ 
+ # The following line converts your jupyter notebook into a python script runnable with Slurm
+-jupyter nbconvert Train.ipynb --to python
++jupyter nbconvert TrainCluster.ipynb --to python
+ 
+ export NUM_GPUS=1  # Set to equal gres=gpu:#!
+ export BATCH_SIZE=21 # 21 for multisubject / 24 for singlesubject (orig. paper used 42 for multisubject / 24 for singlesubject)
+@@ -26,16 +26,16 @@ export GLOBAL_BATCH_SIZE=$((BATCH_SIZE * NUM_GPUS))
+ 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 COUNT_NODE=1 #$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)
+ echo MASTER_ADDR=${MASTER_ADDR}
+ echo MASTER_PORT=${MASTER_PORT}
+ echo WORLD_SIZE=${COUNT_NODE}
+ 
+ # multisubject pretraining
+-model_name="multisubject_excludingsubj01_40sess"
++model_name="rr-test"
+ echo model_name=${model_name}
+-accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --mixed_precision=fp16 Train.py --data_path=/weka/proj-fmri/shared/mindeyev2_dataset --cache_dir=/weka/proj-fmri/shared/cache --model_name=${model_name} --multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --use_prior --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log
+-
++# accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --mixed_precision=fp16 TrainCluster.py --data_path=/weka/proj-medarc/shared/mindeyev2_dataset --cache_dir=/weka/proj-medarc/shared/cache --model_name=${model_name} --multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log
++python TrainCluster.py --data_path=/weka/proj-medarc/shared/mindeyev2_dataset --cache_dir=/weka/proj-medarc/shared/cache --model_name=${model_name} --multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log
+ # singlesubject finetuning
+ #model_name="finetuned_subj01_40sess"
+ #echo model_name=${model_name}
+diff --git a/src/recon_inference.ipynb b/src/recon_inference.ipynb
+index 10d7e27..0593e7e 100644
+--- a/src/recon_inference.ipynb
++++ b/src/recon_inference.ipynb
+@@ -12,7 +12,6 @@
+      "name": "stdout",
+      "output_type": "stream",
+      "text": [
+-      "LOCAL RANK  0\n",
+       "device: cuda\n"
+      ]
+     }
+@@ -60,7 +59,7 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 33,
++   "execution_count": 2,
+    "id": "e52985b1-95ff-487b-8b2d-cc1ad1c190b8",
+    "metadata": {
+     "tags": []
+@@ -71,9 +70,7 @@
+      "output_type": "stream",
+      "text": [
+       "model_name: final_subj01_pretrained_40sess_24bs\n",
+-      "--data_path=/weka/proj-medarc/shared/mindeyev2_dataset                     --cache_dir=/weka/proj-medarc/shared/mindeyev2_dataset                     --model_name=final_subj01_pretrained_40sess_24bs --subj=1                     --hidden_dim=4096 --n_blocks=4 --new_test\n",
+-      "The autoreload extension is already loaded. To reload it, use:\n",
+-      "  %reload_ext autoreload\n"
++      "--data_path=/weka/proj-medarc/shared/mindeyev2_dataset                     --cache_dir=/weka/proj-medarc/shared/mindeyev2_dataset                     --model_name=final_subj01_pretrained_40sess_24bs --subj=1                     --hidden_dim=4096 --n_blocks=4 --new_test\n"
+      ]
+     }
+    ],
+@@ -99,7 +96,7 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 34,
++   "execution_count": 3,
+    "id": "49e5dae4-606d-4dc6-b420-df9e4c14737e",
+    "metadata": {
+     "tags": []
+@@ -157,7 +154,7 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 35,
++   "execution_count": 4,
+    "id": "64672583-9f00-46f5-8d4e-00e4c7068a1d",
+    "metadata": {
+     "tags": []
+@@ -221,7 +218,7 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 36,
++   "execution_count": 5,
+    "id": "a3cbeea8-e95b-48d9-9bc2-91af260c93d1",
+    "metadata": {},
+    "outputs": [
+@@ -254,39 +251,10 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 38,
++   "execution_count": null,
+    "id": "3afc4858-b6a6-4a52-9303-b4a50ea5cc0b",
+    "metadata": {},
+-   "outputs": [
+-    {
+-     "name": "stdout",
+-     "output_type": "stream",
+-     "text": [
+-      "param counts:\n",
+-      "83,653,863 total\n",
+-      "0 trainable\n",
+-      "param counts:\n",
+-      "64,409,600 total\n",
+-      "64,409,600 trainable\n",
+-      "param counts:\n",
+-      "1,903,020,028 total\n",
+-      "1,903,020,028 trainable\n",
+-      "param counts:\n",
+-      "1,967,429,628 total\n",
+-      "1,967,429,628 trainable\n",
+-      "param counts:\n",
+-      "259,865,216 total\n",
+-      "259,865,200 trainable\n",
+-      "param counts:\n",
+-      "2,227,294,844 total\n",
+-      "2,227,294,828 trainable\n",
+-      "\n",
+-      "---loading /weka/proj-fmri/paulscotti/MindEye2_git/train_logs/final_subj01_pretrained_40sess_24bs/last.pth ckpt---\n",
+-      "\n",
+-      "ckpt loaded!\n"
+-     ]
+-    }
+-   ],
++   "outputs": [],
+    "source": [
+     "clip_img_embedder = FrozenOpenCLIPImageEmbedder(\n",
+     "    arch=\"ViT-bigG-14\",\n",
+@@ -393,7 +361,7 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 30,
++   "execution_count": null,
+    "id": "295824db-ab3d-450c-90fb-f656e48994ba",
+    "metadata": {
+     "tags": []
+@@ -430,38 +398,10 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 31,
++   "execution_count": null,
+    "id": "f726f617-39f5-49e2-8d0c-d11d27d01c30",
+    "metadata": {},
+-   "outputs": [
+-    {
+-     "name": "stderr",
+-     "output_type": "stream",
+-     "text": [
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now.\n"
+-     ]
+-    },
+-    {
+-     "name": "stdout",
+-     "output_type": "stream",
+-     "text": [
+-      "Initialized embedder #0: FrozenOpenCLIPImageEmbedder with 1909889025 params. Trainable: False\n",
+-      "Initialized embedder #1: ConcatTimestepEmbedderND with 0 params. Trainable: False\n",
+-      "Initialized embedder #2: ConcatTimestepEmbedderND with 0 params. Trainable: False\n",
+-      "vector_suffix torch.Size([1, 1024])\n"
+-     ]
+-    }
+-   ],
++   "outputs": [],
+    "source": [
+     "# prep unCLIP\n",
+     "config = OmegaConf.load(\"generative_models/configs/unclip6.yaml\")\n",
+@@ -715,9 +655,9 @@
+  ],
+  "metadata": {
+   "kernelspec": {
+-   "display_name": "mindeye",
++   "display_name": "Python 3 (ipykernel)",
+    "language": "python",
+-   "name": "mindeye"
++   "name": "python3"
+   },
+   "language_info": {
+    "codemirror_mode": {

MindEyeV2/src/wandb/run-20240618_123614-testing/files/requirements.txt

@@ -0,0 +1,214 @@
+Babel==2.15.0
+CoCa-pytorch==0.1.0
+GitPython==3.1.43
+Jinja2==3.1.4
+MarkupSafe==2.1.5
+PyYAML==6.0.1
+Pygments==2.18.0
+QtPy==2.4.1
+Send2Trash==1.8.3
+absl-py==2.1.0
+accelerate==0.24.1
+aiohttp==3.9.5
+aiosignal==1.3.1
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+anyio==4.4.0
+argon2-cffi-bindings==21.2.0
+argon2-cffi==23.1.0
+arrow==1.3.0
+asttokens==2.4.1
+async-lru==2.0.4
+attrs==23.2.0
+beartype==0.18.5
+beautifulsoup4==4.12.3
+bleach==6.1.0
+braceexpand==0.1.7
+certifi==2024.6.2
+cffi==1.16.0
+charset-normalizer==3.3.2
+click==8.1.7
+clip-anytorch==2.6.0
+clip==1.0
+comm==0.2.2
+contourpy==1.2.1
+cycler==0.12.1
+dalle2-pytorch==1.15.6
+datasets==2.20.0
+debugpy==1.8.1
+decorator==5.1.1
+deepspeed==0.13.1
+defusedxml==0.7.1
+diffusers==0.23.0
+dill==0.3.8
+docker-pycreds==0.4.0
+einops==0.8.0
+einx==0.3.0
+ema-pytorch==0.5.0
+embedding-reader==1.7.0
+evaluate==0.4.1
+executing==2.0.1
+fastjsonschema==2.20.0
+filelock==3.15.1
+fonttools==4.53.0
+fqdn==1.5.1
+frozendict==2.4.4
+frozenlist==1.4.1
+fsspec==2024.5.0
+ftfy==6.2.0
+gitdb==4.0.11
+h11==0.14.0
+h5py==3.10.0
+hjson==3.1.0
+httpcore==1.0.5
+httpx==0.27.0
+huggingface-hub==0.23.4
+idna==3.7
+imageio==2.34.1
+importlib_metadata==7.1.0
+ipykernel==6.29.4
+ipython==8.25.0
+ipywidgets==8.1.3
+isoduration==20.11.0
+jedi==0.19.1
+joblib==1.4.2
+json5==0.9.25
+jsonpointer==3.0.0
+jsonschema-specifications==2023.12.1
+jsonschema==4.22.0
+jupyter-console==6.6.3
+jupyter-events==0.10.0
+jupyter-lsp==2.2.5
+jupyter==1.0.0
+jupyter_client==8.6.2
+jupyter_core==5.7.2
+jupyter_server==2.14.1
+jupyter_server_terminals==0.5.3
+jupyterlab==4.2.2
+jupyterlab_nvdashboard==0.11.0
+jupyterlab_pygments==0.3.0
+jupyterlab_server==2.27.2
+jupyterlab_widgets==3.0.11
+kiwisolver==1.4.5
+kornia==0.7.1
+lazy_loader==0.4
+lightning-utilities==0.11.2
+matplotlib-inline==0.1.7
+matplotlib==3.8.2
+mistune==3.0.2
+mpmath==1.3.0
+multidict==6.0.5
+multiprocess==0.70.16
+nbclient==0.10.0
+nbconvert==7.16.4
+nbformat==5.10.4
+nest-asyncio==1.6.0
+networkx==3.3
+ninja==1.11.1.1
+nltk==3.8.1
+notebook==7.2.1
+notebook_shim==0.2.4
+numpy==1.26.4
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.18.1
+nvidia-nvjitlink-cu12==12.5.40
+nvidia-nvtx-cu12==12.1.105
+omegaconf==2.3.0
+open-clip-torch==2.24.0
+overrides==7.7.0
+packaging==24.1
+pandas==2.2.0
+pandocfilters==1.5.1
+parso==0.8.4
+pexpect==4.9.0
+pillow==10.3.0
+pip==24.0
+platformdirs==4.2.2
+prometheus_client==0.20.0
+prompt_toolkit==3.0.47
+protobuf==5.27.1
+psutil==5.9.8
+ptyprocess==0.7.0
+pure-eval==0.2.2
+py-cpuinfo==9.0.0
+pyarrow-hotfix==0.6
+pyarrow==15.0.2
+pycparser==2.22
+pydantic==2.7.4
+pydantic_core==2.18.4
+pynvml==11.5.0
+pyparsing==3.1.2
+python-dateutil==2.9.0.post0
+python-json-logger==2.0.7
+pytorch-lightning==2.0.1
+pytorch-warmup==0.1.1
+pytz==2024.1
+pyzmq==26.0.3
+qtconsole==5.5.2
+referencing==0.35.1
+regex==2024.5.15
+requests==2.32.3
+resize-right==0.0.2
+responses==0.18.0
+rfc3339-validator==0.1.4
+rfc3986-validator==0.1.1
+rotary-embedding-torch==0.6.2
+rouge_score==0.1.2
+rpds-py==0.18.1
+safetensors==0.4.3
+scikit-image==0.22.0
+scikit-learn==1.5.0
+scipy==1.13.1
+sentence-transformers==2.5.1
+sentencepiece==0.2.0
+sentry-sdk==2.5.1
+setproctitle==1.3.3
+setuptools==65.5.0
+six==1.16.0
+smmap==5.0.1
+sniffio==1.3.1
+soupsieve==2.5
+stack-data==0.6.3
+sympy==1.12.1
+terminado==0.18.1
+threadpoolctl==3.5.0
+tifffile==2024.5.22
+timm==1.0.3
+tinycss2==1.3.0
+tokenizers==0.15.2
+torch-fidelity==0.3.0
+torch==2.1.0
+torchmetrics==1.3.0.post0
+torchvision==0.16.0
+tornado==6.4.1
+tqdm==4.66.4
+traitlets==5.14.3
+transformers==4.37.2
+triton==2.1.0
+types-python-dateutil==2.9.0.20240316
+typing_extensions==4.12.2
+tzdata==2024.1
+umap==0.1.1
+uri-template==1.3.0
+urllib3==2.2.1
+vector-quantize-pytorch==1.14.24
+wandb==0.17.1
+wcwidth==0.2.13
+webcolors==24.6.0
+webdataset==0.2.73
+webencodings==0.5.1
+websocket-client==1.8.0
+widgetsnbextension==4.0.11
+x-clip==0.14.4
+xformers==0.0.22.post7
+xxhash==3.4.1
+yarl==1.9.4
+zipp==3.19.2

MindEyeV2/src/wandb/run-20240618_123614-testing/files/wandb-metadata.json

@@ -0,0 +1,1035 @@
+{
+    "os": "Linux-5.15.0-1058-aws-x86_64-with-glibc2.31",
+    "python": "3.11.9",
+    "heartbeatAt": "2024-06-18T12:36:15.404626",
+    "startedAt": "2024-06-18T12:36:14.302258",
+    "docker": null,
+    "cuda": null,
+    "args": [],
+    "state": "running",
+    "program": "ckadirt/MindEyeV2/src/TrainCluster.ipynb",
+    "codePathLocal": null,
+    "root": "/weka/proj-fmri/ckadirt/MindEyeV2",
+    "git": {
+        "remote": "https://github.com/MedARC-AI/MindEyeV2.git",
+        "commit": "a3588f0fadc81bdb0cc3e15993c01f8dd64dc7fe"
+    },
+    "email": null,
+    "host": "ip-10-0-149-152",
+    "username": "ckadirt",
+    "executable": "/admin/home-ckadirt/fmri/bin/python3.11",
+    "cpu_count": 96,
+    "cpu_count_logical": 192,
+    "cpu_freq": {
+        "current": 2710.1276875,
+        "min": 0.0,
+        "max": 0.0
+    },
+    "cpu_freq_per_core": [
+        {
+            "current": 2925.786,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3596.327,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.513,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3599.095,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.548,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3598.825,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.706,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3599.237,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.748,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3587.835,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.601,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3595.424,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.911,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2650.0,
+            "min": 0.0,
+            "max": 0.0
+        }
+    ],
+    "disk": {
+        "/": {
+            "total": 232.4169158935547,
+            "used": 181.67956161499023
+        }
+    },
+    "gpu": "NVIDIA H100 80GB HBM3",
+    "gpu_count": 8,
+    "gpu_devices": [
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        }
+    ],
+    "memory": {
+        "total": 1999.9625205993652
+    }
+}

MindEyeV2/src/wandb/run-20240618_123614-testing/files/wandb-summary.json

@@ -0,0 +1 @@
+{"train/loss": 149.35307167663575, "test/loss": 1460.2626953125, "train/lr": 0.0002978457718582204, "train/num_steps": 37500, "test/num_steps": 10, "train/fwd_pct_correct": 0.6664333333333333, "train/bwd_pct_correct": 0.6865, "test/test_fwd_pct_correct": 0.5600000023841858, "test/test_bwd_pct_correct": 0.5333333611488342, "train/loss_clip_total": 0.5423962022828559, "train/loss_blurry_total": 0.0, "train/loss_blurry_cont_total": 0.0, "test/loss_clip_total": 2.158545970916748, "train/blurry_pixcorr": 0.0, "test/blurry_pixcorr": 0.0, "train/recon_cossim": 0.0, "test/recon_cossim": 0.0, "train/recon_mse": 0.0, "test/recon_mse": 0.0, "train/loss_prior": 0.0, "test/loss_prior": 0.0, "train/loss_RR": 148.81067535298666, "test/loss_RR": 1458.1041259765625, "_timestamp": 1718717027.775746, "_runtime": 2853.353439092636, "_step": 9, "_wandb": {"runtime": 3000}}

MindEyeV2/src/wandb/run-20240618_123614-testing/logs/debug.log

@@ -0,0 +1,69 @@
+2024-06-18 12:36:14,405 INFO    MainThread:448378 [wandb_setup.py:_flush():76] Current SDK version is 0.17.1
+2024-06-18 12:36:14,406 INFO    MainThread:448378 [wandb_setup.py:_flush():76] Configure stats pid to 448378
+2024-06-18 12:36:14,406 INFO    MainThread:448378 [wandb_setup.py:_flush():76] Loading settings from /admin/home-ckadirt/.config/wandb/settings
+2024-06-18 12:36:14,406 INFO    MainThread:448378 [wandb_setup.py:_flush():76] Loading settings from /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/settings
+2024-06-18 12:36:14,406 INFO    MainThread:448378 [wandb_setup.py:_flush():76] Loading settings from environment variables: {}
+2024-06-18 12:36:14,406 INFO    MainThread:448378 [wandb_setup.py:_flush():76] Inferring run settings from compute environment: {'program': '<python with no main file>'}
+2024-06-18 12:36:14,406 INFO    MainThread:448378 [wandb_setup.py:_flush():76] Applying login settings: {'base_url': 'https://stability.wandb.io'}
+2024-06-18 12:36:14,406 INFO    MainThread:448378 [wandb_setup.py:_flush():76] Applying login settings: {}
+2024-06-18 12:36:14,407 INFO    MainThread:448378 [wandb_init.py:_log_setup():520] Logging user logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_123614-testing/logs/debug.log
+2024-06-18 12:36:14,407 INFO    MainThread:448378 [wandb_init.py:_log_setup():521] Logging internal logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_123614-testing/logs/debug-internal.log
+2024-06-18 12:36:14,407 INFO    MainThread:448378 [wandb_init.py:_jupyter_setup():466] configuring jupyter hooks <wandb.sdk.wandb_init._WandbInit object at 0x7f13c50a0390>
+2024-06-18 12:36:14,407 INFO    MainThread:448378 [wandb_init.py:init():560] calling init triggers
+2024-06-18 12:36:14,407 INFO    MainThread:448378 [wandb_init.py:init():567] wandb.init called with sweep_config: {}
+config: {'model_name': 'testing', 'global_batch_size': 8, 'batch_size': 8, 'num_epochs': 150, 'num_sessions': 40, 'num_params': 486928065, 'clip_scale': 1.0, 'prior_scale': 30.0, 'blur_scale': 0.5, 'use_image_aug': False, 'max_lr': 0.0003, 'mixup_pct': 0.33, 'num_samples_per_epoch': 30000, 'num_test': 3000, 'ckpt_interval': 999, 'ckpt_saving': False, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1, 'train_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar', 'test_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar'}
+2024-06-18 12:36:14,407 INFO    MainThread:448378 [wandb_init.py:init():610] starting backend
+2024-06-18 12:36:14,408 INFO    MainThread:448378 [wandb_init.py:init():614] setting up manager
+2024-06-18 12:36:14,417 INFO    MainThread:448378 [backend.py:_multiprocessing_setup():105] multiprocessing start_methods=fork,spawn,forkserver, using: spawn
+2024-06-18 12:36:14,421 INFO    MainThread:448378 [wandb_init.py:init():622] backend started and connected
+2024-06-18 12:36:14,444 INFO    MainThread:448378 [wandb_run.py:_label_probe_notebook():1334] probe notebook
+2024-06-18 12:36:14,445 INFO    MainThread:448378 [wandb_run.py:_label_probe_notebook():1344] Unable to probe notebook: 'NoneType' object has no attribute 'get'
+2024-06-18 12:36:14,445 INFO    MainThread:448378 [wandb_init.py:init():711] updated telemetry
+2024-06-18 12:36:14,557 INFO    MainThread:448378 [wandb_init.py:init():744] communicating run to backend with 90.0 second timeout
+2024-06-18 12:36:15,149 INFO    MainThread:448378 [wandb_run.py:_on_init():2402] communicating current version
+2024-06-18 12:36:15,199 INFO    MainThread:448378 [wandb_run.py:_on_init():2411] got version response upgrade_message: "wandb version 0.17.2 is available!  To upgrade, please run:\n $ pip install wandb --upgrade"
+
+2024-06-18 12:36:15,200 INFO    MainThread:448378 [wandb_init.py:init():795] starting run threads in backend
+2024-06-18 12:36:16,480 INFO    MainThread:448378 [wandb_run.py:_console_start():2380] atexit reg
+2024-06-18 12:36:16,480 INFO    MainThread:448378 [wandb_run.py:_redirect():2235] redirect: wrap_raw
+2024-06-18 12:36:16,480 INFO    MainThread:448378 [wandb_run.py:_redirect():2300] Wrapping output streams.
+2024-06-18 12:36:16,480 INFO    MainThread:448378 [wandb_run.py:_redirect():2325] Redirects installed.
+2024-06-18 12:36:16,484 INFO    MainThread:448378 [wandb_init.py:init():838] run started, returning control to user process
+2024-06-18 12:36:16,489 INFO    MainThread:448378 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster.ipynb
+2024-06-18 12:36:16,490 INFO    MainThread:448378 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 12:36:16,562 INFO    MainThread:448378 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 12:36:16,563 INFO    MainThread:448378 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster.ipynb
+2024-06-18 12:36:16,563 INFO    MainThread:448378 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 12:36:16,646 INFO    MainThread:448378 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 12:36:16,653 INFO    MainThread:448378 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster.ipynb
+2024-06-18 12:36:16,654 INFO    MainThread:448378 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 12:36:16,746 INFO    MainThread:448378 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 12:36:16,917 INFO    MainThread:448378 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster.ipynb
+2024-06-18 12:36:16,917 INFO    MainThread:448378 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 12:36:17,011 INFO    MainThread:448378 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 12:36:17,012 INFO    MainThread:448378 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster.ipynb
+2024-06-18 12:36:17,013 INFO    MainThread:448378 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 12:36:17,101 INFO    MainThread:448378 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 13:24:59,287 INFO    MainThread:448378 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster.ipynb
+2024-06-18 13:24:59,288 INFO    MainThread:448378 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 13:25:23,444 INFO    MainThread:448378 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 13:25:57,202 INFO    MainThread:448378 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster.ipynb
+2024-06-18 13:25:57,203 INFO    MainThread:448378 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 13:26:25,695 INFO    MainThread:448378 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 13:26:55,886 INFO    MainThread:448378 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster.ipynb
+2024-06-18 13:26:55,887 INFO    MainThread:448378 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 13:28:40,735 INFO    MainThread:448378 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 13:28:40,757 INFO    MainThread:448378 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster.ipynb
+2024-06-18 13:28:40,758 INFO    MainThread:448378 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 13:28:49,551 INFO    MainThread:448378 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 13:28:49,553 INFO    MainThread:448378 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster.ipynb
+2024-06-18 13:28:49,553 INFO    MainThread:448378 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 13:28:58,554 INFO    MainThread:448378 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 13:28:58,558 INFO    MainThread:448378 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster.ipynb
+2024-06-18 13:28:58,559 INFO    MainThread:448378 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 13:29:01,816 INFO    MainThread:448378 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 13:29:01,823 INFO    MainThread:448378 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster.ipynb
+2024-06-18 13:29:01,824 INFO    MainThread:448378 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 13:29:05,134 INFO    MainThread:448378 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 13:29:05,136 INFO    MainThread:448378 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster.ipynb
+2024-06-18 13:29:05,136 INFO    MainThread:448378 [wandb_init.py:_pause_backend():431] pausing backend

MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/config.yaml

@@ -0,0 +1,115 @@
+wandb_version: 1
+
+model_name:
+  desc: null
+  value: testing-tt2
+global_batch_size:
+  desc: null
+  value: 16
+batch_size:
+  desc: null
+  value: 16
+num_epochs:
+  desc: null
+  value: 150
+num_sessions:
+  desc: null
+  value: 40
+num_params:
+  desc: null
+  value: 486928065
+clip_scale:
+  desc: null
+  value: 1.0
+prior_scale:
+  desc: null
+  value: 30.0
+blur_scale:
+  desc: null
+  value: 0.5
+use_image_aug:
+  desc: null
+  value: false
+max_lr:
+  desc: null
+  value: 3.0e-05
+mixup_pct:
+  desc: null
+  value: 0.33
+num_samples_per_epoch:
+  desc: null
+  value: 30000
+num_test:
+  desc: null
+  value: 3000
+ckpt_interval:
+  desc: null
+  value: 999
+ckpt_saving:
+  desc: null
+  value: false
+seed:
+  desc: null
+  value: 42
+distributed:
+  desc: null
+  value: false
+num_devices:
+  desc: null
+  value: 1
+world_size:
+  desc: null
+  value: 1
+train_url:
+  desc: null
+  value: /weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar
+test_url:
+  desc: null
+  value: /weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar
+_wandb:
+  desc: null
+  value:
+    python_version: 3.11.9
+    cli_version: 0.17.1
+    framework: huggingface
+    huggingface_version: 4.37.2
+    is_jupyter_run: true
+    is_kaggle_kernel: false
+    start_time: 1718729436
+    t:
+      1:
+      - 1
+      - 9
+      - 11
+      - 41
+      - 49
+      - 55
+      - 63
+      - 71
+      - 79
+      - 83
+      - 103
+      2:
+      - 1
+      - 9
+      - 11
+      - 41
+      - 49
+      - 55
+      - 63
+      - 71
+      - 79
+      - 83
+      - 103
+      3:
+      - 13
+      - 14
+      - 16
+      - 23
+      4: 3.11.9
+      5: 0.17.1
+      6: 4.37.2
+      8:
+      - 1
+      - 5
+      13: linux-x86_64

MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/diff.patch

@@ -0,0 +1,240 @@
+diff --git a/src/Train.ipynb b/src/Train.ipynb
+index 91922c8..396549d 100644
+--- a/src/Train.ipynb
++++ b/src/Train.ipynb
+@@ -99,7 +99,7 @@
+      "name": "stdout",
+      "output_type": "stream",
+      "text": [
+-      "PID of this process = 3001293\n",
++      "PID of this process = 3520819\n",
+       "device: cuda\n",
+       "Distributed environment: DistributedType.NO\n",
+       "Num processes: 1\n",
+@@ -1362,9 +1362,9 @@
+  ],
+  "metadata": {
+   "kernelspec": {
+-   "display_name": "mindeye",
++   "display_name": "Python 3 (ipykernel)",
+    "language": "python",
+-   "name": "mindeye"
++   "name": "python3"
+   },
+   "language_info": {
+    "codemirror_mode": {
+diff --git a/src/accel.slurm b/src/accel.slurm
+index 1c1dfbc..f0b347f 100644
+--- a/src/accel.slurm
++++ b/src/accel.slurm
+@@ -1,7 +1,8 @@
+ #!/bin/bash
+-#SBATCH --account=topfmri
+-#SBATCH --partition=a40x
+-#SBATCH --job-name=eye
++#SBATCH --account=fmri
++#SBATCH --qos=normal
++#SBATCH --partition=p5
++#SBATCH --job-name=eyeRR
+ #SBATCH --nodes=1           
+ #SBATCH --gres=gpu:1
+ #SBATCH --time=48:00:00          # total run time limit (HH:MM:SS)
+@@ -12,11 +13,11 @@
+ #SBATCH --exclusive
+ 
+ # Make sure you activate your fmri environment created from src/setup.sh
+-cd /weka/proj-fmri/paulscotti/MindEyeV2/src
+-source fmri/bin/activate
++cd /weka/proj-fmri/ckadirt/MindEyeV2/src
++source /admin/home-ckadirt/fmri/bin/activate
+ 
+ # The following line converts your jupyter notebook into a python script runnable with Slurm
+-jupyter nbconvert Train.ipynb --to python
++jupyter nbconvert TrainCluster-Copy2.ipynb --to python
+ 
+ export NUM_GPUS=1  # Set to equal gres=gpu:#!
+ export BATCH_SIZE=21 # 21 for multisubject / 24 for singlesubject (orig. paper used 42 for multisubject / 24 for singlesubject)
+@@ -26,16 +27,16 @@ export GLOBAL_BATCH_SIZE=$((BATCH_SIZE * NUM_GPUS))
+ 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 COUNT_NODE=1 #$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)
+ echo MASTER_ADDR=${MASTER_ADDR}
+ echo MASTER_PORT=${MASTER_PORT}
+ echo WORLD_SIZE=${COUNT_NODE}
+ 
+ # multisubject pretraining
+-model_name="multisubject_excludingsubj01_40sess"
++model_name="rr-test-top"
+ echo model_name=${model_name}
+-accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --mixed_precision=fp16 Train.py --data_path=/weka/proj-fmri/shared/mindeyev2_dataset --cache_dir=/weka/proj-fmri/shared/cache --model_name=${model_name} --multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --use_prior --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log
+-
++# accelerate launch --num_processes=$(($NUM_GPUS * $COUNT_NODE)) --num_machines=$COUNT_NODE --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --mixed_precision=fp16 TrainCluster.py --data_path=/weka/proj-medarc/shared/mindeyev2_dataset --cache_dir=/weka/proj-medarc/shared/cache --model_name=${model_name} --multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log
++python TrainCluster-Copy2.py --data_path=/weka/proj-medarc/shared/mindeyev2_dataset --cache_dir=/weka/proj-medarc/shared/cache --model_name=${model_name} --no-multi_subject --subj=1 --batch_size=${BATCH_SIZE} --max_lr=3e-5 --mixup_pct=.33 --num_epochs=150 --prior_scale=30 --clip_scale=1 --no-blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=1024 --num_sessions=40 --ckpt_interval=999 --ckpt_saving --wandb_log
+ # singlesubject finetuning
+ #model_name="finetuned_subj01_40sess"
+ #echo model_name=${model_name}
+diff --git a/src/recon_inference.ipynb b/src/recon_inference.ipynb
+index 10d7e27..0593e7e 100644
+--- a/src/recon_inference.ipynb
++++ b/src/recon_inference.ipynb
+@@ -12,7 +12,6 @@
+      "name": "stdout",
+      "output_type": "stream",
+      "text": [
+-      "LOCAL RANK  0\n",
+       "device: cuda\n"
+      ]
+     }
+@@ -60,7 +59,7 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 33,
++   "execution_count": 2,
+    "id": "e52985b1-95ff-487b-8b2d-cc1ad1c190b8",
+    "metadata": {
+     "tags": []
+@@ -71,9 +70,7 @@
+      "output_type": "stream",
+      "text": [
+       "model_name: final_subj01_pretrained_40sess_24bs\n",
+-      "--data_path=/weka/proj-medarc/shared/mindeyev2_dataset                     --cache_dir=/weka/proj-medarc/shared/mindeyev2_dataset                     --model_name=final_subj01_pretrained_40sess_24bs --subj=1                     --hidden_dim=4096 --n_blocks=4 --new_test\n",
+-      "The autoreload extension is already loaded. To reload it, use:\n",
+-      "  %reload_ext autoreload\n"
++      "--data_path=/weka/proj-medarc/shared/mindeyev2_dataset                     --cache_dir=/weka/proj-medarc/shared/mindeyev2_dataset                     --model_name=final_subj01_pretrained_40sess_24bs --subj=1                     --hidden_dim=4096 --n_blocks=4 --new_test\n"
+      ]
+     }
+    ],
+@@ -99,7 +96,7 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 34,
++   "execution_count": 3,
+    "id": "49e5dae4-606d-4dc6-b420-df9e4c14737e",
+    "metadata": {
+     "tags": []
+@@ -157,7 +154,7 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 35,
++   "execution_count": 4,
+    "id": "64672583-9f00-46f5-8d4e-00e4c7068a1d",
+    "metadata": {
+     "tags": []
+@@ -221,7 +218,7 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 36,
++   "execution_count": 5,
+    "id": "a3cbeea8-e95b-48d9-9bc2-91af260c93d1",
+    "metadata": {},
+    "outputs": [
+@@ -254,39 +251,10 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 38,
++   "execution_count": null,
+    "id": "3afc4858-b6a6-4a52-9303-b4a50ea5cc0b",
+    "metadata": {},
+-   "outputs": [
+-    {
+-     "name": "stdout",
+-     "output_type": "stream",
+-     "text": [
+-      "param counts:\n",
+-      "83,653,863 total\n",
+-      "0 trainable\n",
+-      "param counts:\n",
+-      "64,409,600 total\n",
+-      "64,409,600 trainable\n",
+-      "param counts:\n",
+-      "1,903,020,028 total\n",
+-      "1,903,020,028 trainable\n",
+-      "param counts:\n",
+-      "1,967,429,628 total\n",
+-      "1,967,429,628 trainable\n",
+-      "param counts:\n",
+-      "259,865,216 total\n",
+-      "259,865,200 trainable\n",
+-      "param counts:\n",
+-      "2,227,294,844 total\n",
+-      "2,227,294,828 trainable\n",
+-      "\n",
+-      "---loading /weka/proj-fmri/paulscotti/MindEye2_git/train_logs/final_subj01_pretrained_40sess_24bs/last.pth ckpt---\n",
+-      "\n",
+-      "ckpt loaded!\n"
+-     ]
+-    }
+-   ],
++   "outputs": [],
+    "source": [
+     "clip_img_embedder = FrozenOpenCLIPImageEmbedder(\n",
+     "    arch=\"ViT-bigG-14\",\n",
+@@ -393,7 +361,7 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 30,
++   "execution_count": null,
+    "id": "295824db-ab3d-450c-90fb-f656e48994ba",
+    "metadata": {
+     "tags": []
+@@ -430,38 +398,10 @@
+   },
+   {
+    "cell_type": "code",
+-   "execution_count": 31,
++   "execution_count": null,
+    "id": "f726f617-39f5-49e2-8d0c-d11d27d01c30",
+    "metadata": {},
+-   "outputs": [
+-    {
+-     "name": "stderr",
+-     "output_type": "stream",
+-     "text": [
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 10. Setting context_dim to [1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now.\n",
+-      "WARNING:sgm.modules.attention:SpatialTransformer: Found context dims [1664] of depth 1, which does not match the specified 'depth' of 2. Setting context_dim to [1664, 1664] now.\n"
+-     ]
+-    },
+-    {
+-     "name": "stdout",
+-     "output_type": "stream",
+-     "text": [
+-      "Initialized embedder #0: FrozenOpenCLIPImageEmbedder with 1909889025 params. Trainable: False\n",
+-      "Initialized embedder #1: ConcatTimestepEmbedderND with 0 params. Trainable: False\n",
+-      "Initialized embedder #2: ConcatTimestepEmbedderND with 0 params. Trainable: False\n",
+-      "vector_suffix torch.Size([1, 1024])\n"
+-     ]
+-    }
+-   ],
++   "outputs": [],
+    "source": [
+     "# prep unCLIP\n",
+     "config = OmegaConf.load(\"generative_models/configs/unclip6.yaml\")\n",
+@@ -715,9 +655,9 @@
+  ],
+  "metadata": {
+   "kernelspec": {
+-   "display_name": "mindeye",
++   "display_name": "Python 3 (ipykernel)",
+    "language": "python",
+-   "name": "mindeye"
++   "name": "python3"
+   },
+   "language_info": {
+    "codemirror_mode": {

MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log

@@ -0,0 +1,19 @@
+1875
+testing-tt2 starting with epoch 0 / 150
+  0%|                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   | 0/150 [00:00<?, ?it/s]
+loss_RR 2.82318115234375
+loss_RR 2.54290771484375
+loss_RR 2.6309814453125
+loss_RR 2.68603515625
+loss_RR 2.663818359375
+loss_RR 2.65960693359375
+loss_RR 2.68072509765625
+loss_RR 2.656005859375
+loss_RR 2.9378662109375
+loss_RR 3.2581787109375
+loss_RR 2.9754638671875
+loss_RR 2.8961181640625
+loss_RR 2.99554443359375
+loss_RR 2.87652587890625
+
+  0%|                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   | 0/150 [08:49<?, ?it/s]

MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/requirements.txt

@@ -0,0 +1,214 @@
+Babel==2.15.0
+CoCa-pytorch==0.1.0
+GitPython==3.1.43
+Jinja2==3.1.4
+MarkupSafe==2.1.5
+PyYAML==6.0.1
+Pygments==2.18.0
+QtPy==2.4.1
+Send2Trash==1.8.3
+absl-py==2.1.0
+accelerate==0.24.1
+aiohttp==3.9.5
+aiosignal==1.3.1
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+anyio==4.4.0
+argon2-cffi-bindings==21.2.0
+argon2-cffi==23.1.0
+arrow==1.3.0
+asttokens==2.4.1
+async-lru==2.0.4
+attrs==23.2.0
+beartype==0.18.5
+beautifulsoup4==4.12.3
+bleach==6.1.0
+braceexpand==0.1.7
+certifi==2024.6.2
+cffi==1.16.0
+charset-normalizer==3.3.2
+click==8.1.7
+clip-anytorch==2.6.0
+clip==1.0
+comm==0.2.2
+contourpy==1.2.1
+cycler==0.12.1
+dalle2-pytorch==1.15.6
+datasets==2.20.0
+debugpy==1.8.1
+decorator==5.1.1
+deepspeed==0.13.1
+defusedxml==0.7.1
+diffusers==0.23.0
+dill==0.3.8
+docker-pycreds==0.4.0
+einops==0.8.0
+einx==0.3.0
+ema-pytorch==0.5.0
+embedding-reader==1.7.0
+evaluate==0.4.1
+executing==2.0.1
+fastjsonschema==2.20.0
+filelock==3.15.1
+fonttools==4.53.0
+fqdn==1.5.1
+frozendict==2.4.4
+frozenlist==1.4.1
+fsspec==2024.5.0
+ftfy==6.2.0
+gitdb==4.0.11
+h11==0.14.0
+h5py==3.10.0
+hjson==3.1.0
+httpcore==1.0.5
+httpx==0.27.0
+huggingface-hub==0.23.4
+idna==3.7
+imageio==2.34.1
+importlib_metadata==7.1.0
+ipykernel==6.29.4
+ipython==8.25.0
+ipywidgets==8.1.3
+isoduration==20.11.0
+jedi==0.19.1
+joblib==1.4.2
+json5==0.9.25
+jsonpointer==3.0.0
+jsonschema-specifications==2023.12.1
+jsonschema==4.22.0
+jupyter-console==6.6.3
+jupyter-events==0.10.0
+jupyter-lsp==2.2.5
+jupyter==1.0.0
+jupyter_client==8.6.2
+jupyter_core==5.7.2
+jupyter_server==2.14.1
+jupyter_server_terminals==0.5.3
+jupyterlab==4.2.2
+jupyterlab_nvdashboard==0.11.0
+jupyterlab_pygments==0.3.0
+jupyterlab_server==2.27.2
+jupyterlab_widgets==3.0.11
+kiwisolver==1.4.5
+kornia==0.7.1
+lazy_loader==0.4
+lightning-utilities==0.11.2
+matplotlib-inline==0.1.7
+matplotlib==3.8.2
+mistune==3.0.2
+mpmath==1.3.0
+multidict==6.0.5
+multiprocess==0.70.16
+nbclient==0.10.0
+nbconvert==7.16.4
+nbformat==5.10.4
+nest-asyncio==1.6.0
+networkx==3.3
+ninja==1.11.1.1
+nltk==3.8.1
+notebook==7.2.1
+notebook_shim==0.2.4
+numpy==1.26.4
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.18.1
+nvidia-nvjitlink-cu12==12.5.40
+nvidia-nvtx-cu12==12.1.105
+omegaconf==2.3.0
+open-clip-torch==2.24.0
+overrides==7.7.0
+packaging==24.1
+pandas==2.2.0
+pandocfilters==1.5.1
+parso==0.8.4
+pexpect==4.9.0
+pillow==10.3.0
+pip==24.0
+platformdirs==4.2.2
+prometheus_client==0.20.0
+prompt_toolkit==3.0.47
+protobuf==5.27.1
+psutil==5.9.8
+ptyprocess==0.7.0
+pure-eval==0.2.2
+py-cpuinfo==9.0.0
+pyarrow-hotfix==0.6
+pyarrow==15.0.2
+pycparser==2.22
+pydantic==2.7.4
+pydantic_core==2.18.4
+pynvml==11.5.0
+pyparsing==3.1.2
+python-dateutil==2.9.0.post0
+python-json-logger==2.0.7
+pytorch-lightning==2.0.1
+pytorch-warmup==0.1.1
+pytz==2024.1
+pyzmq==26.0.3
+qtconsole==5.5.2
+referencing==0.35.1
+regex==2024.5.15
+requests==2.32.3
+resize-right==0.0.2
+responses==0.18.0
+rfc3339-validator==0.1.4
+rfc3986-validator==0.1.1
+rotary-embedding-torch==0.6.2
+rouge_score==0.1.2
+rpds-py==0.18.1
+safetensors==0.4.3
+scikit-image==0.22.0
+scikit-learn==1.5.0
+scipy==1.13.1
+sentence-transformers==2.5.1
+sentencepiece==0.2.0
+sentry-sdk==2.5.1
+setproctitle==1.3.3
+setuptools==65.5.0
+six==1.16.0
+smmap==5.0.1
+sniffio==1.3.1
+soupsieve==2.5
+stack-data==0.6.3
+sympy==1.12.1
+terminado==0.18.1
+threadpoolctl==3.5.0
+tifffile==2024.5.22
+timm==1.0.3
+tinycss2==1.3.0
+tokenizers==0.15.2
+torch-fidelity==0.3.0
+torch==2.1.0
+torchmetrics==1.3.0.post0
+torchvision==0.16.0
+tornado==6.4.1
+tqdm==4.66.4
+traitlets==5.14.3
+transformers==4.37.2
+triton==2.1.0
+types-python-dateutil==2.9.0.20240316
+typing_extensions==4.12.2
+tzdata==2024.1
+umap==0.1.1
+uri-template==1.3.0
+urllib3==2.2.1
+vector-quantize-pytorch==1.14.24
+wandb==0.17.1
+wcwidth==0.2.13
+webcolors==24.6.0
+webdataset==0.2.73
+webencodings==0.5.1
+websocket-client==1.8.0
+widgetsnbextension==4.0.11
+x-clip==0.14.4
+xformers==0.0.22.post7
+xxhash==3.4.1
+yarl==1.9.4
+zipp==3.19.2

MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/wandb-metadata.json

@@ -0,0 +1,1035 @@
+{
+    "os": "Linux-5.15.0-1058-aws-x86_64-with-glibc2.31",
+    "python": "3.11.9",
+    "heartbeatAt": "2024-06-18T16:50:37.724201",
+    "startedAt": "2024-06-18T16:50:36.657596",
+    "docker": null,
+    "cuda": null,
+    "args": [],
+    "state": "running",
+    "program": "ckadirt/MindEyeV2/src/TrainCluster-Copy2.ipynb",
+    "codePathLocal": null,
+    "root": "/weka/proj-fmri/ckadirt/MindEyeV2",
+    "git": {
+        "remote": "https://github.com/MedARC-AI/MindEyeV2.git",
+        "commit": "a3588f0fadc81bdb0cc3e15993c01f8dd64dc7fe"
+    },
+    "email": null,
+    "host": "ip-10-0-159-181",
+    "username": "ckadirt",
+    "executable": "/admin/home-ckadirt/fmri/bin/python3.11",
+    "cpu_count": 96,
+    "cpu_count_logical": 192,
+    "cpu_freq": {
+        "current": 2737.5654166666754,
+        "min": 0.0,
+        "max": 0.0
+    },
+    "cpu_freq_per_core": [
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3599.552,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3571.159,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.82,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3599.455,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3598.843,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3598.67,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3599.596,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.56,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.635,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.882,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.88,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3504.795,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3548.68,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3597.822,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3593.276,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3596.593,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 3596.334,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        },
+        {
+            "current": 2649.998,
+            "min": 0.0,
+            "max": 0.0
+        }
+    ],
+    "disk": {
+        "/": {
+            "total": 232.4169158935547,
+            "used": 182.96668243408203
+        }
+    },
+    "gpu": "NVIDIA H100 80GB HBM3",
+    "gpu_count": 8,
+    "gpu_devices": [
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        },
+        {
+            "name": "NVIDIA H100 80GB HBM3",
+            "memory_total": 85520809984
+        }
+    ],
+    "memory": {
+        "total": 1999.962516784668
+    }
+}

MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/wandb-summary.json

@@ -0,0 +1 @@
+{"_wandb": {"runtime": 533}}

MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/logs/debug-internal.log

@@ -0,0 +1,432 @@
+2024-06-18 16:50:36,801 INFO    StreamThr :3360594 [internal.py:wandb_internal():85] W&B internal server running at pid: 3360594, started at: 2024-06-18 16:50:36.796840
+2024-06-18 16:50:36,801 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status
+2024-06-18 16:50:36,825 INFO    WriterThread:3360594 [datastore.py:open_for_write():87] open: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/run-testing-tt2.wandb
+2024-06-18 16:50:36,829 DEBUG   SenderThread:3360594 [sender.py:send():379] send: header
+2024-06-18 16:50:36,897 DEBUG   SenderThread:3360594 [sender.py:send():379] send: run
+2024-06-18 16:50:36,913 INFO    SenderThread:3360594 [sender.py:_setup_resume():749] checking resume status for None/mindeye_semantic_cluster/testing-tt2
+2024-06-18 16:50:37,461 INFO    SenderThread:3360594 [dir_watcher.py:__init__():211] watching files in: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files
+2024-06-18 16:50:37,462 INFO    SenderThread:3360594 [sender.py:_start_run_threads():1188] run started: testing-tt2 with start time 1718729436.80175
+2024-06-18 16:50:37,497 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: check_version
+2024-06-18 16:50:37,497 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: check_version
+2024-06-18 16:50:37,575 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: run_start
+2024-06-18 16:50:37,659 DEBUG   HandlerThread:3360594 [system_info.py:__init__():26] System info init
+2024-06-18 16:50:37,659 DEBUG   HandlerThread:3360594 [system_info.py:__init__():41] System info init done
+2024-06-18 16:50:37,659 INFO    HandlerThread:3360594 [system_monitor.py:start():194] Starting system monitor
+2024-06-18 16:50:37,659 INFO    SystemMonitor:3360594 [system_monitor.py:_start():158] Starting system asset monitoring threads
+2024-06-18 16:50:37,659 INFO    HandlerThread:3360594 [system_monitor.py:probe():214] Collecting system info
+2024-06-18 16:50:37,660 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started cpu monitoring
+2024-06-18 16:50:37,660 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started disk monitoring
+2024-06-18 16:50:37,661 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started gpu monitoring
+2024-06-18 16:50:37,661 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started memory monitoring
+2024-06-18 16:50:37,661 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started network monitoring
+2024-06-18 16:50:37,724 DEBUG   HandlerThread:3360594 [system_info.py:probe():152] Probing system
+2024-06-18 16:50:37,761 DEBUG   HandlerThread:3360594 [system_info.py:_probe_git():137] Probing git
+2024-06-18 16:50:37,785 DEBUG   HandlerThread:3360594 [system_info.py:_probe_git():145] Probing git done
+2024-06-18 16:50:37,785 DEBUG   HandlerThread:3360594 [system_info.py:probe():200] Probing system done
+2024-06-18 16:50:37,785 DEBUG   HandlerThread:3360594 [system_monitor.py:probe():223] {'os': 'Linux-5.15.0-1058-aws-x86_64-with-glibc2.31', 'python': '3.11.9', 'heartbeatAt': '2024-06-18T16:50:37.724201', 'startedAt': '2024-06-18T16:50:36.657596', 'docker': None, 'cuda': None, 'args': (), 'state': 'running', 'program': 'ckadirt/MindEyeV2/src/TrainCluster-Copy2.ipynb', 'codePathLocal': None, 'root': '/weka/proj-fmri/ckadirt/MindEyeV2', 'git': {'remote': 'https://github.com/MedARC-AI/MindEyeV2.git', 'commit': 'a3588f0fadc81bdb0cc3e15993c01f8dd64dc7fe'}, 'email': None, 'host': 'ip-10-0-159-181', 'username': 'ckadirt', 'executable': '/admin/home-ckadirt/fmri/bin/python3.11', 'cpu_count': 96, 'cpu_count_logical': 192, 'cpu_freq': {'current': 2737.5654166666754, 'min': 0.0, 'max': 0.0}, 'cpu_freq_per_core': [{'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 3599.552, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 3571.159, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 3597.82, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 3599.455, 'min': 0.0, 'max': 0.0}, {'current': 3598.843, 'min': 0.0, 'max': 0.0}, {'current': 3598.67, 'min': 0.0, 'max': 0.0}, {'current': 3599.596, 'min': 0.0, 'max': 0.0}, {'current': 3597.56, 'min': 0.0, 'max': 0.0}, {'current': 3597.635, 'min': 0.0, 'max': 0.0}, {'current': 3597.882, 'min': 0.0, 'max': 0.0}, {'current': 3597.88, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 3504.795, 'min': 0.0, 'max': 0.0}, {'current': 3548.68, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 3597.822, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 3593.276, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 3596.593, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 3596.334, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}, {'current': 2649.998, 'min': 0.0, 'max': 0.0}], 'disk': {'/': {'total': 232.4169158935547, 'used': 182.96668243408203}}, 'gpu': 'NVIDIA H100 80GB HBM3', 'gpu_count': 8, 'gpu_devices': [{'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}, {'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}, {'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}, {'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}, {'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}, {'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}, {'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}, {'name': 'NVIDIA H100 80GB HBM3', 'memory_total': 85520809984}], 'memory': {'total': 1999.962516784668}}
+2024-06-18 16:50:37,785 INFO    HandlerThread:3360594 [system_monitor.py:probe():224] Finished collecting system info
+2024-06-18 16:50:37,785 INFO    HandlerThread:3360594 [system_monitor.py:probe():227] Publishing system info
+2024-06-18 16:50:37,786 DEBUG   HandlerThread:3360594 [system_info.py:_save_code():44] Saving code
+2024-06-18 16:50:37,786 WARNING HandlerThread:3360594 [system_info.py:_save_code():46] unable to save code -- program entry not found
+2024-06-18 16:50:37,786 DEBUG   HandlerThread:3360594 [system_info.py:_save_patches():84] Saving git patches
+2024-06-18 16:50:38,466 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/diff.patch
+2024-06-18 16:50:38,528 DEBUG   HandlerThread:3360594 [system_info.py:_save_patches():126] Saving git patches done
+2024-06-18 16:50:38,539 INFO    HandlerThread:3360594 [system_monitor.py:probe():229] Finished publishing system info
+2024-06-18 16:50:38,551 DEBUG   SenderThread:3360594 [sender.py:send():379] send: files
+2024-06-18 16:50:38,551 INFO    SenderThread:3360594 [sender.py:_save_file():1454] saving file wandb-metadata.json with policy now
+2024-06-18 16:50:38,551 INFO    SenderThread:3360594 [sender.py:_save_file():1454] saving file diff.patch with policy now
+2024-06-18 16:50:39,061 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: python_packages
+2024-06-18 16:50:39,061 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: python_packages
+2024-06-18 16:50:39,062 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:50:39,063 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: internal_messages
+2024-06-18 16:50:39,071 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:50:39,073 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: pause
+2024-06-18 16:50:39,073 INFO    HandlerThread:3360594 [handler.py:handle_request_pause():724] stopping system metrics thread
+2024-06-18 16:50:39,073 INFO    HandlerThread:3360594 [system_monitor.py:finish():203] Stopping system monitor
+2024-06-18 16:50:39,073 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():172] Starting system metrics aggregation loop
+2024-06-18 16:50:39,074 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():179] Finished system metrics aggregation loop
+2024-06-18 16:50:39,074 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined cpu monitor
+2024-06-18 16:50:39,074 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():183] Publishing last batch of metrics
+2024-06-18 16:50:39,074 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined disk monitor
+2024-06-18 16:50:39,088 INFO    wandb-upload_0:3360594 [upload_job.py:push():130] Uploaded file /tmp/tmp1uqtk8vuwandb/mfd8w0e8-wandb-metadata.json
+2024-06-18 16:50:39,247 DEBUG   SenderThread:3360594 [sender.py:send():379] send: telemetry
+2024-06-18 16:50:39,298 INFO    wandb-upload_1:3360594 [upload_job.py:push():130] Uploaded file /tmp/tmp1uqtk8vuwandb/z6tbxjy6-diff.patch
+2024-06-18 16:50:39,471 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/wandb-metadata.json
+2024-06-18 16:50:39,471 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/requirements.txt
+2024-06-18 16:50:41,255 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined gpu monitor
+2024-06-18 16:50:41,255 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined memory monitor
+2024-06-18 16:50:41,255 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined network monitor
+2024-06-18 16:50:41,256 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: resume
+2024-06-18 16:50:41,256 INFO    HandlerThread:3360594 [handler.py:handle_request_resume():715] starting system metrics thread
+2024-06-18 16:50:41,256 INFO    HandlerThread:3360594 [system_monitor.py:start():194] Starting system monitor
+2024-06-18 16:50:41,256 INFO    SystemMonitor:3360594 [system_monitor.py:_start():158] Starting system asset monitoring threads
+2024-06-18 16:50:41,256 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:50:41,256 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: pause
+2024-06-18 16:50:41,257 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started cpu monitoring
+2024-06-18 16:50:41,258 INFO    HandlerThread:3360594 [handler.py:handle_request_pause():724] stopping system metrics thread
+2024-06-18 16:50:41,258 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started disk monitoring
+2024-06-18 16:50:41,259 INFO    HandlerThread:3360594 [system_monitor.py:finish():203] Stopping system monitor
+2024-06-18 16:50:41,259 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started gpu monitoring
+2024-06-18 16:50:41,259 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():172] Starting system metrics aggregation loop
+2024-06-18 16:50:41,259 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():179] Finished system metrics aggregation loop
+2024-06-18 16:50:41,259 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():183] Publishing last batch of metrics
+2024-06-18 16:50:41,261 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined cpu monitor
+2024-06-18 16:50:41,263 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined disk monitor
+2024-06-18 16:50:43,215 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined gpu monitor
+2024-06-18 16:50:43,215 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: resume
+2024-06-18 16:50:43,215 INFO    HandlerThread:3360594 [handler.py:handle_request_resume():715] starting system metrics thread
+2024-06-18 16:50:43,215 INFO    HandlerThread:3360594 [system_monitor.py:start():194] Starting system monitor
+2024-06-18 16:50:43,216 INFO    SystemMonitor:3360594 [system_monitor.py:_start():158] Starting system asset monitoring threads
+2024-06-18 16:50:43,216 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: pause
+2024-06-18 16:50:43,216 INFO    HandlerThread:3360594 [handler.py:handle_request_pause():724] stopping system metrics thread
+2024-06-18 16:50:43,216 INFO    HandlerThread:3360594 [system_monitor.py:finish():203] Stopping system monitor
+2024-06-18 16:50:43,216 WARNING HandlerThread:3360594 [interfaces.py:finish():205] Failed to finish cpu monitoring: cannot join thread before it is started
+2024-06-18 16:50:43,216 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started cpu monitoring
+2024-06-18 16:50:43,216 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():172] Starting system metrics aggregation loop
+2024-06-18 16:50:43,216 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():179] Finished system metrics aggregation loop
+2024-06-18 16:50:43,216 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():183] Publishing last batch of metrics
+2024-06-18 16:50:43,218 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: resume
+2024-06-18 16:50:43,218 INFO    HandlerThread:3360594 [handler.py:handle_request_resume():715] starting system metrics thread
+2024-06-18 16:50:43,218 INFO    HandlerThread:3360594 [system_monitor.py:start():194] Starting system monitor
+2024-06-18 16:50:43,218 INFO    SystemMonitor:3360594 [system_monitor.py:_start():158] Starting system asset monitoring threads
+2024-06-18 16:50:43,218 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: pause
+2024-06-18 16:50:43,218 INFO    HandlerThread:3360594 [handler.py:handle_request_pause():724] stopping system metrics thread
+2024-06-18 16:50:43,218 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started cpu monitoring
+2024-06-18 16:50:43,219 INFO    HandlerThread:3360594 [system_monitor.py:finish():203] Stopping system monitor
+2024-06-18 16:50:43,219 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started disk monitoring
+2024-06-18 16:50:43,219 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():172] Starting system metrics aggregation loop
+2024-06-18 16:50:43,219 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():179] Finished system metrics aggregation loop
+2024-06-18 16:50:43,219 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():183] Publishing last batch of metrics
+2024-06-18 16:50:43,222 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined cpu monitor
+2024-06-18 16:50:43,223 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined disk monitor
+2024-06-18 16:50:43,223 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: resume
+2024-06-18 16:50:43,223 INFO    HandlerThread:3360594 [handler.py:handle_request_resume():715] starting system metrics thread
+2024-06-18 16:50:43,223 INFO    HandlerThread:3360594 [system_monitor.py:start():194] Starting system monitor
+2024-06-18 16:50:43,223 INFO    SystemMonitor:3360594 [system_monitor.py:_start():158] Starting system asset monitoring threads
+2024-06-18 16:50:43,223 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: pause
+2024-06-18 16:50:43,223 INFO    HandlerThread:3360594 [handler.py:handle_request_pause():724] stopping system metrics thread
+2024-06-18 16:50:43,223 INFO    HandlerThread:3360594 [system_monitor.py:finish():203] Stopping system monitor
+2024-06-18 16:50:43,223 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started cpu monitoring
+2024-06-18 16:50:43,231 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined cpu monitor
+2024-06-18 16:50:43,231 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():172] Starting system metrics aggregation loop
+2024-06-18 16:50:43,237 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():179] Finished system metrics aggregation loop
+2024-06-18 16:50:43,237 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():183] Publishing last batch of metrics
+2024-06-18 16:50:43,243 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: resume
+2024-06-18 16:50:43,243 INFO    HandlerThread:3360594 [handler.py:handle_request_resume():715] starting system metrics thread
+2024-06-18 16:50:43,243 INFO    HandlerThread:3360594 [system_monitor.py:start():194] Starting system monitor
+2024-06-18 16:50:43,248 INFO    SystemMonitor:3360594 [system_monitor.py:_start():158] Starting system asset monitoring threads
+2024-06-18 16:50:43,248 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: internal_messages
+2024-06-18 16:50:43,254 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:50:43,260 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started cpu monitoring
+2024-06-18 16:50:43,260 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:50:43,261 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:50:43,263 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:50:43,263 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started disk monitoring
+2024-06-18 16:50:43,276 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started gpu monitoring
+2024-06-18 16:50:43,277 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started memory monitoring
+2024-06-18 16:50:43,278 INFO    SystemMonitor:3360594 [interfaces.py:start():188] Started network monitoring
+2024-06-18 16:50:43,522 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:50:45,550 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:50:48,063 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:50:53,063 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:50:54,062 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:50:54,062 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:50:58,239 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:51:04,063 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:51:09,062 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:51:09,062 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:51:09,276 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:51:09,961 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/config.yaml
+2024-06-18 16:51:15,064 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:51:20,064 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:51:24,062 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:51:24,062 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:51:25,240 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:51:31,065 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:51:36,065 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:51:39,062 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:51:39,062 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:51:41,240 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:51:43,278 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():172] Starting system metrics aggregation loop
+2024-06-18 16:51:43,280 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:51:47,065 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:51:52,066 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:51:54,062 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:51:54,062 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:51:57,241 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:52:03,066 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:52:08,066 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:52:09,062 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:52:09,063 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:52:13,241 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:52:13,298 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:52:19,066 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:52:24,062 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:52:24,063 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:52:24,240 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:52:30,066 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:52:35,068 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:52:39,063 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:52:39,063 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:52:40,241 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:52:43,283 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:52:46,067 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:52:52,067 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:52:54,063 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:52:54,063 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:52:57,240 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:53:03,068 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:53:09,063 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:53:09,063 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:53:09,241 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:53:13,285 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:53:15,068 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:53:21,068 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:53:24,063 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:53:24,063 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:53:26,241 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:53:32,068 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:53:37,068 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:53:39,063 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:53:39,063 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:53:42,241 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:53:43,289 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:53:48,068 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:53:53,070 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:53:54,063 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:53:54,063 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:53:58,241 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:54:04,069 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:54:09,063 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:54:09,063 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:54:09,240 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:54:13,293 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:54:15,069 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:54:20,069 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:54:24,063 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:54:24,064 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:54:25,241 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:54:31,069 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:54:36,070 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:54:39,064 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:54:39,064 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:54:41,342 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:54:43,292 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:54:47,070 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:54:52,070 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:54:54,064 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:54:54,064 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:54:57,241 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:55:03,075 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:55:09,064 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:55:09,064 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:55:09,254 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:55:13,294 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:55:14,062 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:55:15,070 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:55:16,149 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:55:20,070 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:55:24,064 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:55:24,064 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:55:25,242 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:55:29,380 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:55:31,071 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:55:35,464 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:55:36,071 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:55:39,064 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:55:39,065 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:55:41,242 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:55:43,295 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:55:43,606 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:55:47,071 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:55:52,140 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:55:53,773 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:55:54,064 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:55:54,064 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:55:55,798 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:55:57,242 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:56:03,072 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:56:05,966 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:56:08,396 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:56:09,064 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:56:09,065 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:56:13,299 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:56:14,071 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:56:19,072 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:56:24,065 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:56:24,065 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:56:24,242 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:56:30,072 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:56:33,425 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:56:36,072 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:56:39,065 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:56:39,065 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:56:41,243 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:56:43,298 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:56:46,588 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:56:47,652 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:56:52,072 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:56:54,065 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:56:54,065 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:56:57,242 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:57:03,073 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:57:09,065 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:57:09,065 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:57:09,242 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:57:13,300 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:57:15,073 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:57:20,844 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:57:22,223 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:57:24,065 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:57:24,066 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:57:26,243 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:57:29,350 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:57:32,073 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:57:38,073 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:57:39,066 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:57:39,066 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:57:43,270 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:57:43,301 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:57:47,630 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:57:48,728 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:57:53,734 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:57:54,056 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:57:54,066 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:57:54,066 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:57:59,073 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:58:04,074 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:58:09,066 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:58:09,066 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:58:09,168 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:58:13,305 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:58:15,074 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:58:20,074 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:58:24,066 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:58:24,066 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:58:25,245 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:58:31,075 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:58:37,075 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:58:39,066 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:58:39,066 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:58:42,244 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:58:43,306 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:58:48,075 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:58:53,078 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:58:54,066 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:58:54,067 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:58:58,244 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:59:04,079 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:59:09,067 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:59:09,068 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:59:09,246 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:59:13,307 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:59:15,079 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:59:21,079 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:59:24,067 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: stop_status
+2024-06-18 16:59:24,067 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: stop_status
+2024-06-18 16:59:26,244 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:59:29,491 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: pause
+2024-06-18 16:59:29,491 INFO    HandlerThread:3360594 [handler.py:handle_request_pause():724] stopping system metrics thread
+2024-06-18 16:59:29,492 INFO    HandlerThread:3360594 [system_monitor.py:finish():203] Stopping system monitor
+2024-06-18 16:59:29,492 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():179] Finished system metrics aggregation loop
+2024-06-18 16:59:29,492 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined cpu monitor
+2024-06-18 16:59:29,493 DEBUG   SystemMonitor:3360594 [system_monitor.py:_start():183] Publishing last batch of metrics
+2024-06-18 16:59:29,493 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined disk monitor
+2024-06-18 16:59:30,299 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:59:31,533 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined gpu monitor
+2024-06-18 16:59:31,533 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined memory monitor
+2024-06-18 16:59:31,533 INFO    HandlerThread:3360594 [interfaces.py:finish():200] Joined network monitor
+2024-06-18 16:59:31,533 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: poll_exit
+2024-06-18 16:59:31,534 DEBUG   SenderThread:3360594 [sender.py:send():379] send: stats
+2024-06-18 16:59:31,534 DEBUG   SenderThread:3360594 [sender.py:send():379] send: exit
+2024-06-18 16:59:31,534 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:59:31,534 INFO    SenderThread:3360594 [sender.py:send_exit():586] handling exit code: 0
+2024-06-18 16:59:31,535 INFO    SenderThread:3360594 [sender.py:send_exit():588] handling runtime: 533
+2024-06-18 16:59:31,539 INFO    SenderThread:3360594 [sender.py:_save_file():1454] saving file wandb-summary.json with policy end
+2024-06-18 16:59:31,539 INFO    SenderThread:3360594 [sender.py:send_exit():594] send defer
+2024-06-18 16:59:31,539 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: poll_exit
+2024-06-18 16:59:31,539 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: defer
+2024-06-18 16:59:31,540 INFO    HandlerThread:3360594 [handler.py:handle_request_defer():184] handle defer: 0
+2024-06-18 16:59:31,540 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: defer
+2024-06-18 16:59:31,540 INFO    SenderThread:3360594 [sender.py:send_request_defer():610] handle sender defer: 0
+2024-06-18 16:59:31,540 INFO    SenderThread:3360594 [sender.py:transition_state():614] send defer: 1
+2024-06-18 16:59:31,540 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: defer
+2024-06-18 16:59:31,540 INFO    HandlerThread:3360594 [handler.py:handle_request_defer():184] handle defer: 1
+2024-06-18 16:59:31,540 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: defer
+2024-06-18 16:59:31,540 INFO    SenderThread:3360594 [sender.py:send_request_defer():610] handle sender defer: 1
+2024-06-18 16:59:31,540 INFO    SenderThread:3360594 [sender.py:transition_state():614] send defer: 2
+2024-06-18 16:59:31,540 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: defer
+2024-06-18 16:59:31,540 INFO    HandlerThread:3360594 [handler.py:handle_request_defer():184] handle defer: 2
+2024-06-18 16:59:31,540 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: defer
+2024-06-18 16:59:31,540 INFO    SenderThread:3360594 [sender.py:send_request_defer():610] handle sender defer: 2
+2024-06-18 16:59:31,540 INFO    SenderThread:3360594 [sender.py:transition_state():614] send defer: 3
+2024-06-18 16:59:31,540 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: defer
+2024-06-18 16:59:31,540 INFO    HandlerThread:3360594 [handler.py:handle_request_defer():184] handle defer: 3
+2024-06-18 16:59:31,540 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: defer
+2024-06-18 16:59:31,540 INFO    SenderThread:3360594 [sender.py:send_request_defer():610] handle sender defer: 3
+2024-06-18 16:59:31,540 INFO    SenderThread:3360594 [sender.py:transition_state():614] send defer: 4
+2024-06-18 16:59:31,540 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: defer
+2024-06-18 16:59:31,541 INFO    HandlerThread:3360594 [handler.py:handle_request_defer():184] handle defer: 4
+2024-06-18 16:59:31,541 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: defer
+2024-06-18 16:59:31,541 INFO    SenderThread:3360594 [sender.py:send_request_defer():610] handle sender defer: 4
+2024-06-18 16:59:31,541 INFO    SenderThread:3360594 [sender.py:transition_state():614] send defer: 5
+2024-06-18 16:59:31,541 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: defer
+2024-06-18 16:59:31,541 INFO    HandlerThread:3360594 [handler.py:handle_request_defer():184] handle defer: 5
+2024-06-18 16:59:31,541 DEBUG   SenderThread:3360594 [sender.py:send():379] send: summary
+2024-06-18 16:59:31,552 INFO    SenderThread:3360594 [sender.py:_save_file():1454] saving file wandb-summary.json with policy end
+2024-06-18 16:59:31,552 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: defer
+2024-06-18 16:59:31,552 INFO    SenderThread:3360594 [sender.py:send_request_defer():610] handle sender defer: 5
+2024-06-18 16:59:31,552 INFO    SenderThread:3360594 [sender.py:transition_state():614] send defer: 6
+2024-06-18 16:59:31,552 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: defer
+2024-06-18 16:59:31,552 INFO    HandlerThread:3360594 [handler.py:handle_request_defer():184] handle defer: 6
+2024-06-18 16:59:31,552 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: defer
+2024-06-18 16:59:31,552 INFO    SenderThread:3360594 [sender.py:send_request_defer():610] handle sender defer: 6
+2024-06-18 16:59:31,552 INFO    SenderThread:3360594 [sender.py:transition_state():614] send defer: 7
+2024-06-18 16:59:31,552 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: status_report
+2024-06-18 16:59:31,552 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: defer
+2024-06-18 16:59:31,552 INFO    HandlerThread:3360594 [handler.py:handle_request_defer():184] handle defer: 7
+2024-06-18 16:59:31,552 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: defer
+2024-06-18 16:59:31,552 INFO    SenderThread:3360594 [sender.py:send_request_defer():610] handle sender defer: 7
+2024-06-18 16:59:32,325 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_created():271] file/dir created: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/wandb-summary.json
+2024-06-18 16:59:32,499 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: poll_exit
+2024-06-18 16:59:33,598 INFO    SenderThread:3360594 [sender.py:transition_state():614] send defer: 8
+2024-06-18 16:59:33,598 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: poll_exit
+2024-06-18 16:59:33,598 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: defer
+2024-06-18 16:59:33,599 INFO    HandlerThread:3360594 [handler.py:handle_request_defer():184] handle defer: 8
+2024-06-18 16:59:33,599 DEBUG   SenderThread:3360594 [sender.py:send_request():406] send_request: defer
+2024-06-18 16:59:33,599 INFO    SenderThread:3360594 [sender.py:send_request_defer():610] handle sender defer: 8
+2024-06-18 16:59:33,599 INFO    SenderThread:3360594 [job_builder.py:build():440] Attempting to build job artifact
+2024-06-18 16:59:33,602 INFO    SenderThread:3360594 [job_builder.py:_get_source_type():569] is repo sourced job
+2024-06-18 16:59:33,602 INFO    SenderThread:3360594 [job_builder.py:_get_program_relpath():587] run is notebook based run
+2024-06-18 16:59:33,629 INFO    SenderThread:3360594 [job_builder.py:build():545] adding wandb-job metadata file
+2024-06-18 16:59:33,659 INFO    SenderThread:3360594 [sender.py:transition_state():614] send defer: 9
+2024-06-18 16:59:33,659 DEBUG   SenderThread:3360594 [sender.py:send():379] send: artifact
+2024-06-18 16:59:33,659 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: defer
+2024-06-18 16:59:33,661 INFO    HandlerThread:3360594 [handler.py:handle_request_defer():184] handle defer: 9
+2024-06-18 16:59:33,856 WARNING StreamThr :3360594 [internal.py:is_dead():413] Internal process exiting, parent pid 3353055 disappeared
+2024-06-18 16:59:33,856 ERROR   StreamThr :3360594 [internal.py:wandb_internal():151] Internal process shutdown.
+2024-06-18 16:59:34,349 INFO    Thread-12 :3360594 [dir_watcher.py:_on_file_modified():288] file/dir modified: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:59:34,474 INFO    wandb-upload_0:3360594 [upload_job.py:push():88] Uploaded file /admin/home-ckadirt/.local/share/wandb/artifacts/staging/tmpwylhto1m
+2024-06-18 16:59:34,497 INFO    wandb-upload_1:3360594 [upload_job.py:push():88] Uploaded file /tmp/tmpg9bai4xg/wandb-job.json
+2024-06-18 16:59:34,500 DEBUG   HandlerThread:3360594 [handler.py:handle_request():158] handle_request: poll_exit
+2024-06-18 16:59:34,500 INFO    HandlerThread:3360594 [handler.py:finish():882] shutting down handler
+2024-06-18 16:59:34,500 INFO    WriterThread:3360594 [datastore.py:close():296] close: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/run-testing-tt2.wandb
+2024-06-18 16:59:34,579 INFO    wandb-upload_2:3360594 [upload_job.py:push():88] Uploaded file /admin/home-ckadirt/.local/share/wandb/artifacts/staging/tmpf2z3dfgr
+2024-06-18 16:59:35,132 INFO    SenderThread:3360594 [sender.py:send_artifact():1530] sent artifact job-https___github.com_MedARC-AI_MindEyeV2.git_ckadirt_MindEyeV2_src_TrainCluster-Copy2.ipynb - {'id': 'QXJ0aWZhY3Q6ODczNzEy', 'state': 'PENDING', 'artifactSequence': {'id': 'QXJ0aWZhY3RDb2xsZWN0aW9uOjQ2NzQw', 'latestArtifact': None}}
+2024-06-18 16:59:35,132 INFO    SenderThread:3360594 [sender.py:finish():1608] shutting down sender
+2024-06-18 16:59:35,132 INFO    SenderThread:3360594 [dir_watcher.py:finish():358] shutting down directory watcher
+2024-06-18 16:59:35,363 INFO    SenderThread:3360594 [dir_watcher.py:finish():388] scan: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files
+2024-06-18 16:59:35,365 INFO    SenderThread:3360594 [dir_watcher.py:finish():402] scan save: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/config.yaml config.yaml
+2024-06-18 16:59:35,365 INFO    SenderThread:3360594 [dir_watcher.py:finish():402] scan save: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log output.log
+2024-06-18 16:59:35,365 INFO    SenderThread:3360594 [dir_watcher.py:finish():402] scan save: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/wandb-summary.json wandb-summary.json
+2024-06-18 16:59:35,366 INFO    SenderThread:3360594 [dir_watcher.py:finish():402] scan save: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/wandb-metadata.json wandb-metadata.json
+2024-06-18 16:59:35,368 INFO    SenderThread:3360594 [dir_watcher.py:finish():402] scan save: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/diff.patch diff.patch
+2024-06-18 16:59:35,368 INFO    SenderThread:3360594 [dir_watcher.py:finish():402] scan save: /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/requirements.txt requirements.txt
+2024-06-18 16:59:35,369 INFO    SenderThread:3360594 [file_pusher.py:finish():169] shutting down file pusher
+2024-06-18 16:59:35,369 INFO    SenderThread:3360594 [file_pusher.py:join():175] waiting for file pusher
+2024-06-18 16:59:35,541 INFO    wandb-upload_0:3360594 [upload_job.py:push():130] Uploaded file /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/config.yaml
+2024-06-18 16:59:35,722 INFO    wandb-upload_1:3360594 [upload_job.py:push():130] Uploaded file /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/wandb-summary.json
+2024-06-18 16:59:35,731 INFO    wandb-upload_2:3360594 [upload_job.py:push():130] Uploaded file /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/output.log
+2024-06-18 16:59:35,756 INFO    wandb-upload_3:3360594 [upload_job.py:push():130] Uploaded file /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/files/requirements.txt
+2024-06-18 16:59:36,370 INFO    SenderThread:3360594 [file_stream.py:finish():601] file stream finish called
+2024-06-18 16:59:36,602 INFO    SenderThread:3360594 [file_stream.py:finish():605] file stream finish is done

MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/logs/debug.log

@@ -0,0 +1,48 @@
+2024-06-18 16:50:36,773 INFO    MainThread:3353055 [wandb_setup.py:_flush():76] Current SDK version is 0.17.1
+2024-06-18 16:50:36,773 INFO    MainThread:3353055 [wandb_setup.py:_flush():76] Configure stats pid to 3353055
+2024-06-18 16:50:36,773 INFO    MainThread:3353055 [wandb_setup.py:_flush():76] Loading settings from /admin/home-ckadirt/.config/wandb/settings
+2024-06-18 16:50:36,773 INFO    MainThread:3353055 [wandb_setup.py:_flush():76] Loading settings from /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/settings
+2024-06-18 16:50:36,773 INFO    MainThread:3353055 [wandb_setup.py:_flush():76] Loading settings from environment variables: {}
+2024-06-18 16:50:36,773 INFO    MainThread:3353055 [wandb_setup.py:_flush():76] Inferring run settings from compute environment: {'program': '<python with no main file>'}
+2024-06-18 16:50:36,773 INFO    MainThread:3353055 [wandb_setup.py:_flush():76] Applying login settings: {'base_url': 'https://stability.wandb.io'}
+2024-06-18 16:50:36,773 INFO    MainThread:3353055 [wandb_setup.py:_flush():76] Applying login settings: {}
+2024-06-18 16:50:36,774 INFO    MainThread:3353055 [wandb_init.py:_log_setup():520] Logging user logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/logs/debug.log
+2024-06-18 16:50:36,774 INFO    MainThread:3353055 [wandb_init.py:_log_setup():521] Logging internal logs to /weka/proj-fmri/ckadirt/MindEyeV2/src/wandb/run-20240618_165036-testing-tt2/logs/debug-internal.log
+2024-06-18 16:50:36,774 INFO    MainThread:3353055 [wandb_init.py:_jupyter_setup():466] configuring jupyter hooks <wandb.sdk.wandb_init._WandbInit object at 0x7f4377427f10>
+2024-06-18 16:50:36,775 INFO    MainThread:3353055 [wandb_init.py:init():560] calling init triggers
+2024-06-18 16:50:36,775 INFO    MainThread:3353055 [wandb_init.py:init():567] wandb.init called with sweep_config: {}
+config: {'model_name': 'testing-tt2', 'global_batch_size': 16, 'batch_size': 16, 'num_epochs': 150, 'num_sessions': 40, 'num_params': 486928065, 'clip_scale': 1.0, 'prior_scale': 30.0, 'blur_scale': 0.5, 'use_image_aug': False, 'max_lr': 3e-05, 'mixup_pct': 0.33, 'num_samples_per_epoch': 30000, 'num_test': 3000, 'ckpt_interval': 999, 'ckpt_saving': False, 'seed': 42, 'distributed': False, 'num_devices': 1, 'world_size': 1, 'train_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/train/{0..39}.tar', 'test_url': '/weka/proj-medarc/shared/mindeyev2_dataset/wds/subj01/new_test/0.tar'}
+2024-06-18 16:50:36,775 INFO    MainThread:3353055 [wandb_init.py:init():610] starting backend
+2024-06-18 16:50:36,775 INFO    MainThread:3353055 [wandb_init.py:init():614] setting up manager
+2024-06-18 16:50:36,796 INFO    MainThread:3353055 [backend.py:_multiprocessing_setup():105] multiprocessing start_methods=fork,spawn,forkserver, using: spawn
+2024-06-18 16:50:36,800 INFO    MainThread:3353055 [wandb_init.py:init():622] backend started and connected
+2024-06-18 16:50:36,824 INFO    MainThread:3353055 [wandb_run.py:_label_probe_notebook():1334] probe notebook
+2024-06-18 16:50:36,825 INFO    MainThread:3353055 [wandb_run.py:_label_probe_notebook():1344] Unable to probe notebook: 'NoneType' object has no attribute 'get'
+2024-06-18 16:50:36,825 INFO    MainThread:3353055 [wandb_init.py:init():711] updated telemetry
+2024-06-18 16:50:36,896 INFO    MainThread:3353055 [wandb_init.py:init():744] communicating run to backend with 90.0 second timeout
+2024-06-18 16:50:37,496 INFO    MainThread:3353055 [wandb_run.py:_on_init():2402] communicating current version
+2024-06-18 16:50:37,558 INFO    MainThread:3353055 [wandb_run.py:_on_init():2411] got version response upgrade_message: "wandb version 0.17.2 is available!  To upgrade, please run:\n $ pip install wandb --upgrade"
+
+2024-06-18 16:50:37,558 INFO    MainThread:3353055 [wandb_init.py:init():795] starting run threads in backend
+2024-06-18 16:50:39,062 INFO    MainThread:3353055 [wandb_run.py:_console_start():2380] atexit reg
+2024-06-18 16:50:39,062 INFO    MainThread:3353055 [wandb_run.py:_redirect():2235] redirect: wrap_raw
+2024-06-18 16:50:39,062 INFO    MainThread:3353055 [wandb_run.py:_redirect():2300] Wrapping output streams.
+2024-06-18 16:50:39,062 INFO    MainThread:3353055 [wandb_run.py:_redirect():2325] Redirects installed.
+2024-06-18 16:50:39,068 INFO    MainThread:3353055 [wandb_init.py:init():838] run started, returning control to user process
+2024-06-18 16:50:39,072 INFO    MainThread:3353055 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster-Copy2.ipynb
+2024-06-18 16:50:39,073 INFO    MainThread:3353055 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 16:50:39,157 INFO    MainThread:3353055 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 16:50:39,158 INFO    MainThread:3353055 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster-Copy2.ipynb
+2024-06-18 16:50:39,159 INFO    MainThread:3353055 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 16:50:39,237 INFO    MainThread:3353055 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 16:50:39,249 INFO    MainThread:3353055 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster-Copy2.ipynb
+2024-06-18 16:50:39,253 INFO    MainThread:3353055 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 16:50:39,338 INFO    MainThread:3353055 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 16:50:39,509 INFO    MainThread:3353055 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster-Copy2.ipynb
+2024-06-18 16:50:39,509 INFO    MainThread:3353055 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 16:50:39,592 INFO    MainThread:3353055 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 16:50:39,593 INFO    MainThread:3353055 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster-Copy2.ipynb
+2024-06-18 16:50:39,594 INFO    MainThread:3353055 [wandb_init.py:_pause_backend():431] pausing backend
+2024-06-18 16:50:39,686 INFO    MainThread:3353055 [wandb_init.py:_resume_backend():436] resuming backend
+2024-06-18 16:59:29,490 INFO    MainThread:3353055 [jupyter.py:_save_ipynb():383] looking for notebook: ckadirt/MindEyeV2/src/TrainCluster-Copy2.ipynb
+2024-06-18 16:59:29,491 INFO    MainThread:3353055 [wandb_init.py:_pause_backend():431] pausing backend