from transformers import PretrainedConfig

class CustomViTNanoConfig(PretrainedConfig):
    model_type = "custom_vit_nano"

    def __init__(
        self,
        image_size=224,
        patch_size=16,
        in_chans=3,
        num_classes=1000,
        embed_dim=224,
        depth=6,
        num_heads=4,
        mlp_hidden_dim=608,
        stem_channels=(32, 64, 128),
        dropout=0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.image_size = image_size
        self.patch_size = patch_size
        self.in_chans = in_chans
        self.num_classes = num_classes
        self.embed_dim = embed_dim
        self.depth = depth
        self.num_heads = num_heads
        self.mlp_hidden_dim = mlp_hidden_dim
        self.stem_channels = stem_channels
        self.dropout = dropout

CustomViTNanoConfig.register_for_auto_class()