from transformers import PretrainedConfig


class UpscalerConfig(PretrainedConfig):
    model_type = "upscaler"

    def __init__(
        self,
        scale: int = 4,
        in_channels: int = 3,
        width: int = 32,
        num_blocks: int = 3,
        espcn_channels: int = 64,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.scale = int(scale)
        self.in_channels = int(in_channels)
        self.width = int(width)
        self.num_blocks = int(num_blocks)
        self.espcn_channels = int(espcn_channels)