from transformers import PretrainedConfig
from typing import List


class AUNetConfig(PretrainedConfig):
    model_type = "s2l8hModel"
    def __init__(
            self,
            in_channels:int = 7,
            out_channels:int = 6,
            depth:int = 5,
            spatial_attention:str = 'None',
            growth_factor:int = 6,
            interp_mode:str = 'bicubic',
            up_mode:str = 'upsample',
            ca_layer:bool = False,
            **kwargs,
    ):
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.depth = depth
        self.spatial_attention = spatial_attention
        self.growth_factor = growth_factor
        self.interp_mode = interp_mode
        self.up_mode = up_mode
        self.ca_layer = ca_layer

        super().__init__(**kwargs)