from transformers import PretrainedConfig

class UpscalerConfig(PretrainedConfig):
    model_type = "upscaler"

    def __init__(
        self,
        scale: int = 2,
        in_channels: int = 3,
        width: int = 32,
        num_blocks: int = 3,
        feat1: int = 64,
        feat2: int = 32,
        use_refine: bool = False,
        **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.feat1 = int(feat1)
        self.feat2 = int(feat2)
        self.use_refine = bool(use_refine)