stable-diffusion-webui

1
import os
2
import sys
3

4
from modules import modelloader, devices
5
from modules.shared import opts
6
from modules.upscaler import Upscaler, UpscalerData
7
from modules.upscaler_utils import upscale_with_model
8

9

10
class UpscalerHAT(Upscaler):
11
    def __init__(self, dirname):
12
        self.name = "HAT"
13
        self.scalers = []
14
        self.user_path = dirname
15
        super().__init__()
16
        for file in self.find_models(ext_filter=[".pt", ".pth"]):
17
            name = modelloader.friendly_name(file)
18
            scale = 4  # TODO: scale might not be 4, but we can't know without loading the model
19
            scaler_data = UpscalerData(name, file, upscaler=self, scale=scale)
20
            self.scalers.append(scaler_data)
21

22
    def do_upscale(self, img, selected_model):
23
        try:
24
            model = self.load_model(selected_model)
25
        except Exception as e:
26
            print(f"Unable to load HAT model {selected_model}: {e}", file=sys.stderr)
27
            return img
28
        model.to(devices.device_esrgan)  # TODO: should probably be device_hat
29
        return upscale_with_model(
30
            model,
31
            img,
32
            tile_size=opts.ESRGAN_tile,  # TODO: should probably be HAT_tile
33
            tile_overlap=opts.ESRGAN_tile_overlap,  # TODO: should probably be HAT_tile_overlap
34
        )
35

36
    def load_model(self, path: str):
37
        if not os.path.isfile(path):
38
            raise FileNotFoundError(f"Model file {path} not found")
39
        return modelloader.load_spandrel_model(
40
            path,
41
            device=devices.device_esrgan,  # TODO: should probably be device_hat
42
            expected_architecture='HAT',
43
        )
44