vjepa2-encoder-basic / configuration_vjepa2_fmri_encoder.py
dawidtang's picture
Fix HF config and extractor compatibility
77cfc6c verified
Raw
History Blame Contribute Delete
1.26 kB
"""Transformers config for the V-JEPA2 fMRI encoder."""
from __future__ import annotations
from transformers import PretrainedConfig
class VJEPA2FMRIEncoderConfig(PretrainedConfig):
model_type = "vjepa2_fmri_encoder"
def __init__(
self,
checkpoint_filename: str = "vjepa2_offline_encoder.pth",
output_dim: int = 20484,
input_duration_seconds: float = 3.0,
input_format: str = "video_tensor_b_t_c_h_w",
output_description: str = "z_scored_fmri_betas_no_time_dimension",
backbone_filename: str = "vitl.pt",
vjepa_size: str = "large",
load_vjepa: bool = True,
image_size: int = 224,
normalize_input: bool = True,
**kwargs,
) -> None:
super().__init__(**kwargs)
self.checkpoint_filename = checkpoint_filename
self.output_dim = int(output_dim)
self.input_duration_seconds = float(input_duration_seconds)
self.input_format = input_format
self.output_description = output_description
self.backbone_filename = backbone_filename
self.vjepa_size = vjepa_size
self.load_vjepa = bool(load_vjepa)
self.image_size = int(image_size)
self.normalize_input = bool(normalize_input)