from transformers import PreTrainedModel

from .config import RealESRGANConfig
from .rrdbnet import RRDBNet


class RealESRGANModel(PreTrainedModel):
    config_class = RealESRGANConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = RRDBNet(
            num_in_ch=config.num_in_ch,
            num_out_ch=config.num_out_ch,
            num_feat=config.num_feat,
            num_block=config.num_block,
            num_grow_ch=config.num_grow_ch,
            scale=config.scale,
        )

    def forward(self, tensor):
        return self.model.forward(tensor)