from transformers import PretrainedConfig


class FineViTConfig(PretrainedConfig):
    model_type = "finevit"

    def __init__(
        self,
        backbone_model_name: str = "facebook/dinov2-with-registers-base",
        teacher_encoder_type: str = "dinov2_with_registers",
        backbone_config: dict | None = None,
        num_latents: int = 64,
        latent_encoder_num_layers: int = 12,
        latent_encoder_head_dim: int = 64,
        latent_encoder_num_attention_heads: int = 12,
        latent_encoder_mlp_hidden_dim: int = 3072,
        patch_decoder_num_layers: int = 8,
        patch_decoder_head_dim: int = 64,
        patch_decoder_num_attention_heads: int = 12,
        patch_decoder_mlp_hidden_dim: int = 3072,
        dropout: float = 0.0,
        image_size: int = 224,
        train_encoder: bool = False,
        init_backbone_from_pretrained: bool = True,
        initializer_range: float = 0.02,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.backbone_model_name = backbone_model_name
        self.teacher_encoder_type = teacher_encoder_type
        self.backbone_config = backbone_config
        self.num_latents = num_latents
        self.latent_encoder_num_layers = latent_encoder_num_layers
        self.latent_encoder_head_dim = latent_encoder_head_dim
        self.latent_encoder_num_attention_heads = latent_encoder_num_attention_heads
        self.latent_encoder_mlp_hidden_dim = latent_encoder_mlp_hidden_dim
        self.patch_decoder_num_layers = patch_decoder_num_layers
        self.patch_decoder_head_dim = patch_decoder_head_dim
        self.patch_decoder_num_attention_heads = patch_decoder_num_attention_heads
        self.patch_decoder_mlp_hidden_dim = patch_decoder_mlp_hidden_dim
        self.dropout = dropout
        self.image_size = image_size
        self.train_encoder = train_encoder
        self.init_backbone_from_pretrained = init_backbone_from_pretrained
        self.initializer_range = initializer_range


FineViTConfig.register_for_auto_class("AutoConfig")