from transformers import PretrainedConfig


class ViTv2Config(PretrainedConfig):
    model_type = "vitv2"

    def __init__(
        self,
        img_size=224,
        patch_size=16,
        embed_dim=384,
        depth=12,
        num_heads=6,
        mlp_ratio=4,
        num_register_tokens=0,
        init_values=None,
        **ignored_kwargs,
    ):
        super().__init__(**ignored_kwargs)

        self.depth = depth
        self.img_size = img_size
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.mlp_ratio = mlp_ratio
        self.patch_size = patch_size
        self.init_values = init_values
        self.num_register_tokens = num_register_tokens