stable-diffusion-webui

1
import importlib
2
import torch
3

4
from modules import shared
5

6

7
def check_for_npu():
8
    if importlib.util.find_spec("torch_npu") is None:
9
        return False
10
    import torch_npu
11

12
    try:
13
        # Will raise a RuntimeError if no NPU is found
14
        _ = torch_npu.npu.device_count()
15
        return torch.npu.is_available()
16
    except RuntimeError:
17
        return False
18

19

20
def get_npu_device_string():
21
    if shared.cmd_opts.device_id is not None:
22
        return f"npu:{shared.cmd_opts.device_id}"
23
    return "npu:0"
24

25

26
def torch_npu_gc():
27
    with torch.npu.device(get_npu_device_string()):
28
        torch.npu.empty_cache()
29

30

31
has_npu = check_for_npu()
32