Upload folder using huggingface_hub

.gitignore

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

LICENSE

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

README.md

@@ -0,0 +1,11 @@
+# MindEyeV2
+
+In-progress
+
+1. Download all of https://huggingface.co/datasets/pscotti/mindeyev2 and place them in a folder. You will need to specify the path to this folder as "data_path" variable.
+
+2. Run setup.sh to install a new "fmri" conda environment.
+
+3. Activate the conda environment with "conda activate fmri"
+
+4. Run Train.ipynb or Train.py (they are the same code)

src/Train.py

@@ -0,0 +1,708 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+# # Code to convert this notebook to .py if you want to run it via command line or with Slurm
+# from subprocess import call
+# command = "jupyter nbconvert Train.ipynb --to python"
+# call(command,shell=True)
+
+
+# # Import packages & functions
+
+# In[2]:
+
+
+import os
+import sys
+import json
+import argparse
+import numpy as np
+import time
+import random
+import h5py
+from tqdm import tqdm
+
+import webdataset as wds
+import gc
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torchvision import transforms
+
+# tf32 data type is faster than standard float32
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# custom functions #
+import utils
+
+
+# In[ ]:
+
+
+local_rank = os.getenv('RANK')
+if local_rank is None: 
+    local_rank = 0
+else:
+    local_rank = int(local_rank)
+print("LOCAL RANK ", local_rank)
+
+### Single-GPU config ###
+## Feel free to uncomment the below 4 lines and comment out all the multi-gpu config code to simplify things for single-gpu
+# from accelerate import Accelerator
+# num_devices = torch.cuda.device_count()
+# if num_devices==0: num_devices = 1
+# accelerator = Accelerator(split_batches=False)
+# global_batch_size = 128
+                          
+### Multi-GPU config ###
+from accelerate import Accelerator, DeepSpeedPlugin
+num_devices = torch.cuda.device_count()
+if num_devices==0: num_devices = 1
+if num_devices <= 1 and utils.is_interactive():
+    # can emulate a distributed environment for deepspeed to work in jupyter notebook
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = str(np.random.randint(10000)+9000)
+    os.environ["RANK"] = "0"
+    os.environ["LOCAL_RANK"] = "0"
+    os.environ["WORLD_SIZE"] = "1"
+    os.environ["GLOBAL_BATCH_SIZE"] = "128" # set this to your batch size!
+    global_batch_size = os.environ["GLOBAL_BATCH_SIZE"]
+
+# alter the deepspeed config according to your global and local batch size
+if local_rank == 0:
+    with open('deepspeed_config_stage2.json', 'r') as file:
+        config = json.load(file)
+    config['train_batch_size'] = int(os.environ["GLOBAL_BATCH_SIZE"])
+    config['train_micro_batch_size_per_gpu'] = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices
+    with open('deepspeed_config_stage2.json', 'w') as file:
+        json.dump(config, file)
+else:
+    # give some time for the local_rank=0 gpu to prep new deepspeed config file
+    time.sleep(10)
+deepspeed_plugin = DeepSpeedPlugin("deepspeed_config_stage2.json")
+accelerator = Accelerator(split_batches=False, deepspeed_plugin=deepspeed_plugin)
+
+
+# In[ ]:
+
+
+print("PID of this process =",os.getpid())
+device = accelerator.device
+print("device:",device)
+num_workers = num_devices
+print(accelerator.state)
+world_size = accelerator.state.num_processes
+distributed = not accelerator.state.distributed_type == 'NO'
+print("distributed =",distributed, "num_devices =", num_devices, "local rank =", local_rank, "world size =", world_size)
+print = accelerator.print # only print if local_rank=0
+
+
+# # Configurations
+
+# In[3]:
+
+
+# if running this interactively, can specify jupyter_args here for argparser to use
+if utils.is_interactive():
+    # Example use
+    jupyter_args = f"--data_path=/fsx/proj-fmri/shared/mindeyev2_dataset \
+                    --model_name=test \
+                    --subj=1 --batch_size={global_batch_size} --n_samples_save=0 \
+                    --max_lr=3e-4 --mixup_pct=.66 --num_epochs=12 --ckpt_interval=999 --no-use_image_aug"
+
+    jupyter_args = jupyter_args.split()
+    print(jupyter_args)
+    
+    from IPython.display import clear_output # function to clear print outputs in cell
+    get_ipython().run_line_magic('load_ext', 'autoreload')
+    # this allows you to change functions in models.py or utils.py and have this notebook automatically update with your revisions
+    get_ipython().run_line_magic('autoreload', '2')
+
+
+# In[4]:
+
+
+parser = argparse.ArgumentParser(description="Model Training Configuration")
+parser.add_argument(
+    "--model_name", type=str, default="testing",
+    help="name of model, used for ckpt saving and wandb logging (if enabled)",
+)
+parser.add_argument(
+    "--data_path", type=str, default="/fsx/proj-fmri/shared/natural-scenes-dataset",
+    help="Path to where NSD data is stored / where to download it to",
+)
+parser.add_argument(
+    "--subj",type=int, default=1, choices=[1,2,5,7],
+)
+parser.add_argument(
+    "--batch_size", type=int, default=32,
+    help="Batch size can be increased by 10x if only training v2c and not diffusion prior",
+)
+parser.add_argument(
+    "--wandb_log",action=argparse.BooleanOptionalAction,default=False,
+    help="whether to log to wandb",
+)
+parser.add_argument(
+    "--resume_from_ckpt",action=argparse.BooleanOptionalAction,default=False,
+    help="if not using wandb and want to resume from a ckpt",
+)
+parser.add_argument(
+    "--wandb_project",type=str,default="stability",
+    help="wandb project name",
+)
+parser.add_argument(
+    "--mixup_pct",type=float,default=.33,
+    help="proportion of way through training when to switch from BiMixCo to SoftCLIP",
+)
+parser.add_argument(
+    "--use_image_aug",action=argparse.BooleanOptionalAction,default=True,
+    help="whether to use image augmentation",
+)
+parser.add_argument(
+    "--num_epochs",type=int,default=240,
+    help="number of epochs of training",
+)
+parser.add_argument(
+    "--lr_scheduler_type",type=str,default='cycle',choices=['cycle','linear'],
+)
+parser.add_argument(
+    "--ckpt_saving",action=argparse.BooleanOptionalAction,default=True,
+)
+parser.add_argument(
+    "--ckpt_interval",type=int,default=5,
+    help="save backup ckpt and reconstruct every x epochs",
+)
+parser.add_argument(
+    "--seed",type=int,default=42,
+)
+parser.add_argument(
+    "--max_lr",type=float,default=3e-4,
+)
+parser.add_argument(
+    "--n_samples_save",type=int,default=0,choices=[0,1],
+    help="Number of reconstructions for monitoring progress, 0 will speed up training",
+)
+
+if utils.is_interactive():
+    args = parser.parse_args(jupyter_args)
+else:
+    args = parser.parse_args()
+
+# create global variables without the args prefix
+for attribute_name in vars(args).keys():
+    globals()[attribute_name] = getattr(args, attribute_name)
+
+print("global batch_size", batch_size)
+batch_size = int(batch_size / num_devices)
+print("batch_size", batch_size)
+
+
+# In[5]:
+
+
+outdir = os.path.abspath(f'../train_logs/{model_name}')
+if not os.path.exists(outdir):
+    os.makedirs(outdir,exist_ok=True)
+if use_image_aug:
+    import kornia
+    from kornia.augmentation.container import AugmentationSequential
+    img_augment = AugmentationSequential(
+        kornia.augmentation.RandomResizedCrop((224,224), (0.6,1), p=0.3),
+        kornia.augmentation.Resize((224, 224)),
+        kornia.augmentation.RandomHorizontalFlip(p=0.3),
+        kornia.augmentation.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1, p=0.3),
+        kornia.augmentation.RandomGrayscale(p=0.3),
+        same_on_batch=False,
+        data_keys=["input"],
+    )
+
+
+# # Prep data, models, and dataloaders
+
+# ## Dataloader
+
+# In[6]:
+
+
+if subj==1:
+    num_train = 24958
+    num_test = 2770
+test_batch_size = num_test
+
+def my_split_by_node(urls): return urls
+    
+train_url = f"{data_path}/wds/subj0{subj}/train/" + "{0..36}.tar"
+print(train_url)
+
+train_data = wds.WebDataset(train_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+test_url = f"{data_path}/wds/subj0{subj}/test/" + "0.tar"
+print(test_url)
+
+test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\
+                    .shuffle(750, initial=1500, rng=random.Random(42))\
+                    .decode("torch")\
+                    .rename(behav="behav.npy", past_behav="past_behav.npy", future_behav="future_behav.npy", olds_behav="olds_behav.npy")\
+                    .to_tuple(*["behav", "past_behav", "future_behav", "olds_behav"])
+test_dl = torch.utils.data.DataLoader(test_data, batch_size=test_batch_size, shuffle=False, drop_last=False, pin_memory=True)
+
+
+# ### check dataloaders are working
+
+# In[7]:
+
+
+# test_indices = []
+# test_images = []
+# for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+#     test_indices = np.append(test_indices, behav[:,0,5].numpy())
+#     test_images = np.append(test_images, behav[:,0,0].numpy())
+# test_indices = test_indices.astype(np.int16)
+# print(test_i, (test_i+1) * test_batch_size, len(test_indices))
+# print("---\n")
+
+# train_indices = []
+# train_images = []
+# for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+#     train_indices = np.append(train_indices, behav[:,0,5].long().numpy())
+#     train_images = np.append(train_images, behav[:,0,0].numpy())
+# train_indices = train_indices.astype(np.int16)
+# print(train_i, (train_i+1) * batch_size, len(train_indices))
+
+
+# ## Load voxel betas, K-means clustering model, and images
+
+# In[8]:
+
+
+# load betas
+f = h5py.File(f'{data_path}/betas_all_subj0{subj}.hdf5', 'r')
+voxels = f['betas'][:]
+print(f"subj0{subj} betas loaded into memory")
+voxels = torch.Tensor(voxels).to("cpu").half()
+if subj==1:
+    voxels = torch.hstack((voxels, torch.zeros((len(voxels), 5))))
+print("voxels", voxels.shape)
+num_voxels = voxels.shape[-1]
+
+# load orig images
+f = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
+images = f['images'][:]
+images = torch.Tensor(images).to("cpu").half()
+print("images", images.shape)
+
+
+# In[9]:
+
+
+from models import Clipper
+eva02_model = Clipper("ViT-L/14", device=torch.device(f"cuda:{local_rank}"), hidden_state=True, norm_embs=True)
+
+clip_seq_dim = 257
+clip_emb_dim = 768
+hidden_dim = 4096
+
+
+# In[10]:
+
+
+class MindEyeModule(nn.Module):
+    def __init__(self):
+        super(MindEyeModule, self).__init__()
+    def forward(self, x):
+        return x
+        
+model = MindEyeModule()
+model
+
+
+# In[11]:
+
+
+class RidgeRegression(torch.nn.Module):
+    # make sure to add weight_decay when initializing optimizer
+    def __init__(self, input_size, out_features): 
+        super(RidgeRegression, self).__init__()
+        self.linear = torch.nn.Linear(input_size, out_features)
+    def forward(self, x):
+        return self.linear(x)
+        
+model.ridge = RidgeRegression(voxels.shape[1], out_features=hidden_dim)
+utils.count_params(model.ridge)
+utils.count_params(model)
+
+b = torch.randn((2,voxels.shape[1]))
+print(b.shape, model.ridge(b).shape)
+
+
+# In[12]:
+
+
+from functools import partial
+class BrainNetwork(nn.Module):
+    def __init__(self, out_dim=768, in_dim=15724, clip_size=768, h=4096, n_blocks=4, norm_type='ln', act_first=False, use_projector=True, drop1=.5, drop2=.15):
+        super().__init__()
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop2)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, out_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.clip_size = clip_size
+        self.use_projector = use_projector
+        if use_projector:
+            self.projector = nn.Sequential(
+                nn.LayerNorm(clip_size),
+                nn.GELU(),
+                nn.Linear(clip_size, 2048),
+                nn.LayerNorm(2048),
+                nn.GELU(),
+                nn.Linear(2048, 2048),
+                nn.LayerNorm(2048),
+                nn.GELU(),
+                nn.Linear(2048, clip_size)
+            )
+        
+    def forward(self, x):
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        if self.use_projector:
+            return self.projector(x.reshape(len(x), -1, self.clip_size))
+        return x
+
+model.backbone = BrainNetwork(in_dim=hidden_dim, clip_size=clip_emb_dim, out_dim=clip_seq_dim*clip_emb_dim, use_projector=True)
+utils.count_params(model.backbone)
+utils.count_params(model)
+
+b = torch.randn((2,hidden_dim))
+print(b.shape, model.backbone(b).shape)
+
+
+# In[13]:
+
+
+no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+opt_grouped_parameters = [
+    {'params': [p for n, p in model.ridge.named_parameters()], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 1e-2},
+    {'params': [p for n, p in model.backbone.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
+]
+
+optimizer = torch.optim.AdamW(opt_grouped_parameters, lr=max_lr)
+
+if lr_scheduler_type == 'linear':
+    lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+        optimizer,
+        total_iters=int(num_epochs*(num_train*num_devices//batch_size)),
+        last_epoch=-1
+    )
+elif lr_scheduler_type == 'cycle':
+    total_steps=int(num_epochs*(num_train*num_devices//batch_size))
+    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer, 
+        max_lr=max_lr,
+        total_steps=total_steps,
+        final_div_factor=1000,
+        last_epoch=-1, pct_start=2/num_epochs
+    )
+    
+def save_ckpt(tag):    
+    ckpt_path = outdir+f'/{tag}.pth'
+    print(f'saving {ckpt_path}',flush=True)
+    unwrapped_model = accelerator.unwrap_model(model)
+    try:
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': unwrapped_model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'lr_scheduler': lr_scheduler.state_dict(),
+            'train_losses': losses,
+            'test_losses': test_losses,
+            'lrs': lrs,
+            }, ckpt_path)
+    except:
+        print("Couldn't save... moving on to prevent crashing.")
+    del unwrapped_model
+        
+print("\nDone with model preparations!")
+
+
+# # Weights and Biases
+
+# In[14]:
+
+
+# params for wandb
+if local_rank==0 and wandb_log: # only use main process for wandb logging
+    import wandb
+    
+    wandb_project = 'stability'
+    wandb_run = model_name
+    wandb_notes = ''
+    
+    print(f"wandb {wandb_project} run {wandb_run}")
+    wandb.login(host='https://stability.wandb.io')#, relogin=True)
+    wandb_config = {
+      "model_name": model_name,
+      "clip_variant": clip_variant,
+      "batch_size": batch_size,
+      "num_epochs": num_epochs,
+      "use_image_aug": use_image_aug,
+      "max_lr": max_lr,
+      "lr_scheduler_type": lr_scheduler_type,
+      "mixup_pct": mixup_pct,
+      "num_train": num_train,
+      "num_test": num_test,
+      "seed": seed,
+      "distributed": distributed,
+      "num_devices": num_devices,
+      "world_size": world_size,
+    }
+    print("wandb_config:\n",wandb_config)
+    if True: # wandb_auto_resume
+        print("wandb_id:",model_name)
+        wandb.init(
+            id = model_name,
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+            resume="allow",
+        )
+    else:
+        wandb.init(
+            project=wandb_project,
+            name=wandb_run,
+            config=wandb_config,
+            notes=wandb_notes,
+        )
+else:
+    wandb_log = False
+
+
+# # Main
+
+# In[15]:
+
+
+epoch = 0
+losses, test_losses, lrs = [], [], []
+best_test_loss = 1e9
+soft_loss_temps = utils.cosine_anneal(0.004, 0.0075, num_epochs - int(mixup_pct * num_epochs))
+
+# Optionally resume from checkpoint #
+if resume_from_ckpt:
+    print("\n---resuming from last.pth ckpt---\n")
+    try:
+        checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+    except:
+        print('last.pth failed... trying last_backup.pth')
+        checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+    epoch = checkpoint['epoch']
+    print("Epoch",epoch)
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    diffusion_prior.load_state_dict(checkpoint['model_state_dict'])
+    del checkpoint
+elif wandb_log:
+    if wandb.run.resumed:
+        print("\n---resuming from last.pth ckpt---\n")
+        try:
+            checkpoint = torch.load(outdir+'/last.pth', map_location='cpu')
+        except:
+            print('last.pth failed... trying last_backup.pth')
+            checkpoint = torch.load(outdir+'/last_backup.pth', map_location='cpu')
+        epoch = checkpoint['epoch']
+        print("Epoch",epoch)
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+        diffusion_prior.load_state_dict(checkpoint['model_state_dict'])
+        del checkpoint
+torch.cuda.empty_cache()
+
+
+# In[16]:
+
+
+model, optimizer, train_dl, test_dl, lr_scheduler = accelerator.prepare(
+model, optimizer, train_dl, test_dl, lr_scheduler
+)
+
+
+# In[17]:
+
+
+print(f"{model_name} starting with epoch {epoch} / {num_epochs}")
+progress_bar = tqdm(range(epoch,num_epochs), ncols=1200, disable=(local_rank!=0))
+test_image, test_voxel = None, None
+mse = nn.MSELoss()
+for epoch in progress_bar:
+    model.train()
+    
+    fwd_percent_correct = 0.
+    bwd_percent_correct = 0.
+    test_fwd_percent_correct = 0.
+    test_bwd_percent_correct = 0.
+
+    for train_i, (behav, past_behav, future_behav, old_behav) in enumerate(train_dl):
+        with torch.cuda.amp.autocast():
+            optimizer.zero_grad()
+
+            voxel = voxels[behav[:,0,5].cpu().long()].to(device)
+            image = images[behav[:,0,0].cpu().long()].to(device)
+            
+            if use_image_aug: image = img_augment(image)
+            
+            clip_target = eva02_model.embed_image(image.float())
+            assert not torch.any(torch.isnan(clip_target))
+
+            if epoch < int(mixup_pct * num_epochs):
+                voxel, perm, betas, select = utils.mixco(voxel)
+
+            voxel_ridge = model.ridge(voxel)
+            
+            clip_voxels = model.backbone(voxel_ridge)
+            
+            clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+            clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+
+            if epoch < int(mixup_pct * num_epochs):                
+                loss_clip = utils.mixco_nce(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006, 
+                    perm=perm, betas=betas, select=select)
+            else:
+                epoch_temp = soft_loss_temps[epoch-int(mixup_pct*num_epochs)]
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=epoch_temp)
+                
+            loss = loss_clip
+            
+            utils.check_loss(loss)
+
+            accelerator.backward(loss)
+            optimizer.step()
+    
+            losses.append(loss.item())
+            lrs.append(optimizer.param_groups[0]['lr'])
+    
+            # forward and backward top 1 accuracy        
+            labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+            fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)
+            bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)
+
+            if lr_scheduler_type is not None:
+                lr_scheduler.step()
+
+    model.eval()
+    for test_i, (behav, past_behav, future_behav, old_behav) in enumerate(test_dl):
+        with torch.no_grad():
+            with torch.cuda.amp.autocast():    
+                # all test samples should be loaded per batch such that test_i should never exceed 0
+                if len(behav) != num_test: print("!",len(behav),num_test)
+                
+                ## Average same-image repeats ##
+                if test_image is None:
+                    voxel = voxels[behav[:,0,5].cpu().long()]
+                    image = behav[:,0,0].cpu().long()
+                    
+                    unique_image, sort_indices = torch.unique(image, return_inverse=True)
+                    for im in unique_image:
+                        locs = torch.where(im == image)[0]
+                        if test_image is None:
+                            test_image = images[im][None]
+                            test_voxel = torch.mean(voxel[locs],axis=0)[None]
+                        else:
+                            test_image = torch.vstack((test_image, images[im][None]))
+                            test_voxel = torch.vstack((test_voxel, torch.mean(voxel[locs],axis=0)[None]))
+    
+                # random sample of 300
+                random_indices = torch.randperm(len(test_voxel))[:300]
+                voxel = test_voxel[random_indices].to(device)
+                image = test_image[random_indices].to(device)
+                assert len(image) == 300
+        
+                clip_target = eva02_model.embed_image(image.float())
+                
+                voxel_ridge = model.ridge(voxel)
+                
+                clip_voxels = model.backbone(voxel_ridge)
+                
+                clip_voxels_norm = nn.functional.normalize(clip_voxels.flatten(1), dim=-1)
+                clip_target_norm = nn.functional.normalize(clip_target.flatten(1), dim=-1)
+        
+                loss_clip = utils.soft_clip_loss(
+                    clip_voxels_norm,
+                    clip_target_norm,
+                    temp=.006)
+                    
+                loss = loss_clip
+                
+                utils.check_loss(loss)
+        
+                test_losses.append(loss.item())
+        
+                # forward and backward top 1 accuracy        
+                labels = torch.arange(len(clip_target_norm)).to(clip_voxels_norm.device) 
+                test_fwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_voxels_norm, clip_target_norm), labels, k=1)
+                test_bwd_percent_correct += utils.topk(utils.batchwise_cosine_similarity(clip_target_norm, clip_voxels_norm), labels, k=1)
+
+    if local_rank==0:      
+        if utils.is_interactive():
+            # clear_output(wait=True)
+            print("---")
+
+        assert (test_i+1) == 1
+        logs = {"train/loss": np.mean(losses[-(train_i+1):]),
+            "test/loss": np.mean(test_losses[-(test_i+1):]),
+            "train/lr": lrs[-1],
+            "train/num_steps": len(losses),
+            "test/num_steps": len(test_losses),
+            "train/fwd_pct_correct": fwd_percent_correct.item() / (train_i + 1),
+            "train/bwd_pct_correct": bwd_percent_correct.item() / (train_i + 1),
+            "test/test_fwd_pct_correct": test_fwd_percent_correct.item() / (test_i + 1),
+            "test/test_bwd_pct_correct": test_bwd_percent_correct.item() / (test_i + 1),
+            }
+        progress_bar.set_postfix(**logs)
+
+        # Save model checkpoint and reconstruct
+        if epoch % ckpt_interval == 0:
+            if not utils.is_interactive():
+                save_ckpt(f'last')
+                
+        if wandb_log: wandb.log(logs)
+
+    # wait for other GPUs to catch up if needed
+    accelerator.wait_for_everyone()
+    torch.cuda.empty_cache()
+    gc.collect()
+
+print("\n===Finished!===\n")
+if ckpt_saving:
+    save_ckpt(f'last')
+if not utils.is_interactive():
+    sys.exit(0)
+

src/accel.slurm

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

src/checking_models.ipynb

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

src/cnd_prov/cnd_prov-Copy1.py

@@ -0,0 +1,148 @@
+#new one
+import pandas as pd
+import datasets
+import os
+import pickle
+import numpy as np
+from PIL import Image
+import matplotlib
+
+_VERSION = datasets.Version("0.0.3")
+
+_DESCRIPTION = "TODO"
+_HOMEPAGE = "TODO"
+_LICENSE = "TODO"
+_CITATION = "TODO"
+
+_FEATURES = datasets.Features(
+    {
+        "target": datasets.Image(),
+        "source": datasets.Image(),
+	    "heatmap": datasets.Image(),
+        "depth": datasets.Image(),
+        "prompt": datasets.Value("string"),
+    },
+)
+
+METADATA_DIR = "/fsx/proj-fmri/ckadirt/MindEyeV2/src/cnd_prov/data.pkl"
+SOURCE_DIR = "/fsx/proj-fmri/shared/controlNetData/source"
+TARGET_DIR = "/fsx/proj-fmri/shared/controlNetData/target"
+HEATMAP_DIR = "/fsx/proj-fmri/shared/controlNetData/seg"
+DEPTH_DIR = "/fsx/proj-fmri/shared/dinov2_depth"
+
+# METADATA_URL = hf_hub_url(
+#     "fusing/fill50k",
+#     filename="train.jsonl",
+#     repo_type="dataset",
+# )
+
+# IMAGES_URL = hf_hub_url(
+#     "fusing/fill50k",
+#     filename="images.zip",
+#     repo_type="dataset",
+# )
+
+# CONDITIONING_IMAGES_URL = hf_hub_url(
+#     "fusing/fill50k",
+#     filename="conditioning_images.zip",
+#     repo_type="dataset",
+# )
+
+_DEFAULT_CONFIG = datasets.BuilderConfig(name="default", version=_VERSION)
+
+
+class CocoTest(datasets.GeneratorBasedBuilder):
+    BUILDER_CONFIGS = [_DEFAULT_CONFIG]
+    DEFAULT_CONFIG_NAME = "default"
+
+    def _info(self):
+        return datasets.DatasetInfo(
+            description=_DESCRIPTION,
+            features=_FEATURES,
+            supervised_keys=None,
+            homepage=_HOMEPAGE,
+            license=_LICENSE,
+            citation=_CITATION,
+        )
+
+    def _split_generators(self, dl_manager):
+        metadata_path = METADATA_DIR
+        target_dir = TARGET_DIR
+        source_dir = SOURCE_DIR
+        heatmap_dir = HEATMAP_DIR
+        depth_dir = DEPTH_DIR
+
+        return [
+            datasets.SplitGenerator(
+                name=datasets.Split.TRAIN,
+                # These kwargs will be passed to _generate_examples
+                gen_kwargs={
+                    "metadata_path": metadata_path,
+                    "target_dir": TARGET_DIR,
+                    "source_dir": SOURCE_DIR,
+                    "heatmap_dir": HEATMAP_DIR,
+                    "depth_dir": DEPTH_DIR,
+                    "num_examples": 120, 
+                },
+            ),
+            datasets.SplitGenerator(
+                name=datasets.Split.VALIDATION,
+                # These kwargs will be passed to _generate_examples
+                gen_kwargs={
+                    "metadata_path": metadata_path,
+                    "target_dir": TARGET_DIR,
+                    "source_dir": SOURCE_DIR,
+                    "heatmap_dir": HEATMAP_DIR,
+                    "depth_dir": DEPTH_DIR,
+                    "num_examples": 120, 
+                },
+            ),
+        ]
+
+
+
+    def _generate_examples(self, metadata_path, target_dir, source_dir, heatmap_dir, depth_dir, num_examples):
+        data = []
+        with open(metadata_path, 'rb') as f:
+            loaded_data = pickle.load(f)
+            for line in loaded_data[:num_examples]:
+                data.append(line)
+
+        for _, item in enumerate(data):
+            source_filename = item['source']
+            target_filename = item['target']
+            heatmap_filename = item['h_map']
+            depth_filename = item['depth']
+            prompt = item['prompt']
+            
+            
+
+            tgt_img = open(target_filename, "rb").read()
+            src_img = open(source_filename, "rb").read()
+            h_img = open(heatmap_filename, "rb").read()
+            
+
+            # one channel depth image
+            d_img = Image.open(depth_filename).convert('1')
+
+
+
+            yield item["target"], {
+                "prompt": prompt,
+                "target": {
+                    "path": target_filename,
+                    "bytes": tgt_img,
+                },
+                "source": {
+                    "path": source_filename,
+                    "bytes": src_img,
+                },
+                "heatmap": {
+                    "path": heatmap_filename,
+                    "bytes": h_img,
+                },
+                "depth": {
+                    "path": depth_filename,
+                    "bytes": d_img,
+                },
+            }

src/cnd_prov/cnd_prov.py

@@ -0,0 +1,138 @@
+#new one
+import pandas as pd
+import datasets
+import os
+import pickle
+
+_VERSION = datasets.Version("0.0.3")
+
+_DESCRIPTION = "TODO"
+_HOMEPAGE = "TODO"
+_LICENSE = "TODO"
+_CITATION = "TODO"
+
+_FEATURES = datasets.Features(
+    {
+        "target": datasets.Image(),
+        "source": datasets.Image(),
+	    "heatmap": datasets.Image(),
+        "depth": datasets.Image(),
+        "prompt": datasets.Value("string"),
+    },
+)
+
+METADATA_DIR = "/fsx/proj-fmri/ckadirt/MindEyeV2/src/cnd_prov/data.pkl"
+SOURCE_DIR = "/fsx/proj-fmri/shared/controlNetData/source"
+TARGET_DIR = "/fsx/proj-fmri/shared/controlNetData/target"
+HEATMAP_DIR = "/fsx/proj-fmri/shared/controlNetData/seg"
+DEPTH_DIR = "/fsx/proj-fmri/shared/dinov2_depth"
+
+# METADATA_URL = hf_hub_url(
+#     "fusing/fill50k",
+#     filename="train.jsonl",
+#     repo_type="dataset",
+# )
+
+# IMAGES_URL = hf_hub_url(
+#     "fusing/fill50k",
+#     filename="images.zip",
+#     repo_type="dataset",
+# )
+
+# CONDITIONING_IMAGES_URL = hf_hub_url(
+#     "fusing/fill50k",
+#     filename="conditioning_images.zip",
+#     repo_type="dataset",
+# )
+
+_DEFAULT_CONFIG = datasets.BuilderConfig(name="default", version=_VERSION)
+
+
+class CocoTest(datasets.GeneratorBasedBuilder):
+    BUILDER_CONFIGS = [_DEFAULT_CONFIG]
+    DEFAULT_CONFIG_NAME = "default"
+
+    def _info(self):
+        return datasets.DatasetInfo(
+            description=_DESCRIPTION,
+            features=_FEATURES,
+            supervised_keys=None,
+            homepage=_HOMEPAGE,
+            license=_LICENSE,
+            citation=_CITATION,
+        )
+
+    def _split_generators(self, dl_manager):
+        metadata_path = METADATA_DIR
+        target_dir = TARGET_DIR
+        source_dir = SOURCE_DIR
+        heatmap_dir = HEATMAP_DIR
+        depth_dir = DEPTH_DIR
+
+        return [
+            datasets.SplitGenerator(
+                name=datasets.Split.TRAIN,
+                # These kwargs will be passed to _generate_examples
+                gen_kwargs={
+                    "metadata_path": metadata_path,
+                    "target_dir": TARGET_DIR,
+                    "source_dir": SOURCE_DIR,
+                    "heatmap_dir": HEATMAP_DIR,
+                    "depth_dir": DEPTH_DIR,
+                    "num_examples": 190573, 
+                },
+            ),
+            datasets.SplitGenerator(
+                name=datasets.Split.VALIDATION,
+                # These kwargs will be passed to _generate_examples
+                gen_kwargs={
+                    "metadata_path": metadata_path,
+                    "target_dir": TARGET_DIR,
+                    "source_dir": SOURCE_DIR,
+                    "heatmap_dir": HEATMAP_DIR,
+                    "depth_dir": DEPTH_DIR,
+                    "num_examples": 20000, 
+                },
+            ),
+        ]
+
+    def _generate_examples(self, metadata_path, target_dir, source_dir, heatmap_dir, depth_dir, num_examples):
+        data = []
+        with open(metadata_path, 'rb') as f:
+            loaded_data = pickle.load(f)
+            for line in loaded_data[:num_examples]:
+                data.append(line)
+
+        for _, item in enumerate(data):
+            source_filename = item['source']
+            target_filename = item['target']
+            heatmap_filename = item['h_map']
+            depth_filename = item['depth']
+            prompt = item['prompt']
+            
+            
+
+            tgt_img = open(target_filename, "rb").read()
+            src_img = open(source_filename, "rb").read()
+            h_img = open(heatmap_filename, "rb").read()
+            d_img = open(depth_filename, "rb").read()
+
+            yield item["target"], {
+                "prompt": prompt,
+                "target": {
+                    "path": target_filename,
+                    "bytes": tgt_img,
+                },
+                "source": {
+                    "path": source_filename,
+                    "bytes": src_img,
+                },
+                "heatmap": {
+                    "path": heatmap_filename,
+                    "bytes": h_img,
+                },
+                "depth": {
+                    "path": depth_filename,
+                    "bytes": d_img,
+                },
+            }

src/cnd_prov/data.pkl

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

src/deepspeed_config_stage1.json

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

src/deepspeed_config_stage2.json

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

src/deepspeed_config_stage3.json

@@ -0,0 +1,44 @@
+{
+    "bf16": {
+        "enabled": false
+    },
+    "fp16": {
+        "enabled": true
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "contiguous_gradients": true,
+        "stage3_gather_16bit_weights_on_model_save": true,
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_prefetch_bucket_size": 1e7,
+        "stage3_param_persistence_threshold": 1e5,
+        "reduce_bucket_size": 1e7,
+        "sub_group_size": 1e9,
+        "offload_optimizer": {
+            "device": "none",
+            "nvme_path": "/scratch",
+            "pin_memory": true
+         },
+        "offload_param": {
+            "device": "none",
+            "nvme_path": "/scratch",
+            "buffer_size": 4e9,
+            "pin_memory": true
+       }
+    },
+    "aio": {
+        "block_size": 26214400,
+        "queue_depth": 32,
+        "thread_count": 1,
+        "single_submit": false,
+        "overlap_events": true
+    },
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 1.0,
+    "steps_per_print": 20000,
+    "train_batch_size": 128,
+    "train_micro_batch_size_per_gpu": 16,
+    "wall_clock_breakdown": false,
+    "zero_allow_untested_optimizer": true
+}

src/huggingface_to_s3.ipynb

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

src/models.py

@@ -0,0 +1,1344 @@
+import os
+import numpy as np
+from torchvision import transforms
+import torch
+import torch.nn as nn
+import PIL
+import clip
+import open_clip
+from functools import partial
+
+# for prior
+from dalle2_pytorch import DiffusionPrior
+from dalle2_pytorch.dalle2_pytorch import l2norm, default, exists
+from tqdm.auto import tqdm
+import random
+import json
+from dalle2_pytorch.train_configs import DiffusionPriorNetworkConfig
+# vd prior
+from dalle2_pytorch.dalle2_pytorch import RotaryEmbedding, CausalTransformer, SinusoidalPosEmb, MLP, Rearrange, repeat, rearrange, prob_mask_like, LayerNorm, RelPosBias, FeedForward, Attention
+
+# for pipeline
+from diffusers import StableDiffusionImageVariationPipeline, VersatileDiffusionDualGuidedPipeline
+from typing import Callable, List, Optional, Union
+
+from diffusers.models.vae import Decoder
+
+# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+# class BrainMLP(nn.Module):
+#     def __init__(self, out_dim=257*768, in_dim=15724, clip_size=768, h=4096):
+#         super().__init__()
+#         self.lin0 = nn.Sequential(
+#             nn.Linear(in_dim, h, bias=False),
+#             nn.LayerNorm(h),
+#             nn.GELU(inplace=True),
+#             nn.Dropout(0.5))
+#         self.mlp = nn.ModuleList([
+#             nn.Sequential(
+#                 nn.Linear(h, h),
+#                 nn.LayerNorm(h),
+#                 nn.GELU(inplace=True),
+#                 nn.Dropout(0.15)
+#             ) for _ in range(4)])
+#         self.lin1 = nn.Linear(h, out_dim, bias=True)
+#         self.proj = nn.Sequential(
+#             nn.LayerNorm(clip_size),
+#             nn.GELU(),
+#             nn.Linear(clip_size, 2048),
+#             nn.LayerNorm(2048),
+#             nn.GELU(),
+#             nn.Linear(2048, 2048),
+#             nn.LayerNorm(2048),
+#             nn.GELU(),
+#             nn.Linear(2048, clip_size))
+#     def forward(self, x):
+#         x = self.lin0(x)
+#         residual = x
+#         for res_block in range(self.n_blocks):
+#             x = self.mlp[res_block](x)
+#             x += residual
+#             residual = x
+#         diffusion_prior_input = self.lin1(x.reshape(len(x), -1))
+#         disjointed_clip_fmri = self.proj(diffusion_prior_input.reshape(
+#                                         len(x),-1, self.clip_size))
+#         return diffusion_prior_input, disjointed_clip_fmri
+
+
+
+class Clipper(torch.nn.Module):
+    def __init__(self, clip_variant, clamp_embs=False, norm_embs=False,
+                 hidden_state=False, device=torch.device('cpu')):
+        super().__init__()
+        assert clip_variant in ("RN50", "ViT-L/14", "ViT-B/32", "RN50x64"), \
+            "clip_variant must be one of RN50, ViT-L/14, ViT-B/32, RN50x64"
+        print(clip_variant, device)
+        
+        if clip_variant=="ViT-L/14" and hidden_state:
+            # from transformers import CLIPVisionModelWithProjection
+            # image_encoder = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-large-patch14",cache_dir="/fsx/proj-medarc/fmri/cache")
+            from transformers import CLIPVisionModelWithProjection
+            sd_cache_dir = '/fsx/proj-fmri/shared/cache/models--shi-labs--versatile-diffusion/snapshots/2926f8e11ea526b562cd592b099fcf9c2985d0b7'
+            image_encoder = CLIPVisionModelWithProjection.from_pretrained(sd_cache_dir, subfolder='image_encoder').eval()
+            image_encoder = image_encoder.to(device)
+            for param in image_encoder.parameters():
+                param.requires_grad = False # dont need to calculate gradients
+            self.image_encoder = image_encoder
+        elif hidden_state:
+            raise Exception("hidden_state embeddings only works with ViT-L/14 right now")
+        
+        clip_model, preprocess = clip.load(clip_variant, device=device)
+        clip_model.eval() # dont want to train model
+        for param in clip_model.parameters():
+            param.requires_grad = False # dont need to calculate gradients
+            
+        self.clip = clip_model
+        self.clip_variant = clip_variant
+        if clip_variant == "RN50x64":
+            self.clip_size = (448,448)
+        else:
+            self.clip_size = (224,224)
+            
+        preproc = transforms.Compose([
+            transforms.Resize(size=self.clip_size[0], interpolation=transforms.InterpolationMode.BICUBIC),
+            transforms.CenterCrop(size=self.clip_size),
+            transforms.Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711))
+        ])
+        self.preprocess = preproc
+        self.hidden_state = hidden_state
+        self.mean = np.array([0.48145466, 0.4578275, 0.40821073])
+        self.std = np.array([0.26862954, 0.26130258, 0.27577711])
+        self.normalize = transforms.Normalize(self.mean, self.std)
+        self.denormalize = transforms.Normalize((-self.mean / self.std).tolist(), (1.0 / self.std).tolist())
+        self.clamp_embs = clamp_embs
+        self.norm_embs = norm_embs
+        self.device= device
+        
+        def versatile_normalize_embeddings(encoder_output):
+            embeds = encoder_output.last_hidden_state
+            embeds = image_encoder.vision_model.post_layernorm(embeds)
+            embeds = image_encoder.visual_projection(embeds)
+            return embeds
+        self.versatile_normalize_embeddings = versatile_normalize_embeddings
+
+    def resize_image(self, image):
+        # note: antialias should be False if planning to use Pinkney's Image Variation SD model
+        return transforms.Resize(self.clip_size)(image.to(self.device))
+
+    def embed_image(self, image):
+        """Expects images in -1 to 1 range"""
+        if self.hidden_state:
+            # clip_emb = self.preprocess((image/1.5+.25).to(self.device)) # for some reason the /1.5+.25 prevents oversaturation
+            clip_emb = self.preprocess((image).to(self.device))
+            clip_emb = self.image_encoder(clip_emb)
+            clip_emb = self.versatile_normalize_embeddings(clip_emb)
+        else:
+            clip_emb = self.preprocess(image.to(self.device))
+            clip_emb = self.clip.encode_image(clip_emb)
+        # input is now in CLIP space, but mind-reader preprint further processes embeddings:
+        if self.clamp_embs:
+            clip_emb = torch.clamp(clip_emb, -1.5, 1.5)
+        if self.norm_embs:
+            if self.hidden_state:        
+                # normalize all tokens by cls token's norm
+                clip_emb = clip_emb / torch.norm(clip_emb[:, 0], dim=-1).reshape(-1, 1, 1)
+            else:
+                clip_emb = nn.functional.normalize(clip_emb, dim=-1)
+        return clip_emb
+
+    def embed_text(self, text_samples):
+        clip_text = clip.tokenize(text_samples).to(self.device)
+        clip_text = self.clip.encode_text(clip_text)
+        if self.clamp_embs:
+            clip_text = torch.clamp(clip_text, -1.5, 1.5)
+        if self.norm_embs:
+            clip_text = nn.functional.normalize(clip_text, dim=-1)
+        return clip_text
+
+    def embed_curated_annotations(self, annots):
+        for i,b in enumerate(annots):
+            t = ''
+            while t == '':
+                rand = torch.randint(5,(1,1))[0][0]
+                t = b[0,rand]
+            if i==0:
+                txt = np.array(t)
+            else:
+                txt = np.vstack((txt,t))
+        txt = txt.flatten()
+        return self.embed_text(txt)
+    
+class OpenClipper(torch.nn.Module):
+    def __init__(self, clip_variant, norm_embs=False, device=torch.device('cpu')):
+        super().__init__()
+        print(clip_variant, device)
+        assert clip_variant == 'ViT-H-14' # not setup for other models yet
+                
+        clip_model, _, preprocess = open_clip.create_model_and_transforms('ViT-H-14', 
+                                        pretrained='laion2b_s32b_b79k', device=device)
+        clip_model.eval() # dont want to train model
+        for param in clip_model.parameters():
+            param.requires_grad = False # dont need to calculate gradients
+            
+        # overwrite preprocess to accept torch inputs instead of PIL Image
+        preprocess = transforms.Compose([
+                transforms.Resize(224, interpolation=transforms.InterpolationMode.BICUBIC, antialias=None),
+                transforms.CenterCrop(224),
+                transforms.Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711))
+        ])
+        
+        tokenizer = open_clip.get_tokenizer('ViT-H-14')
+            
+        self.clip = clip_model
+        self.norm_embs = norm_embs
+        self.preprocess = preprocess
+        self.tokenizer = tokenizer
+        self.device = device
+        
+    def embed_image(self, image):
+        """Expects images in -1 to 1 range"""
+        image = self.preprocess(image).to(self.device)
+        with torch.no_grad(), torch.cuda.amp.autocast():
+            image_features = self.clip.encode_image(image)
+        if self.norm_embs:
+            image_features = nn.functional.normalize(image_features, dim=-1)
+        return image_features
+
+    def embed_text(self, text_samples):
+        text = self.tokenizer(text_samples).to(self.device)
+        with torch.no_grad(), torch.cuda.amp.autocast():
+            text_features = self.clip.encode_text(text)
+        if self.norm_embs:
+            text_features = nn.functional.normalize(text_features, dim=-1)
+        return text_features
+
+    def embed_curated_annotations(self, annots):
+        for i,b in enumerate(annots):
+            t = ''
+            while t == '':
+                rand = torch.randint(5,(1,1))[0][0]
+                t = b[0,rand]
+            if i==0:
+                txt = np.array(t)
+            else:
+                txt = np.vstack((txt,t))
+        txt = txt.flatten()
+        return self.embed_text(txt)
+    
+class BrainNetwork(nn.Module):
+    def __init__(self, out_dim=768, in_dim=15724, clip_size=768, h=4096, n_blocks=4, norm_type='ln', act_first=False, use_projector=True, drop1=.5, drop2=.15):
+        super().__init__()
+        norm_func = partial(nn.BatchNorm1d, num_features=h) if norm_type == 'bn' else partial(nn.LayerNorm, normalized_shape=h)
+        act_fn = partial(nn.ReLU, inplace=True) if norm_type == 'bn' else nn.GELU
+        act_and_norm = (act_fn, norm_func) if act_first else (norm_func, act_fn)
+        # self.temp = nn.Parameter(torch.tensor(.006))
+        self.lin0 = nn.Sequential(
+            nn.Linear(in_dim, h),
+            *[item() for item in act_and_norm],
+            nn.Dropout(drop1),
+        )
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h),
+                *[item() for item in act_and_norm],
+                nn.Dropout(drop2)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, out_dim, bias=True)
+        self.n_blocks = n_blocks
+        self.clip_size = clip_size
+        
+        self.use_projector = use_projector
+        if use_projector:
+            self.projector = nn.Sequential(
+                nn.LayerNorm(clip_size),
+                nn.GELU(),
+                nn.Linear(clip_size, 2048),
+                nn.LayerNorm(2048),
+                nn.GELU(),
+                nn.Linear(2048, 2048),
+                nn.LayerNorm(2048),
+                nn.GELU(),
+                nn.Linear(2048, clip_size)
+            )
+        
+    def forward(self, x):
+        '''
+            bs, 1, 15724 -> bs, 32, h
+            bs, 32, h -> bs, 32h
+            b2, 32h -> bs, 768
+        '''
+        if x.ndim == 4:
+            # case when we passed 3D data of shape [N, 81, 104, 83]
+            assert x.shape[1] == 81 and x.shape[2] == 104 and x.shape[3] == 83
+            # [N, 699192]
+            x = x.reshape(x.shape[0], -1)
+        x = self.lin0(x)  # bs, h
+        residual = x
+        for res_block in range(self.n_blocks):
+            x = self.mlp[res_block](x)
+            x += residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)
+        if self.use_projector:
+            return x, self.projector(x.reshape(len(x), -1, self.clip_size))
+        return x
+    
+class BrainDiffusionPriorOld(DiffusionPrior):
+    """ 
+    Differences from original:
+    - Allow for passing of generators to torch random functions
+    - Option to include the voxel2clip model and pass voxels into forward method
+    - Return predictions when computing loss
+    - Load pretrained model from @nousr trained on LAION aesthetics
+    """
+    def __init__(self, *args, **kwargs):
+        voxel2clip = kwargs.pop('voxel2clip', None)
+        super().__init__(*args, **kwargs)
+        self.voxel2clip = voxel2clip
+
+    @torch.no_grad()
+    def p_sample(self, x, t, text_cond = None, self_cond = None, clip_denoised = True, cond_scale = 1.,
+                generator=None):
+        b, *_, device = *x.shape, x.device
+        model_mean, _, model_log_variance, x_start = self.p_mean_variance(x = x, t = t, text_cond = text_cond, self_cond = self_cond, clip_denoised = clip_denoised, cond_scale = cond_scale)
+        if generator is None:
+            noise = torch.randn_like(x)
+        else:
+            #noise = torch.randn_like(x)
+            noise = torch.randn(x.size(), device=x.device, dtype=x.dtype, generator=generator)
+        # no noise when t == 0
+        nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
+        pred = model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
+        return pred, x_start
+
+    @torch.no_grad()
+    def p_sample_loop_ddpm(self, shape, text_cond, cond_scale = 1., generator=None):
+        batch, device = shape[0], self.device
+
+        if generator is None:
+            image_embed = torch.randn(shape, device = device)
+        else:
+            image_embed = torch.randn(shape, device = device, generator=generator)
+        x_start = None # for self-conditioning
+
+        if self.init_image_embed_l2norm:
+            image_embed = l2norm(image_embed) * self.image_embed_scale
+
+        for i in tqdm(reversed(range(0, self.noise_scheduler.num_timesteps)), desc='sampling loop time step', total=self.noise_scheduler.num_timesteps, disable=True):
+            times = torch.full((batch,), i, device = device, dtype = torch.long)
+
+            self_cond = x_start if self.net.self_cond else None
+            image_embed, x_start = self.p_sample(image_embed, times, text_cond = text_cond, self_cond = self_cond, cond_scale = cond_scale, 
+                                                 generator=generator)
+
+        if self.sampling_final_clamp_l2norm and self.predict_x_start:
+            image_embed = self.l2norm_clamp_embed(image_embed)
+
+        return image_embed
+
+    def p_losses(self, image_embed, times, text_cond, noise = None):
+        noise = default(noise, lambda: torch.randn_like(image_embed))
+
+        image_embed_noisy = self.noise_scheduler.q_sample(x_start = image_embed, t = times, noise = noise)
+
+        self_cond = None
+        if self.net.self_cond and random.random() < 0.5:
+            with torch.no_grad():
+                self_cond = self.net(image_embed_noisy, times, **text_cond).detach()
+
+        pred = self.net(
+            image_embed_noisy,
+            times,
+            self_cond = self_cond,
+            text_cond_drop_prob = self.text_cond_drop_prob,
+            image_cond_drop_prob = self.image_cond_drop_prob,
+            **text_cond
+        )
+
+        if self.predict_x_start and self.training_clamp_l2norm:
+            pred = self.l2norm_clamp_embed(pred)
+
+        if self.predict_v:
+            target = self.noise_scheduler.calculate_v(image_embed, times, noise)
+        elif self.predict_x_start:
+            target = image_embed
+        else:
+            target = noise
+
+        loss = self.noise_scheduler.loss_fn(pred, target)
+        return loss, pred
+
+    def forward(
+        self,
+        text = None,
+        image = None,
+        voxel = None,
+        text_embed = None,      # allow for training on preprocessed CLIP text and image embeddings
+        image_embed = None,
+        text_encodings = None,  # as well as CLIP text encodings
+        *args,
+        **kwargs
+    ):
+        assert exists(text) ^ exists(text_embed) ^ exists(voxel), 'either text, text embedding, or voxel must be supplied'
+        assert exists(image) ^ exists(image_embed), 'either image or image embedding must be supplied'
+        assert not (self.condition_on_text_encodings and (not exists(text_encodings) and not exists(text))), 'text encodings must be present if you specified you wish to condition on it on initialization'
+
+        if exists(voxel):
+            assert exists(self.voxel2clip), 'voxel2clip must be trained if you wish to pass in voxels'
+            assert not exists(text_embed), 'cannot pass in both text and voxels'
+            text_embed = self.voxel2clip(voxel)
+
+        if exists(image):
+            image_embed, _ = self.clip.embed_image(image)
+
+        # calculate text conditionings, based on what is passed in
+
+        if exists(text):
+            text_embed, text_encodings = self.clip.embed_text(text)
+
+        text_cond = dict(text_embed = text_embed)
+
+        if self.condition_on_text_encodings:
+            assert exists(text_encodings), 'text encodings must be present for diffusion prior if specified'
+            text_cond = {**text_cond, 'text_encodings': text_encodings}
+
+        # timestep conditioning from ddpm
+
+        batch, device = image_embed.shape[0], image_embed.device
+        times = self.noise_scheduler.sample_random_times(batch)
+
+        # scale image embed (Katherine)
+
+        image_embed *= self.image_embed_scale
+
+        # calculate forward loss
+
+        loss, pred = self.p_losses(image_embed, times, text_cond = text_cond, *args, **kwargs)
+
+        return loss, pred#, text_embed
+   
+    @staticmethod
+    def from_pretrained(net_kwargs={}, prior_kwargs={}, voxel2clip_path=None, ckpt_dir='./checkpoints'):
+        # "https://huggingface.co/nousr/conditioned-prior/raw/main/vit-l-14/aesthetic/prior_config.json"
+        config_url = os.path.join(ckpt_dir, "prior_config.json")
+        config = json.load(open(config_url))
+        
+        config['prior']['net']['max_text_len'] = 256
+        config['prior']['net'].update(net_kwargs)
+        # print('net_config', config['prior']['net'])
+        net_config = DiffusionPriorNetworkConfig(**config['prior']['net'])
+
+        kwargs = config['prior']
+        kwargs.pop('clip')
+        kwargs.pop('net')
+        kwargs.update(prior_kwargs)
+        # print('prior_config', kwargs)
+
+        diffusion_prior_network = net_config.create()
+        diffusion_prior = BrainDiffusionPriorOld(net=diffusion_prior_network, clip=None, **kwargs).to(torch.device('cpu'))
+        
+        # 'https://huggingface.co/nousr/conditioned-prior/resolve/main/vit-l-14/aesthetic/best.pth'
+        ckpt_url = os.path.join(ckpt_dir, 'best.pth')
+        ckpt = torch.load(ckpt_url, map_location=torch.device('cpu'))
+
+        # Note these keys will be missing (maybe due to an update to the code since training):
+        # "net.null_text_encodings", "net.null_text_embeds", "net.null_image_embed"
+        # I don't think these get used if `cond_drop_prob = 0` though (which is the default here)
+        diffusion_prior.load_state_dict(ckpt, strict=False)
+        # keys = diffusion_prior.load_state_dict(ckpt, strict=False)
+        # print("missing keys in prior checkpoint (probably ok)", keys.missing_keys)
+
+        if voxel2clip_path:
+            # load the voxel2clip weights
+            checkpoint = torch.load(voxel2clip_path, map_location=torch.device('cpu'))
+            
+            state_dict = checkpoint['model_state_dict']
+            for key in list(state_dict.keys()):
+                if 'module.' in key:
+                    state_dict[key.replace('module.', '')] = state_dict[key]
+                    del state_dict[key]
+            diffusion_prior.voxel2clip.load_state_dict(state_dict)
+        
+        return diffusion_prior
+
+class BrainDiffusionPrior(DiffusionPrior):
+    """ 
+    Differences from original:
+    - Allow for passing of generators to torch random functions
+    - Option to include the voxel2clip model and pass voxels into forward method
+    - Return predictions when computing loss
+    - Load pretrained model from @nousr trained on LAION aesthetics
+    """
+    def __init__(self, *args, **kwargs):
+        voxel2clip = kwargs.pop('voxel2clip', None)
+        super().__init__(*args, **kwargs)
+        self.voxel2clip = voxel2clip
+
+    @torch.no_grad()
+    def p_sample(self, x, t, text_cond = None, self_cond = None, clip_denoised = True, cond_scale = 1.,
+                generator=None):
+        b, *_, device = *x.shape, x.device
+        model_mean, _, model_log_variance, x_start = self.p_mean_variance(x = x, t = t, text_cond = text_cond, self_cond = self_cond, clip_denoised = clip_denoised, cond_scale = cond_scale)
+        if generator is None:
+            noise = torch.randn_like(x)
+        else:
+            #noise = torch.randn_like(x)
+            noise = torch.randn(x.size(), device=x.device, dtype=x.dtype, generator=generator)
+        # no noise when t == 0
+        nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
+        pred = model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
+        return pred, x_start
+
+    @torch.no_grad()
+    def p_sample_loop_ddpm(self, shape, text_cond, cond_scale = 1., generator=None):
+        batch, device = shape[0], self.device
+
+        if generator is None:
+            image_embed = torch.randn(shape, device = device)
+        else:
+            image_embed = torch.randn(shape, device = device, generator=generator)
+        x_start = None # for self-conditioning
+
+        if self.init_image_embed_l2norm:
+            image_embed = l2norm(image_embed) * self.image_embed_scale
+
+        for i in tqdm(reversed(range(0, self.noise_scheduler.num_timesteps)), desc='sampling loop time step', total=self.noise_scheduler.num_timesteps, disable=True):
+            times = torch.full((batch,), i, device = device, dtype = torch.long)
+
+            self_cond = x_start if self.net.self_cond else None
+            image_embed, x_start = self.p_sample(image_embed, times, text_cond = text_cond, self_cond = self_cond, cond_scale = cond_scale, 
+                                                 generator=generator)
+
+        if self.sampling_final_clamp_l2norm and self.predict_x_start:
+            image_embed = self.l2norm_clamp_embed(image_embed)
+
+        return image_embed
+
+    def p_losses(self, image_embed, times, text_cond, noise = None):
+        noise = default(noise, lambda: torch.randn_like(image_embed))
+
+        image_embed_noisy = self.noise_scheduler.q_sample(x_start = image_embed, t = times, noise = noise)
+
+        self_cond = None
+        if self.net.self_cond and random.random() < 0.5:
+            with torch.no_grad():
+                self_cond = self.net(image_embed_noisy, times, **text_cond).detach()
+
+        pred = self.net(
+            image_embed_noisy,
+            times,
+            self_cond = self_cond,
+            text_cond_drop_prob = self.text_cond_drop_prob,
+            image_cond_drop_prob = self.image_cond_drop_prob,
+            **text_cond
+        )
+
+        if self.predict_x_start and self.training_clamp_l2norm:
+            pred = self.l2norm_clamp_embed(pred)
+
+        if self.predict_v:
+            target = self.noise_scheduler.calculate_v(image_embed, times, noise)
+        elif self.predict_x_start:
+            target = image_embed
+        else:
+            target = noise
+
+        loss = self.noise_scheduler.loss_fn(pred, target)
+        return loss, pred
+
+    def forward(
+        self,
+        text = None,
+        image = None,
+        voxel = None,
+        text_embed = None,      # allow for training on preprocessed CLIP text and image embeddings
+        image_embed = None,
+        text_encodings = None,  # as well as CLIP text encodings
+        *args,
+        **kwargs
+    ):
+        assert exists(text) ^ exists(text_embed) ^ exists(voxel), 'either text, text embedding, or voxel must be supplied'
+        assert exists(image) ^ exists(image_embed), 'either image or image embedding must be supplied'
+        assert not (self.condition_on_text_encodings and (not exists(text_encodings) and not exists(text))), 'text encodings must be present if you specified you wish to condition on it on initialization'
+
+        if exists(voxel):
+            assert exists(self.voxel2clip), 'voxel2clip must be trained if you wish to pass in voxels'
+            assert not exists(text_embed), 'cannot pass in both text and voxels'
+            if self.voxel2clip.use_projector:
+                clip_voxels_mse, clip_voxels = self.voxel2clip(voxel)
+                text_embed = clip_voxels_mse
+            else:
+                clip_voxels = self.voxel2clip(voxel)
+                text_embed = clip_voxels_mse = clip_voxels
+            # text_embed = self.voxel2clip(voxel)
+
+        if exists(image):
+            image_embed, _ = self.clip.embed_image(image)
+
+        # calculate text conditionings, based on what is passed in
+
+        if exists(text):
+            text_embed, text_encodings = self.clip.embed_text(text)
+
+        text_cond = dict(text_embed = text_embed)
+
+        if self.condition_on_text_encodings:
+            assert exists(text_encodings), 'text encodings must be present for diffusion prior if specified'
+            text_cond = {**text_cond, 'text_encodings': text_encodings}
+
+        # timestep conditioning from ddpm
+
+        batch, device = image_embed.shape[0], image_embed.device
+        times = self.noise_scheduler.sample_random_times(batch)
+
+        # PS: I dont think we need this? also if uncommented this does in-place global variable change
+        # scale image embed (Katherine)
+        # image_embed *= self.image_embed_scale
+
+        # calculate forward loss
+
+        loss, pred = self.p_losses(image_embed*self.image_embed_scale, times, text_cond = text_cond, *args, **kwargs)
+        
+        # undo the scaling so we can directly use it for real mse loss and reconstruction
+        return loss, pred
+    
+class BrainSD(StableDiffusionImageVariationPipeline):
+    """ 
+    Differences from original:
+    - Keep generated images on GPU and return tensors
+    - No NSFW checker
+    - Can pass in image or image_embedding to generate a variation
+    NOTE: requires latest version of diffusers to avoid the latent dims not being correct.
+    """
+
+    def decode_latents(self, latents):
+        latents = 1 / 0.18215 * latents
+        image = self.vae.decode(latents).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+        # image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        return image
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        image: Optional[Union[PIL.Image.Image, List[PIL.Image.Image], torch.FloatTensor]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: Optional[int] = 1,
+        image_embeddings: Optional[torch.FloatTensor] = None,
+    ):
+
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        if image_embeddings is None:
+            assert image is not None, "If image_embeddings is None, image must not be None"
+    
+            # resize and normalize the way that's recommended
+            tform = transforms.Compose([
+                #transforms.ToTensor(), ## don't need this since we've already got tensors
+                transforms.Resize(
+                    (224, 224),
+                    interpolation=transforms.InterpolationMode.BICUBIC,
+                    antialias=False,
+                    ),
+                transforms.Normalize(
+                    [0.48145466, 0.4578275, 0.40821073],
+                    [0.26862954, 0.26130258, 0.27577711]),
+            ])
+            image = tform(image)
+
+            # 1. Check inputs. Raise error if not correct
+            self.check_inputs(image, height, width, callback_steps)
+
+            # 2. Define call parameters
+            if isinstance(image, PIL.Image.Image):
+                batch_size = 1
+            elif isinstance(image, list):
+                batch_size = len(image)
+            else:
+                batch_size = image.shape[0]
+
+            # 3. Encode input image
+            image_embeddings = self._encode_image(image, device, num_images_per_prompt, do_classifier_free_guidance)
+        else:
+            batch_size = image_embeddings.shape[0] // 2
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.in_channels
+
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            image_embeddings.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                # predict the noise residual
+                noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=image_embeddings).sample
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+
+        # 8. Post-processing
+        image = self.decode_latents(latents)
+
+        # # 9. Run safety checker
+        # image, has_nsfw_concept = self.run_safety_checker(image, device, image_embeddings.dtype)
+
+        # # 10. Convert to PIL
+        # if output_type == "pil":
+        #     image = self.numpy_to_pil(image)
+
+        # if not return_dict:
+        #     return (image, has_nsfw_concept)
+
+        # return StableDiffusionPipelineOutput(images=image)
+
+        return image
+
+class Voxel2StableDiffusionModel(torch.nn.Module):
+    def __init__(self, in_dim=15724, h=4096, n_blocks=4, use_cont=False):
+        super().__init__()
+        self.lin0 = nn.Sequential(
+            nn.Linear(in_dim, h, bias=False),
+            nn.LayerNorm(h),
+            nn.SiLU(inplace=True),
+            nn.Dropout(0.5),
+        )
+
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(h, h, bias=False),
+                nn.LayerNorm(h),
+                nn.SiLU(inplace=True),
+                nn.Dropout(0.25)
+            ) for _ in range(n_blocks)
+        ])
+        self.lin1 = nn.Linear(h, 16384, bias=False)
+        self.norm = nn.LayerNorm(512)
+
+        self.register_parameter('queries', nn.Parameter(torch.randn(1, 256, 512) * 0.044))
+        self.transformer = nn.TransformerDecoder(
+            nn.TransformerDecoderLayer(d_model=512, nhead=8, norm_first=True,
+                                        dim_feedforward=1024, activation=nn.functional.gelu,
+                                        batch_first=True, dropout=0.25),
+            num_layers=n_blocks
+        )
+
+        # option 1  -> 124.56M
+        # self.lin1 = nn.Linear(h, 32768, bias=True)
+        # self.upsampler = Decoder(
+        #     in_channels=64,
+        #     out_channels=4,
+        #     up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D"],
+        #     block_out_channels=[64, 128, 256, 256],
+        #     layers_per_block=1,
+        # )
+
+        # option2  -> 132.52M
+        # self.lin1 = nn.Linear(h, 1024, bias=True)
+        # self.upsampler = Decoder(
+        #     in_channels=64,
+        #     out_channels=4,
+        #     up_block_types=["UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D","UpDecoderBlock2D", "UpDecoderBlock2D"],
+        #     block_out_channels=[64, 128, 256, 256, 512],
+        #     layers_per_block=1,
+        # )
+        
+        if use_cont:
+            self.maps_projector = nn.Sequential(
+                nn.LayerNorm(512),
+                nn.Linear(512, 512),
+                nn.LayerNorm(512),
+                nn.ReLU(True),
+                nn.Linear(512, 512),
+                nn.LayerNorm(512),
+                nn.ReLU(True),
+                nn.Linear(512, 512)
+            )
+        else:
+            self.maps_projector = nn.Identity()
+
+        self.upsampler = nn.Sequential(
+            nn.GroupNorm(1, 32),
+            nn.SiLU(inplace=True),
+            nn.Conv2d(32, 320, 3, padding=1),
+            nn.GroupNorm(32, 320),
+            nn.SiLU(inplace=True),
+            nn.Conv2d(320, 320, 3, padding=1),
+            nn.GroupNorm(32, 320),
+            nn.SiLU(inplace=True),
+            nn.Conv2d(320, 4, 3, padding=1)
+        )
+
+    def forward(self, x, return_transformer_feats=False):
+        x = self.lin0(x)
+        residual = x
+        for res_block in self.mlp:
+            x = res_block(x)
+            x = x + residual
+            residual = x
+        x = x.reshape(len(x), -1)
+        x = self.lin1(x)  # bs, 4096
+
+        # # x = x.reshape(x.shape[0], -1, 8, 8).contiguous()  # bs, 64, 8, 8
+        # x = x.reshape(x.shape[0], -1, 64, 64).contiguous()
+        # return self.upsampler(x)
+
+        # decoder
+        x = self.norm(x.reshape(x.shape[0], 32, 512))
+        preds = self.transformer(self.queries.expand(x.shape[0], -1, -1), x)
+        sd_embeds = preds.permute(0,2,1).reshape(-1, 512, 16, 16)
+        sd_embeds = nn.functional.pixel_shuffle(sd_embeds, 4)  # bs, 32, 32, 32
+
+        # contrastive
+        if return_transformer_feats:
+            return self.upsampler(sd_embeds), self.maps_projector(preds)
+        
+        return self.upsampler(sd_embeds)
+    
+class BrainNetworkDETR(BrainNetwork):
+    # 133M
+    def __init__(self, out_dim=768, in_dim=15724, h=4096, n_blocks=4, norm_type='ln', act_first=False, 
+                encoder_tokens=32, decoder_tokens=257):
+        # encoder
+        super().__init__(out_dim*encoder_tokens, in_dim, h, n_blocks, norm_type, act_first)
+        self.norm = nn.LayerNorm(out_dim)
+        self.encoder_tokens = encoder_tokens
+        
+        self.register_parameter('queries', nn.Parameter(torch.randn(1, decoder_tokens, out_dim)))
+        self.transformer = nn.TransformerDecoder(
+            nn.TransformerDecoderLayer(d_model=out_dim, nhead=8,
+                                        dim_feedforward=1024, 
+                                        batch_first=True, dropout=0.25),
+            num_layers=n_blocks
+        )
+        self.decoder_projector = nn.Sequential(
+            nn.LayerNorm(out_dim),
+            nn.Linear(out_dim, out_dim)
+        )
+
+
+    def forward(self, x):
+        enc = super().forward(x)
+        enc = self.norm(enc.reshape(enc.shape[0], self.encoder_tokens, -1))
+
+        dec = self.transformer(self.queries.expand(x.shape[0], -1, -1), enc)
+        dec = self.decoder_projector(dec)
+        return dec
+
+class VersatileDiffusionPriorNetwork(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_timesteps = None,
+        num_time_embeds = 1,
+        # num_image_embeds = 1,
+        # num_brain_embeds = 1,
+        num_tokens = 257,
+        causal = True,
+        learned_query_mode = 'none',
+        **kwargs
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_time_embeds = num_time_embeds
+        self.continuous_embedded_time = not exists(num_timesteps)
+        self.learned_query_mode = learned_query_mode
+
+        self.to_time_embeds = nn.Sequential(
+            nn.Embedding(num_timesteps, dim * num_time_embeds) if exists(num_timesteps) else nn.Sequential(SinusoidalPosEmb(dim), MLP(dim, dim * num_time_embeds)), # also offer a continuous version of timestep embeddings, with a 2 layer MLP
+            Rearrange('b (n d) -> b n d', n = num_time_embeds)
+        )
+
+        if self.learned_query_mode == 'token':
+            self.learned_query = nn.Parameter(torch.randn(num_tokens, dim))
+        if self.learned_query_mode == 'pos_emb':
+            scale = dim ** -0.5
+            self.learned_query = nn.Parameter(torch.randn(num_tokens, dim) * scale)
+        if self.learned_query_mode == 'all_pos_emb':
+            scale = dim ** -0.5
+            self.learned_query = nn.Parameter(torch.randn(num_tokens*2+1, dim) * scale)
+        self.causal_transformer = FlaggedCausalTransformer(dim = dim, causal=causal, **kwargs)
+
+        self.null_brain_embeds = nn.Parameter(torch.randn(num_tokens, dim))
+        self.null_image_embed = nn.Parameter(torch.randn(num_tokens, dim))
+
+        self.num_tokens = num_tokens
+        self.self_cond = False
+
+    def forward_with_cond_scale(
+        self,
+        *args,
+        cond_scale = 1.,
+        **kwargs
+    ):
+        logits = self.forward(*args, **kwargs)
+
+        if cond_scale == 1:
+            return logits
+
+        null_logits = self.forward(*args, brain_cond_drop_prob = 1., image_cond_drop_prob = 1, **kwargs)
+        return null_logits + (logits - null_logits) * cond_scale
+
+    def forward(
+        self,
+        image_embed,
+        diffusion_timesteps,
+        *,
+        self_cond=None,
+        brain_embed=None,
+        text_embed=None,
+        brain_cond_drop_prob = 0.,
+        text_cond_drop_prob = None,
+        image_cond_drop_prob = 0.
+    ):
+        if text_embed is not None:
+            brain_embed = text_embed
+        if text_cond_drop_prob is not None:
+            brain_cond_drop_prob = text_cond_drop_prob
+        
+        image_embed = image_embed.view(len(image_embed),-1,768)
+        # text_embed = text_embed.view(len(text_embed),-1,768)
+        brain_embed = brain_embed.view(len(brain_embed),-1,768)
+        # print(*image_embed.shape)
+        # print(*image_embed.shape, image_embed.device, image_embed.dtype)
+        
+        batch, _, dim, device, dtype = *image_embed.shape, image_embed.device, image_embed.dtype
+        # num_time_embeds, num_image_embeds, num_brain_embeds = self.num_time_embeds, self.num_image_embeds, self.num_brain_embeds
+        
+        # classifier free guidance masks
+        brain_keep_mask = prob_mask_like((batch,), 1 - brain_cond_drop_prob, device = device)
+        brain_keep_mask = rearrange(brain_keep_mask, 'b -> b 1 1')
+
+        image_keep_mask = prob_mask_like((batch,), 1 - image_cond_drop_prob, device = device)
+        image_keep_mask = rearrange(image_keep_mask, 'b -> b 1 1')
+
+        # mask out brain embeddings with null brain embeddings
+
+        # import pdb; pdb.set_trace()
+        null_brain_embeds = self.null_brain_embeds.to(brain_embed.dtype)
+        brain_embed = torch.where(
+            brain_keep_mask,
+            brain_embed,
+            null_brain_embeds[None]
+        )
+
+        # mask out image embeddings with null image embeddings
+        null_image_embed = self.null_image_embed.to(image_embed.dtype)
+        image_embed = torch.where(
+            image_keep_mask,
+            image_embed,
+            null_image_embed[None]
+        )
+
+        # whether brain embedding is used for conditioning depends on whether brain encodings are available for attention (for classifier free guidance, even though it seems from the paper it was not used in the prior ddpm, as the objective is different)
+        # but let's just do it right
+        if self.continuous_embedded_time:
+            # if continuous cast to flat, else keep int for indexing embeddings
+            diffusion_timesteps = diffusion_timesteps.type(dtype)
+        time_embed = self.to_time_embeds(diffusion_timesteps)
+
+        if self.learned_query_mode == 'token':
+            learned_queries = repeat(self.learned_query, 'n d -> b n d', b = batch)
+        elif self.learned_query_mode == 'pos_emb':
+            pos_embs = repeat(self.learned_query, 'n d -> b n d', b = batch)
+            image_embed = image_embed + pos_embs
+            learned_queries = torch.empty((batch, 0, dim), device=brain_embed.device)
+        elif self.learned_query_mode == 'all_pos_emb':
+            pos_embs = repeat(self.learned_query, 'n d -> b n d', b = batch)
+            learned_queries = torch.empty((batch, 0, dim), device=brain_embed.device)
+        else:
+            learned_queries = torch.empty((batch, 0, dim), device=brain_embed.device)
+        
+        tokens = torch.cat((
+            brain_embed,  # 257
+            time_embed,  # 1
+            image_embed,  # 257
+            learned_queries  # 257
+        ), dim = -2)
+        if self.learned_query_mode == 'all_pos_emb':
+            tokens = tokens + pos_embs
+
+        # attend
+        tokens = self.causal_transformer(tokens)
+
+        # get learned query, which should predict the image embedding (per DDPM timestep)
+        pred_image_embed = tokens[..., -self.num_tokens:, :]
+
+        return pred_image_embed
+
+# import math
+# from collections import namedtuple
+# from einops import rearrange, repeat, reduce, pack, unpack
+# from einops.layers.torch import Rearrange
+# from torch import einsum
+# class Attention(nn.Module):
+#     def __init__(
+#         self,
+#         dim,
+#         *,
+#         dim_head = 64,
+#         heads = 8,
+#         dropout = 0.,
+#         causal = False,
+#         rotary_emb = None,
+#         cosine_sim = True,
+#         cosine_sim_scale = 16
+#     ):
+#         super().__init__()
+#         self.scale = cosine_sim_scale if cosine_sim else (dim_head ** -0.5)
+#         self.cosine_sim = cosine_sim
+
+#         self.heads = heads
+#         inner_dim = dim_head * heads
+#         self.dim = dim
+#         self.inner_dim = inner_dim
+
+#         self.causal = causal
+#         self.norm = LayerNorm(dim)
+#         self.dropout = nn.Dropout(dropout)
+
+#         self.null_kv = nn.Parameter(torch.randn(2, dim_head))
+#         self.to_q = nn.Linear(dim, inner_dim, bias = False)
+#         self.to_kv = nn.Linear(dim, dim_head * 2, bias = False)
+
+#         self.rotary_emb = rotary_emb
+
+#         self.to_out = nn.Sequential(
+#             nn.Linear(inner_dim, dim, bias = False),
+#             LayerNorm(dim)
+#         )
+
+#     def forward(self, x, mask = None, attn_bias = None):
+#         b, n, device = *x.shape[:2], x.device
+#         print("xinit", torch.any(torch.isnan(x)))
+#         x = self.norm(x)
+#         print("xnorm", torch.any(torch.isnan(x)))
+#         print("xnorm.shape", x.shape)
+
+#         q = self.to_q(x)
+#         print("q0", torch.any(torch.isnan(q)))
+        
+#         k, v = self.to_kv(x).chunk(2, dim = -1)
+#         print("k0", torch.any(torch.isnan(k)))
+#         print("v0", torch.any(torch.isnan(v)))
+        
+
+#         q = rearrange(q, 'b n (h d) -> b h n d', h = self.heads)
+#         q = q * self.scale
+
+#         # rotary embeddings
+
+#         if exists(self.rotary_emb):
+#             q, k = map(self.rotary_emb.rotate_queries_or_keys, (q, k))
+
+#         # add null key / value for classifier free guidance in prior net
+
+#         nk, nv = map(lambda t: repeat(t, 'd -> b 1 d', b = b), self.null_kv.unbind(dim = -2))
+#         k = torch.cat((nk, k), dim = -2)
+#         v = torch.cat((nv, v), dim = -2)
+
+#         # whether to use cosine sim
+
+#         if self.cosine_sim:
+#             q, k = map(l2norm, (q, k))
+
+#         q, k = map(lambda t: t * math.sqrt(self.scale), (q, k))
+#         print("q2", torch.any(torch.isnan(q)))
+#         print("k2", torch.any(torch.isnan(k)))
+
+#         # calculate query / key similarities
+
+#         sim = einsum('b h i d, b j d -> b h i j', q, k)
+
+#         # relative positional encoding (T5 style)
+
+#         if exists(attn_bias):
+#             sim = sim + attn_bias
+
+#         # masking
+
+#         max_neg_value = -torch.finfo(sim.dtype).max
+
+#         print("sim1", torch.any(torch.isnan(sim)))
+
+#         if exists(mask):
+#             mask = F.pad(mask, (1, 0), value = True)
+#             mask = rearrange(mask, 'b j -> b 1 1 j')
+#             sim = sim.masked_fill(~mask, max_neg_value)
+
+#         print("sim2", torch.any(torch.isnan(sim)))
+        
+#         if self.causal:
+#             i, j = sim.shape[-2:]
+#             causal_mask = torch.ones((i, j), dtype = torch.bool, device = device).triu(j - i + 1)
+#             sim = sim.masked_fill(causal_mask, max_neg_value)
+
+#         # attention
+
+#         print("simFinal", torch.any(torch.isnan(sim)))
+#         attn = sim.softmax(dim = -1, dtype = torch.float32)
+#         print("attn", torch.any(torch.isnan(attn)))
+#         attn = attn.type(sim.dtype)
+
+#         attn = self.dropout(attn)
+
+#         # aggregate values
+
+#         out = einsum('b h i j, b j d -> b h i d', attn, v)
+
+#         out = rearrange(out, 'b h n d -> b n (h d)')
+#         return self.to_out(out)
+
+class FlaggedCausalTransformer(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth,
+        dim_head = 64,
+        heads = 8,
+        ff_mult = 4,
+        norm_in = False,
+        norm_out = True,
+        attn_dropout = 0.,
+        ff_dropout = 0.,
+        final_proj = True,
+        normformer = False,
+        rotary_emb = False,
+        causal=True
+    ):
+        super().__init__()
+        self.init_norm = LayerNorm(dim) if norm_in else nn.Identity() # from latest BLOOM model and Yandex's YaLM
+
+        self.rel_pos_bias = RelPosBias(heads = heads)
+
+        rotary_emb = RotaryEmbedding(dim = min(32, dim_head)) if rotary_emb else None
+
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                Attention(dim = dim, causal = causal, dim_head = dim_head, heads = heads, dropout = attn_dropout, rotary_emb = rotary_emb),
+                FeedForward(dim = dim, mult = ff_mult, dropout = ff_dropout, post_activation_norm = normformer)
+            ]))
+
+        self.norm = LayerNorm(dim, stable = True) if norm_out else nn.Identity()  # unclear in paper whether they projected after the classic layer norm for the final denoised image embedding, or just had the transformer output it directly: plan on offering both options
+        self.project_out = nn.Linear(dim, dim, bias = False) if final_proj else nn.Identity()
+
+    def forward(self, x):
+        n, device = x.shape[1], x.device
+
+        x = self.init_norm(x)
+
+        attn_bias = self.rel_pos_bias(n, n + 1, device = device)
+
+        for attn, ff in self.layers:
+            x = attn(x, attn_bias = attn_bias) + x
+            x = ff(x) + x
+
+        out = self.norm(x)
+        return self.project_out(out)
+    
+class BrainVD(VersatileDiffusionDualGuidedPipeline):
+    """ 
+    Differences from original:
+    - Keep generated images on GPU and return tensors
+    - No NSFW checker
+    - Can pass in image or image_embedding to generate a variation
+    NOTE: requires latest version of diffusers to avoid the latent dims not being correct.
+    """
+
+    def decode_latents(self, latents):
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        # image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        return image
+
+    def check_inputs(self, prompt, image, height, width, callback_steps):
+        if prompt is not None and not isinstance(prompt, str) and not isinstance(prompt, list):
+            raise ValueError(f"`prompt` has to be of type None, `str` or `list` but is {type(prompt)}")
+        if image is not None and not isinstance(image, PIL.Image.Image) and not isinstance(image, list):
+            raise ValueError(f"`image` has to be of type None, `PIL.Image` or `list` but is {type(image)}")
+
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if (callback_steps is None) or (
+            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
+        ):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}."
+            )
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[PIL.Image.Image, List[PIL.Image.Image]] = None,
+        image: Union[str, List[str]] = None,
+        text_to_image_strength: float = 0.5,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: Optional[int] = 1,
+        image_embeddings: Optional[torch.FloatTensor] = None,
+        prompt_embeddings: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ):
+
+        height = height or self.image_unet.config.sample_size * self.vae_scale_factor
+        width = width or self.image_unet.config.sample_size * self.vae_scale_factor
+
+        self.check_inputs(prompt, image, height, width, callback_steps)
+
+        prompt = [prompt] if prompt is not None and not isinstance(prompt, list) else prompt
+        image = [image] if image is not None and not isinstance(image, list) else image
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+
+        # 3. Encode input prompt
+        if image_embeddings is None:
+            if image is not None:
+                image_embeddings = self._encode_image_prompt(
+                    image, device, num_images_per_prompt, do_classifier_free_guidance
+                )
+                batch_size = len(image)
+            else:
+                image_embeddings = None
+        
+        if prompt_embeddings is None:
+            if prompt is not None:
+                prompt_embeddings = self._encode_text_prompt(
+                    prompt, device, num_images_per_prompt, do_classifier_free_guidance
+                )
+                batch_size = len(prompt)
+            else:
+                prompt_embeddings = None
+        if image_embeddings is not None:
+            batch_size = image_embeddings.shape[0] // 2
+        elif prompt_embeddings is not None:
+            batch_size = prompt_embeddings.shape[0] // 2
+        
+        if image_embeddings is not None and prompt_embeddings is not None:
+            dual_prompt_embeddings = torch.cat([prompt_embeddings, image_embeddings], dim=1)
+        elif image_embeddings is None:
+            dual_prompt_embeddings = prompt_embeddings
+            text_to_image_strength = 1.
+        elif prompt_embeddings is None:
+            dual_prompt_embeddings = image_embeddings
+            text_to_image_strength = 0.
+        else:
+            raise ValueError()
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.image_unet.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            dual_prompt_embeddings.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs.
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Combine the attention blocks of the image and text UNets
+        self.set_transformer_params(text_to_image_strength, ("text", "image"))
+
+        # 8. Denoising loop
+        for i, t in enumerate(self.progress_bar(timesteps)):
+            # expand the latents if we are doing classifier free guidance
+            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+            # predict the noise residual
+            noise_pred = self.image_unet(latent_model_input, t, encoder_hidden_states=dual_prompt_embeddings).sample
+
+            # perform guidance
+            if do_classifier_free_guidance:
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+
+            # call the callback, if provided
+            if callback is not None and i % callback_steps == 0:
+                callback(i, t, latents)
+
+        # 8. Post-processing
+        image = self.decode_latents(latents)
+
+        return image

src/setup.sh

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

src/utils.py

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