DOCSEARCHbyAlgolia
Submit Docs

MindEye2: fMRI-to-Image Reconstruction and Retrieval

repository·main·Indexed 42 days ago

https://github.com/medarc-ai/mindeyev2

A state-of-the-art framework for reconstructing and retrieving images from fMRI data. It utilizes shared-subject models to achieve high performance with limited training data and includes tools for multi-subject pretraining, single-subject fine-tuning, and a multi-stage inference pipeline for image refinement and evaluation.

Tokens
24.7K
Snippets
64
Records
80
Agent score
53%

What's inside MindEye2

  1. How the Generative Models codebase works

    main

    The codebase follows a config-driven approach centered on modularity. Submodules are built and combined by calling instantiate_from_config() on objects defined in YAML configuration files (found in the configs/ directory).

    Core Architectural Changes (vs. old ldm codebase):

    • DiffusionEngine: The core diffusion model class (formerly LatentDiffusion) now handles all types of conditioning (vectors, sequences, and spatial conditionings) via a single GeneralConditioner class.
    • Decoupled Components: Guiders (e.g., classifier-free guidance in sgm/modules/diffusionmodules/guiders.py) are separated from samplers (sgm/modules/diffusionmodules/sampling.py). Samplers are independent of the model.
    • Denoiser Framework: The codebase adopts a denoiser framework where discrete-time models are treated as special cases of continuous-time models. Features like diffusion loss weighting, network preconditioning, and noise level sampling are now independent components.

  2. Configure Network, Loss, and Sampler

    main

    Generative model training is controlled by three primary configuration blocks:

    1. Network (network_config): Defines the neural network backbone (e.g., diffusion backbones, including transformers).
    2. Loss (loss_config): Defines the loss function. For standard diffusion training, you must also set sigma_sampler_config.
    3. Sampler (sampler_config): Defines the numerical solver, number of steps, discretization type, and guidance wrappers (like classifier-free guidance). The sampler is independent of the model architecture.

  3. Handle datasets for training

    main

    The project supports two main data handling approaches:

    1. Large-scale training: Recommended to use the data pipelines project (included via installation). This expects data in webdataset-format.
    2. Small map-style datasets: Defined within the repository. These must return a dictionary of data keys/values. Images must be in channel-first format with pixel values in the range [-1, 1].

    Example data dictionary structure:

    example = {
    "jpg": x,  # x is a tensor in range -1...1, format: [C, H, W]
    "txt": "a beautiful image"
}
    example = {"jpg": x,  # this is a tensor -1...1 chw
           "txt": "a beautiful image"}

  4. Train and fine-tune MindEye2 models

    main

    Training is primarily handled via src/Train.ipynb. You can run this as a notebook or convert it to a .py script using argparser arguments.

    Training Modes

    • Multi-subject Pretraining: Set --multi_subject and --subj=# (where # is the subject ID from NSD to exclude from pretraining).
    • Single-subject Training: Set --no-multi_subject and --subj=# (where # is the subject you wish to train on).
    • Fine-tuning: To fine-tune a single-subject model from a pretrained multi-subject checkpoint, set --no-multi_subject and --multisubject_ckpt=path_to_your_pretrained_ckpt_folder.

    Resource Management & Optimization

    • Low Compute: If you lack an 8xA100 80GB node, avoid OOM by:
      • Lowering hidden_dim to 1024 or 512.
      • Removing the low-level submodule with --no-blurry_recon.
      • Lowering the batch size.
    • Multi-GPU: If using multiple GPUs, configure accelerate to use DeepSpeed Stage 2 with CPU offloading via accelerate config.
    # Example: Multi-subject pretraining
accelerate launch --mixed_precision=fp16 Train.py --model_name=final_multisubject_subj0# --multi_subject --subj=# --batch_size=42 --max_lr=3e-4 --mixup_pct=.33 --num_epochs=150 --use_prior --prior_scale=30 --clip_scale=1 --blurry_recon --blur_scale=.5 --no-use_image_aug --n_blocks=4 --hidden_dim=4096 --num_sessions=40

  5. Install MindEye2

    main

    Follow these steps to set up the MindEye2 environment:

    1. Data Access: Agree to the Natural Scenes Dataset's Terms and Conditions and complete the NSD Data Access form.
    2. Clone Repository:
      git clone https://github.com/MedARC-AI/MindEyeV2.git
cd MindEyeV2/src
    3. Download Dataset Subset: Download the necessary files from Hugging Face into the src folder. Note that downloading the full dataset exceeds 100 GB; use the provided Python script to download only the required subset for training, inference, and evaluation.
    4. Environment Setup: Run the setup script to create a virtual environment named fmri.
      ./setup.sh
source fmri/bin/activate
    git clone https://github.com/MedARC-AI/MindEyeV2.git
cd MindEyeV2/src
./setup.sh
source fmri/bin/activate

  6. Run the Streamlit sampling demo

    main

    The repository provides a Streamlit-based demo for text-to-image and image-to-image sampling via scripts/demo/sampling.py.

    Before running the demo, you must download the model weights (e.g., from Hugging Face) and place them in the checkpoints/ directory.

    Run the demo with the following command:

    streamlit run scripts/demo/sampling.py --server.port <your_port>

  7. Download MindEye2 dataset subset via Python

    main

    Use this script to download only the files required for training, inference, and evaluation from the pscotti/mindeyev2 Hugging Face dataset, avoiding the full 100GB+ download.

    import os
from huggingface_hub import list_repo_files, hf_hub_download

repo_id, branch, exclude_dirs, exclude_files = "pscotti/mindeyev2", "main", ["train_logs", "evals"], ["human_trials_mindeye2.ipynb", "subj01_annots.npy", "shared1000.npy"]

include_specific_files = ["evals/all_images.pt", "evals/all_captions.pt", "evals/all_git_generated_captions.pt"]

def download_files(repo_id, branch, exclude_dirs, exclude_files, include_specific_files):
    files = list_repo_files(repo_id, repo_type="dataset", revision=branch)
    for file_path in files:
        if (not any(ex_dir in file_path for ex_dir in exclude_dirs) or file_path in include_specific_files) and not any(ex_file in file_path for ex_file in exclude_files):
            hf_hub_download(repo_id, filename=file_path, repo_type="dataset", revision=branch, local_dir=os.getcwd())

download_files(repo_id, branch, exclude_dirs, exclude_files, include_specific_files)

  8. Run MindEye2 inference and evaluation

    main

    The MindEye2 pipeline consists of several notebooks for different stages of the workflow:

    1. Inference: src/recon_inference.ipynb runs inference on a pretrained model, outputting tensors for reconstructions and predicted captions (saved as recons.pt).
    2. Refinement: src/enhanced_recon_inference.ipynb runs a refinement stage to produce higher-quality reconstructions (saved as enhancedrecons.pt).
    3. Evaluation: src/final_evaluations.ipynb visualizes the saved reconstructions and computes quantitative metrics.

    For batch processing, refer to the .slurm files in the repository for example scripts to submit these notebooks as Slurm jobs.

  9. Install Generative Models (sgm) with PyTorch 2.0

    main

    To install the sgm package for use with PyTorch 2.0, follow these steps. It is recommended to use python3.10 to avoid version conflicts.

    1. Clone the repository:
      git clone https://github.com/Stability-AI/generative-models.git
cd generative-models
    2. Create and activate a virtual environment:
      python3 -m venv .pt2
source .pt2/bin/activate
    3. Install the required PyTorch 2.0 dependencies:
      pip3 install -r requirements/pt2.txt
    4. Install the sgm package:
      pip3 install .
    5. (Optional) If you are performing training, install sdata:
      pip3 install -e git+https://github.com/Stability-AI/datapipelines.git@main#egg=sdata
    git clone https://github.com/Stability-AI/generative-models.git
cd generative-models
python3 -m venv .pt2
source .pt2/bin/activate
pip3 install -r requirements/pt2.txt
pip3 install .

  10. Detect invisible watermarks in generated images

    main

    Images generated with this codebase use the invisible-watermark library. You can detect these watermarks using scripts/demo/detect.py.

    Minimal Installation (Experimental)

    If you want to run detection without a full sgm installation, you can use a minimal environment:

    python -m venv .detect
source .detect/bin/activate
pip install "numpy>=1.17" "PyWavelets>=1.1.1" "opencv-python>=4.1.0.25"
pip install --no-deps invisible-watermark

    Usage

    Ensure your virtual environment is active, then run the detection script using one of the following patterns:

    • Single file: python scripts/demo/detect.py <filename>
    • Multiple files: python scripts/demo/detect.py <file1> <file2> ... <fileN>
    • Entire folder: python scripts/demo/detect.py <folder_name>/*
    # Example: testing a single file
python scripts/demo/detect.py my_image.png

  11. Build a distributable wheel using Hatch

    main

    The repository uses PEP 517 compliant packaging with Hatch. To build a distributable wheel package:

    1. Install hatch:
      pip install hatch
    2. Build the wheel:
   ```bash
hatch build -t wheel

    The built package will be located in the dist/ directory. You can install it using pip install dist/*.whl.

    Note: The package does not specify dependencies; you must install required packages (like PyTorch) manually based on your use case.

    pip install hatch
hatch build -t wheel

  12. Launch training for generative models

    main

    To start training, use main.py with the --base flag. You can provide multiple configuration files which will be merged from left to right (later configs overwrite values from earlier ones). This allows you to combine model, training, and data configurations separately or use a single monolithic config.

    Important Notes:

    • Dataset Formats: For large-scale datasets (e.g., ImageNet), the system expects data in the webdataset-format. Check configs for USER: comments to see which parts require adaptation for your specific dataset.
    • PyTorch Compatibility: Supports pytorch1.13 and pytorch2. However, autoencoder training (e.g., configs/example_training/autoencoder/kl-f4/imagenet-attnfree-logvar.yaml) only supports pytorch1.13.
    • Latent Models: For latent generative models (e.g., configs/example_training/imagenet-f8_cond.yaml), you must download the VAE checkpoint from Hugging Face and replace the CKPT_PATH placeholder in the config file.
    # Example: Running a single config for MNIST
python main.py --base configs/example_training/toy/mnist_cond.yaml

# Example: Merging multiple configs (left to right overwrite)
python main.py --base configs/<config1.yaml> configs/<config2.yaml>