from transformers import PretrainedConfig

class ProbUNetConfig(PretrainedConfig):
    model_type = "ProbUNet"
    def __init__(
            self,
            dim=2,
            in_channels=1,
            out_channels=1,
            num_feature_maps=24,
            latent_size=3,
            depth=5,
            latent_distribution="normal",
            no_outact_op=False,
            prob_injection_at="end",
            **kwargs):
        self.dim = dim
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.num_feature_maps = num_feature_maps
        self.latent_size = latent_size
        self.depth = depth
        self.latent_distribution = latent_distribution
        self.no_outact_op = no_outact_op
        self.prob_injection_at = prob_injection_at
        super().__init__(**kwargs)