from typing import List

import torch

import models.ultrasharp.arch as arch
from models.ultrasharp.util import infer_params, upscale


class Ultrasharp:
    def __init__(self, model_path, tile_pad=0, tile=0):
        self.filename = model_path
        self.tile_pad = tile_pad
        self.tile = tile

    def enhance(self, img, outscale=4):
        state_dict = torch.load(self.filename, map_location="cpu")

        in_nc, out_nc, nf, nb, plus, mscale = infer_params(state_dict)

        model = arch.RRDBNet(
            in_nc=in_nc, out_nc=out_nc, nf=nf, nb=nb, upscale=mscale, plus=plus
        )
        model.load_state_dict(state_dict)
        model.eval()

        model.to("cuda")

        img = upscale(model, img, self.tile_pad, self.tile)
        return img, None