stable-diffusion-webui

1
from modules import shared
2
from modules.sd_hijack_utils import CondFunc
3

4
has_ipex = False
5
try:
6
    import torch
7
    import intel_extension_for_pytorch as ipex # noqa: F401
8
    has_ipex = True
9
except Exception:
10
    pass
11

12

13
def check_for_xpu():
14
    return has_ipex and hasattr(torch, 'xpu') and torch.xpu.is_available()
15

16

17
def get_xpu_device_string():
18
    if shared.cmd_opts.device_id is not None:
19
        return f"xpu:{shared.cmd_opts.device_id}"
20
    return "xpu"
21

22

23
def torch_xpu_gc():
24
    with torch.xpu.device(get_xpu_device_string()):
25
        torch.xpu.empty_cache()
26

27

28
has_xpu = check_for_xpu()
29

30

31
# Arc GPU cannot allocate a single block larger than 4GB: https://github.com/intel/compute-runtime/issues/627
32
# Here we implement a slicing algorithm to split large batch size into smaller chunks,
33
# so that SDPA of each chunk wouldn't require any allocation larger than ARC_SINGLE_ALLOCATION_LIMIT.
34
# The heuristic limit (TOTAL_VRAM // 8) is tuned for Intel Arc A770 16G and Arc A750 8G,
35
# which is the best trade-off between VRAM usage and performance.
36
ARC_SINGLE_ALLOCATION_LIMIT = {}
37
orig_sdp_attn_func = torch.nn.functional.scaled_dot_product_attention
38
def torch_xpu_scaled_dot_product_attention(
39
    query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, *args, **kwargs
40
):
41
    # cast to same dtype first
42
    key = key.to(query.dtype)
43
    value = value.to(query.dtype)
44
    if attn_mask is not None and attn_mask.dtype != torch.bool:
45
        attn_mask = attn_mask.to(query.dtype)
46

47
    N = query.shape[:-2]  # Batch size
48
    L = query.size(-2)  # Target sequence length
49
    E = query.size(-1)  # Embedding dimension of the query and key
50
    S = key.size(-2)  # Source sequence length
51
    Ev = value.size(-1)  # Embedding dimension of the value
52

53
    total_batch_size = torch.numel(torch.empty(N))
54
    device_id = query.device.index
55
    if device_id not in ARC_SINGLE_ALLOCATION_LIMIT:
56
        ARC_SINGLE_ALLOCATION_LIMIT[device_id] = min(torch.xpu.get_device_properties(device_id).total_memory // 8, 4 * 1024 * 1024 * 1024)
57
    batch_size_limit = max(1, ARC_SINGLE_ALLOCATION_LIMIT[device_id] // (L * S * query.element_size()))
58

59
    if total_batch_size <= batch_size_limit:
60
        return orig_sdp_attn_func(
61
            query,
62
            key,
63
            value,
64
            attn_mask,
65
            dropout_p,
66
            is_causal,
67
            *args, **kwargs
68
        )
69

70
    query = torch.reshape(query, (-1, L, E))
71
    key = torch.reshape(key, (-1, S, E))
72
    value = torch.reshape(value, (-1, S, Ev))
73
    if attn_mask is not None:
74
        attn_mask = attn_mask.view(-1, L, S)
75
    chunk_count = (total_batch_size + batch_size_limit - 1) // batch_size_limit
76
    outputs = []
77
    for i in range(chunk_count):
78
        attn_mask_chunk = (
79
            None
80
            if attn_mask is None
81
            else attn_mask[i * batch_size_limit : (i + 1) * batch_size_limit, :, :]
82
        )
83
        chunk_output = orig_sdp_attn_func(
84
            query[i * batch_size_limit : (i + 1) * batch_size_limit, :, :],
85
            key[i * batch_size_limit : (i + 1) * batch_size_limit, :, :],
86
            value[i * batch_size_limit : (i + 1) * batch_size_limit, :, :],
87
            attn_mask_chunk,
88
            dropout_p,
89
            is_causal,
90
            *args, **kwargs
91
        )
92
        outputs.append(chunk_output)
93
    result = torch.cat(outputs, dim=0)
94
    return torch.reshape(result, (*N, L, Ev))
95

96

97
def is_xpu_device(device: str | torch.device = None):
98
    if device is None:
99
        return False
100
    if isinstance(device, str):
101
        return device.startswith("xpu")
102
    return device.type == "xpu"
103

104

105
if has_xpu:
106
    try:
107
        # torch.Generator supports "xpu" device since 2.1
108
        torch.Generator("xpu")
109
    except RuntimeError:
110
        # W/A for https://github.com/intel/intel-extension-for-pytorch/issues/452: torch.Generator API doesn't support XPU device (for torch < 2.1)
111
        CondFunc('torch.Generator',
112
            lambda orig_func, device=None: torch.xpu.Generator(device),
113
            lambda orig_func, device=None: is_xpu_device(device))
114

115
    # W/A for some OPs that could not handle different input dtypes
116
    CondFunc('torch.nn.functional.layer_norm',
117
        lambda orig_func, input, normalized_shape=None, weight=None, *args, **kwargs:
118
        orig_func(input.to(weight.data.dtype), normalized_shape, weight, *args, **kwargs),
119
        lambda orig_func, input, normalized_shape=None, weight=None, *args, **kwargs:
120
        weight is not None and input.dtype != weight.data.dtype)
121
    CondFunc('torch.nn.modules.GroupNorm.forward',
122
        lambda orig_func, self, input: orig_func(self, input.to(self.weight.data.dtype)),
123
        lambda orig_func, self, input: input.dtype != self.weight.data.dtype)
124
    CondFunc('torch.nn.modules.linear.Linear.forward',
125
        lambda orig_func, self, input: orig_func(self, input.to(self.weight.data.dtype)),
126
        lambda orig_func, self, input: input.dtype != self.weight.data.dtype)
127
    CondFunc('torch.nn.modules.conv.Conv2d.forward',
128
        lambda orig_func, self, input: orig_func(self, input.to(self.weight.data.dtype)),
129
        lambda orig_func, self, input: input.dtype != self.weight.data.dtype)
130
    CondFunc('torch.bmm',
131
        lambda orig_func, input, mat2, out=None: orig_func(input.to(mat2.dtype), mat2, out=out),
132
        lambda orig_func, input, mat2, out=None: input.dtype != mat2.dtype)
133
    CondFunc('torch.cat',
134
        lambda orig_func, tensors, dim=0, out=None: orig_func([t.to(tensors[0].dtype) for t in tensors], dim=dim, out=out),
135
        lambda orig_func, tensors, dim=0, out=None: not all(t.dtype == tensors[0].dtype for t in tensors))
136
    CondFunc('torch.nn.functional.scaled_dot_product_attention',
137
        lambda orig_func, *args, **kwargs: torch_xpu_scaled_dot_product_attention(*args, **kwargs),
138
        lambda orig_func, query, *args, **kwargs: query.is_xpu)
139