plutosss/ImageProcessing

1import sys
2import contextlib
3from functools import lru_cache
4
5import torch
6from modules import errors, shared, npu_specific
7
8if sys.platform == "darwin":
9    from modules import mac_specific
10
11if shared.cmd_opts.use_ipex:
12    from modules import xpu_specific
13
14
15def has_xpu() -> bool:
16    return shared.cmd_opts.use_ipex and xpu_specific.has_xpu
17
18
19def has_mps() -> bool:
20    if sys.platform != "darwin":
21        return False
22    else:
23        return mac_specific.has_mps
24
25
26def cuda_no_autocast(device_id=None) -> bool:
27    if device_id is None:
28        device_id = get_cuda_device_id()
29    return (
30        torch.cuda.get_device_capability(device_id) == (7, 5)
31        and torch.cuda.get_device_name(device_id).startswith("NVIDIA GeForce GTX 16")
32    )
33
34
35def get_cuda_device_id():
36    return (
37        int(shared.cmd_opts.device_id)
38        if shared.cmd_opts.device_id is not None and shared.cmd_opts.device_id.isdigit()
39        else 0
40    ) or torch.cuda.current_device()
41
42
43def get_cuda_device_string():
44    if shared.cmd_opts.device_id is not None:
45        return f"cuda:{shared.cmd_opts.device_id}"
46
47    return "cuda"
48
49
50def get_optimal_device_name():
51    if torch.cuda.is_available():
52        return get_cuda_device_string()
53
54    if has_mps():
55        return "mps"
56
57    if has_xpu():
58        return xpu_specific.get_xpu_device_string()
59
60    if npu_specific.has_npu:
61        return npu_specific.get_npu_device_string()
62
63    return "cpu"
64
65
66def get_optimal_device():
67    return torch.device(get_optimal_device_name())
68
69
70def get_device_for(task):
71    if task in shared.cmd_opts.use_cpu or "all" in shared.cmd_opts.use_cpu:
72        return cpu
73
74    return get_optimal_device()
75
76
77def torch_gc():
78
79    if torch.cuda.is_available():
80        with torch.cuda.device(get_cuda_device_string()):
81            torch.cuda.empty_cache()
82            torch.cuda.ipc_collect()
83
84    if has_mps():
85        mac_specific.torch_mps_gc()
86
87    if has_xpu():
88        xpu_specific.torch_xpu_gc()
89
90    if npu_specific.has_npu:
91        torch_npu_set_device()
92        npu_specific.torch_npu_gc()
93
94
95def torch_npu_set_device():
96    # Work around due to bug in torch_npu, revert me after fixed, @see https://gitee.com/ascend/pytorch/issues/I8KECW?from=project-issue
97    if npu_specific.has_npu:
98        torch.npu.set_device(0)
99
100
101def enable_tf32():
102    if torch.cuda.is_available():
103
104        # enabling benchmark option seems to enable a range of cards to do fp16 when they otherwise can't
105        # see https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/4407
106        if cuda_no_autocast():
107            torch.backends.cudnn.benchmark = True
108
109        torch.backends.cuda.matmul.allow_tf32 = True
110        torch.backends.cudnn.allow_tf32 = True
111
112
113errors.run(enable_tf32, "Enabling TF32")
114
115cpu: torch.device = torch.device("cpu")
116fp8: bool = False
117device: torch.device = None
118device_interrogate: torch.device = None
119device_gfpgan: torch.device = None
120device_esrgan: torch.device = None
121device_codeformer: torch.device = None
122dtype: torch.dtype = torch.float16
123dtype_vae: torch.dtype = torch.float16
124dtype_unet: torch.dtype = torch.float16
125dtype_inference: torch.dtype = torch.float16
126unet_needs_upcast = False
127
128
129def cond_cast_unet(input):
130    return input.to(dtype_unet) if unet_needs_upcast else input
131
132
133def cond_cast_float(input):
134    return input.float() if unet_needs_upcast else input
135
136
137nv_rng = None
138patch_module_list = [
139    torch.nn.Linear,
140    torch.nn.Conv2d,
141    torch.nn.MultiheadAttention,
142    torch.nn.GroupNorm,
143    torch.nn.LayerNorm,
144]
145
146
147def manual_cast_forward(target_dtype):
148    def forward_wrapper(self, *args, **kwargs):
149        if any(
150            isinstance(arg, torch.Tensor) and arg.dtype != target_dtype
151            for arg in args
152        ):
153            args = [arg.to(target_dtype) if isinstance(arg, torch.Tensor) else arg for arg in args]
154            kwargs = {k: v.to(target_dtype) if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()}
155
156        org_dtype = target_dtype
157        for param in self.parameters():
158            if param.dtype != target_dtype:
159                org_dtype = param.dtype
160                break
161
162        if org_dtype != target_dtype:
163            self.to(target_dtype)
164        result = self.org_forward(*args, **kwargs)
165        if org_dtype != target_dtype:
166            self.to(org_dtype)
167
168        if target_dtype != dtype_inference:
169            if isinstance(result, tuple):
170                result = tuple(
171                    i.to(dtype_inference)
172                    if isinstance(i, torch.Tensor)
173                    else i
174                    for i in result
175                )
176            elif isinstance(result, torch.Tensor):
177                result = result.to(dtype_inference)
178        return result
179    return forward_wrapper
180
181
182@contextlib.contextmanager
183def manual_cast(target_dtype):
184    applied = False
185    for module_type in patch_module_list:
186        if hasattr(module_type, "org_forward"):
187            continue
188        applied = True
189        org_forward = module_type.forward
190        if module_type == torch.nn.MultiheadAttention:
191            module_type.forward = manual_cast_forward(torch.float32)
192        else:
193            module_type.forward = manual_cast_forward(target_dtype)
194        module_type.org_forward = org_forward
195    try:
196        yield None
197    finally:
198        if applied:
199            for module_type in patch_module_list:
200                if hasattr(module_type, "org_forward"):
201                    module_type.forward = module_type.org_forward
202                    delattr(module_type, "org_forward")
203
204
205def autocast(disable=False):
206    if disable:
207        return contextlib.nullcontext()
208
209    if fp8 and device==cpu:
210        return torch.autocast("cpu", dtype=torch.bfloat16, enabled=True)
211
212    if fp8 and dtype_inference == torch.float32:
213        return manual_cast(dtype)
214
215    if dtype == torch.float32 or dtype_inference == torch.float32:
216        return contextlib.nullcontext()
217
218    if has_xpu() or has_mps() or cuda_no_autocast():
219        return manual_cast(dtype)
220
221    return torch.autocast("cuda")
222
223
224def without_autocast(disable=False):
225    return torch.autocast("cuda", enabled=False) if torch.is_autocast_enabled() and not disable else contextlib.nullcontext()
226
227
228class NansException(Exception):
229    pass
230
231
232def test_for_nans(x, where):
233    if shared.cmd_opts.disable_nan_check:
234        return
235
236    if not torch.all(torch.isnan(x)).item():
237        return
238
239    if where == "unet":
240        message = "A tensor with all NaNs was produced in Unet."
241
242        if not shared.cmd_opts.no_half:
243            message += " This could be either because there's not enough precision to represent the picture, or because your video card does not support half type. Try setting the \"Upcast cross attention layer to float32\" option in Settings > Stable Diffusion or using the --no-half commandline argument to fix this."
244
245    elif where == "vae":
246        message = "A tensor with all NaNs was produced in VAE."
247
248        if not shared.cmd_opts.no_half and not shared.cmd_opts.no_half_vae:
249            message += " This could be because there's not enough precision to represent the picture. Try adding --no-half-vae commandline argument to fix this."
250    else:
251        message = "A tensor with all NaNs was produced."
252
253    message += " Use --disable-nan-check commandline argument to disable this check."
254
255    raise NansException(message)
256
257
258@lru_cache
259def first_time_calculation():
260    """
261    just do any calculation with pytorch layers - the first time this is done it allocaltes about 700MB of memory and
262    spends about 2.7 seconds doing that, at least with NVidia.
263    """
264
265    x = torch.zeros((1, 1)).to(device, dtype)
266    linear = torch.nn.Linear(1, 1).to(device, dtype)
267    linear(x)
268
269    x = torch.zeros((1, 1, 3, 3)).to(device, dtype)
270    conv2d = torch.nn.Conv2d(1, 1, (3, 3)).to(device, dtype)
271    conv2d(x)
272