from transformers import PretrainedConfig
from typing import List

class WNet3DConfig(PretrainedConfig):
    model_type = "WNet"
    def __init__(
            self,
            in_ch=1, 
            out_ch=5, 
            init_features=64,
            **kwargs):
        self.in_ch = in_ch
        self.out_ch = out_ch
        self.init_features = init_features
        super().__init__(**kwargs)

class AttWNet3DConfig(PretrainedConfig):
    model_type = "AttWNet"
    def __init__(
            self,
            in_ch=1, 
            out_ch=5, 
            init_features=64,
            **kwargs):
        self.in_ch = in_ch
        self.out_ch = out_ch
        self.init_features = init_features
        super().__init__(**kwargs)

class WNetMSS3DConfig(PretrainedConfig):
    model_type = "WNetMSS"
    def __init__(
            self,
            in_ch=1, 
            out_ch=5, 
            init_features=64,
            **kwargs):
        self.in_ch = in_ch
        self.out_ch = out_ch
        self.init_features = init_features
        super().__init__(**kwargs)