stable-diffusion-webui

1
import os
2
from abc import abstractmethod
3

4
import PIL
5
from PIL import Image
6

7
import modules.shared
8
from modules import modelloader, shared
9

10
LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
11
NEAREST = (Image.Resampling.NEAREST if hasattr(Image, 'Resampling') else Image.NEAREST)
12

13

14
class Upscaler:
15
    name = None
16
    model_path = None
17
    model_name = None
18
    model_url = None
19
    enable = True
20
    filter = None
21
    model = None
22
    user_path = None
23
    scalers: []
24
    tile = True
25

26
    def __init__(self, create_dirs=False):
27
        self.mod_pad_h = None
28
        self.tile_size = modules.shared.opts.ESRGAN_tile
29
        self.tile_pad = modules.shared.opts.ESRGAN_tile_overlap
30
        self.device = modules.shared.device
31
        self.img = None
32
        self.output = None
33
        self.scale = 1
34
        self.half = not modules.shared.cmd_opts.no_half
35
        self.pre_pad = 0
36
        self.mod_scale = None
37
        self.model_download_path = None
38

39
        if self.model_path is None and self.name:
40
            self.model_path = os.path.join(shared.models_path, self.name)
41
        if self.model_path and create_dirs:
42
            os.makedirs(self.model_path, exist_ok=True)
43

44
        try:
45
            import cv2  # noqa: F401
46
            self.can_tile = True
47
        except Exception:
48
            pass
49

50
    @abstractmethod
51
    def do_upscale(self, img: PIL.Image, selected_model: str):
52
        return img
53

54
    def upscale(self, img: PIL.Image, scale, selected_model: str = None):
55
        self.scale = scale
56
        dest_w = int((img.width * scale) // 8 * 8)
57
        dest_h = int((img.height * scale) // 8 * 8)
58

59
        for _ in range(3):
60
            if img.width >= dest_w and img.height >= dest_h:
61
                break
62

63
            shape = (img.width, img.height)
64

65
            img = self.do_upscale(img, selected_model)
66

67
            if shape == (img.width, img.height):
68
                break
69

70
        if img.width != dest_w or img.height != dest_h:
71
            img = img.resize((int(dest_w), int(dest_h)), resample=LANCZOS)
72

73
        return img
74

75
    @abstractmethod
76
    def load_model(self, path: str):
77
        pass
78

79
    def find_models(self, ext_filter=None) -> list:
80
        return modelloader.load_models(model_path=self.model_path, model_url=self.model_url, command_path=self.user_path, ext_filter=ext_filter)
81

82
    def update_status(self, prompt):
83
        print(f"\nextras: {prompt}", file=shared.progress_print_out)
84

85

86
class UpscalerData:
87
    name = None
88
    data_path = None
89
    scale: int = 4
90
    scaler: Upscaler = None
91
    model: None
92

93
    def __init__(self, name: str, path: str, upscaler: Upscaler = None, scale: int = 4, model=None):
94
        self.name = name
95
        self.data_path = path
96
        self.local_data_path = path
97
        self.scaler = upscaler
98
        self.scale = scale
99
        self.model = model
100

101
    def __repr__(self):
102
        return f"<UpscalerData name={self.name} path={self.data_path} scale={self.scale}>"
103

104

105
class UpscalerNone(Upscaler):
106
    name = "None"
107
    scalers = []
108

109
    def load_model(self, path):
110
        pass
111

112
    def do_upscale(self, img, selected_model=None):
113
        return img
114

115
    def __init__(self, dirname=None):
116
        super().__init__(False)
117
        self.scalers = [UpscalerData("None", None, self)]
118

119

120
class UpscalerLanczos(Upscaler):
121
    scalers = []
122

123
    def do_upscale(self, img, selected_model=None):
124
        return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=LANCZOS)
125

126
    def load_model(self, _):
127
        pass
128

129
    def __init__(self, dirname=None):
130
        super().__init__(False)
131
        self.name = "Lanczos"
132
        self.scalers = [UpscalerData("Lanczos", None, self)]
133

134

135
class UpscalerNearest(Upscaler):
136
    scalers = []
137

138
    def do_upscale(self, img, selected_model=None):
139
        return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=NEAREST)
140

141
    def load_model(self, _):
142
        pass
143

144
    def __init__(self, dirname=None):
145
        super().__init__(False)
146
        self.name = "Nearest"
147
        self.scalers = [UpscalerData("Nearest", None, self)]
148