model_management.py 9.33 KB
Newer Older
1
2
3
import psutil
from enum import Enum
from cli_args import args
4

5
6
7
8
9
10
11
class VRAMState(Enum):
    CPU = 0
    NO_VRAM = 1
    LOW_VRAM = 2
    NORMAL_VRAM = 3
    HIGH_VRAM = 4
    MPS = 5
12

13
14
15
# Determine VRAM State
vram_state = VRAMState.NORMAL_VRAM
set_vram_to = VRAMState.NORMAL_VRAM
16

17
total_vram = 0
18
19
total_vram_available_mb = -1

20
accelerate_enabled = False
藍+85CD's avatar
藍+85CD committed
21
xpu_available = False
22

23
24
try:
    import torch
藍+85CD's avatar
藍+85CD committed
25
26
27
28
29
30
    try:
        import intel_extension_for_pytorch as ipex
        if torch.xpu.is_available():
            xpu_available = True
            total_vram = torch.xpu.get_device_properties(torch.xpu.current_device()).total_memory / (1024 * 1024)
    except:
31
        total_vram = torch.cuda.mem_get_info(torch.cuda.current_device())[1] / (1024 * 1024)
32
    total_ram = psutil.virtual_memory().total / (1024 * 1024)
33
    if not args.normalvram and not args.cpu:
34
35
        if total_vram <= 4096:
            print("Trying to enable lowvram mode because your GPU seems to have 4GB or less. If you don't want this use: --normalvram")
36
            set_vram_to = VRAMState.LOW_VRAM
comfyanonymous's avatar
comfyanonymous committed
37
        elif total_vram > total_ram * 1.1 and total_vram > 14336:
38
            print("Enabling highvram mode because your GPU has more vram than your computer has ram. If you don't want this use: --normalvram")
39
            vram_state = VRAMState.HIGH_VRAM
40
41
42
except:
    pass

43
44
45
46
47
try:
    OOM_EXCEPTION = torch.cuda.OutOfMemoryError
except:
    OOM_EXCEPTION = Exception

48
49
if args.disable_xformers:
    XFORMERS_IS_AVAILABLE = False
50
51
52
53
else:
    try:
        import xformers
        import xformers.ops
54
        XFORMERS_IS_AVAILABLE = True
55
    except:
56
        XFORMERS_IS_AVAILABLE = False
57

58
59
ENABLE_PYTORCH_ATTENTION = args.use_pytorch_cross_attention
if ENABLE_PYTORCH_ATTENTION:
60
61
62
    torch.backends.cuda.enable_math_sdp(True)
    torch.backends.cuda.enable_flash_sdp(True)
    torch.backends.cuda.enable_mem_efficient_sdp(True)
63
    XFORMERS_IS_AVAILABLE = False
64

65
66
67
68
69
70
if args.lowvram:
    set_vram_to = VRAMState.LOW_VRAM
elif args.novram:
    set_vram_to = VRAMState.NO_VRAM
elif args.highvram:
    vram_state = VRAMState.HIGH_VRAM
71

72

73
if set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM):
74
75
76
77
78
79
80
81
    try:
        import accelerate
        accelerate_enabled = True
        vram_state = set_vram_to
    except Exception as e:
        import traceback
        print(traceback.format_exc())
        print("ERROR: COULD NOT ENABLE LOW VRAM MODE.")
82
83

    total_vram_available_mb = (total_vram - 1024) // 2
84
    total_vram_available_mb = int(max(256, total_vram_available_mb))
85

86
87
try:
    if torch.backends.mps.is_available():
88
        vram_state = VRAMState.MPS
89
90
91
except:
    pass

92
93
if args.cpu:
    vram_state = VRAMState.CPU
94

95
print(f"Set vram state to: {vram_state.name}")
96

97
98

current_loaded_model = None
comfyanonymous's avatar
comfyanonymous committed
99
current_gpu_controlnets = []
100

101
102
103
model_accelerated = False


104
105
def unload_model():
    global current_loaded_model
106
    global model_accelerated
comfyanonymous's avatar
comfyanonymous committed
107
    global current_gpu_controlnets
108
109
    global vram_state

110
    if current_loaded_model is not None:
111
112
113
114
        if model_accelerated:
            accelerate.hooks.remove_hook_from_submodules(current_loaded_model.model)
            model_accelerated = False

115
        #never unload models from GPU on high vram
116
        if vram_state != VRAMState.HIGH_VRAM:
117
            current_loaded_model.model.cpu()
118
119
        current_loaded_model.unpatch_model()
        current_loaded_model = None
120

121
    if vram_state != VRAMState.HIGH_VRAM:
122
123
124
125
        if len(current_gpu_controlnets) > 0:
            for n in current_gpu_controlnets:
                n.cpu()
            current_gpu_controlnets = []
126
127
128
129


def load_model_gpu(model):
    global current_loaded_model
130
131
    global vram_state
    global model_accelerated
132
    global xpu_available
133

134
135
136
137
138
139
140
141
142
    if model is current_loaded_model:
        return
    unload_model()
    try:
        real_model = model.patch_model()
    except Exception as e:
        model.unpatch_model()
        raise e
    current_loaded_model = model
143
    if vram_state == VRAMState.CPU:
144
        pass
145
    elif vram_state == VRAMState.MPS:
Yurii Mazurevich's avatar
Yurii Mazurevich committed
146
147
148
        mps_device = torch.device("mps")
        real_model.to(mps_device)
        pass
149
    elif vram_state == VRAMState.NORMAL_VRAM or vram_state == VRAMState.HIGH_VRAM:
150
        model_accelerated = False
151
152
153
154
        if xpu_available:
            real_model.to("xpu")
        else:
            real_model.cuda()
155
    else:
156
        if vram_state == VRAMState.NO_VRAM:
157
            device_map = accelerate.infer_auto_device_map(real_model, max_memory={0: "256MiB", "cpu": "16GiB"})
158
        elif vram_state == VRAMState.LOW_VRAM:
159
            device_map = accelerate.infer_auto_device_map(real_model, max_memory={0: "{}MiB".format(total_vram_available_mb), "cpu": "16GiB"})
comfyanonymous's avatar
comfyanonymous committed
160

161
        accelerate.dispatch_model(real_model, device_map=device_map, main_device="xpu" if xpu_available else "cuda")
162
        model_accelerated = True
163
    return current_loaded_model
164

comfyanonymous's avatar
comfyanonymous committed
165
166
def load_controlnet_gpu(models):
    global current_gpu_controlnets
167
    global vram_state
168
    if vram_state == VRAMState.CPU:
169
        return
170

171
    if vram_state == VRAMState.LOW_VRAM or vram_state == VRAMState.NO_VRAM:
172
173
174
        #don't load controlnets like this if low vram because they will be loaded right before running and unloaded right after
        return

comfyanonymous's avatar
comfyanonymous committed
175
176
177
178
    for m in current_gpu_controlnets:
        if m not in models:
            m.cpu()

179
    device = get_torch_device()
comfyanonymous's avatar
comfyanonymous committed
180
181
    current_gpu_controlnets = []
    for m in models:
182
        current_gpu_controlnets.append(m.to(device))
comfyanonymous's avatar
comfyanonymous committed
183

184

185
186
def load_if_low_vram(model):
    global vram_state
187
    global xpu_available
188
    if vram_state == VRAMState.LOW_VRAM or vram_state == VRAMState.NO_VRAM:
189
190
191
192
        if xpu_available:
            return model.to("xpu")
        else:
            return model.cuda()
193
194
195
196
    return model

def unload_if_low_vram(model):
    global vram_state
197
    if vram_state == VRAMState.LOW_VRAM or vram_state == VRAMState.NO_VRAM:
198
199
200
        return model.cpu()
    return model

201
def get_torch_device():
202
    global xpu_available
203
    if vram_state == VRAMState.MPS:
Yurii Mazurevich's avatar
Yurii Mazurevich committed
204
        return torch.device("mps")
205
    if vram_state == VRAMState.CPU:
206
207
        return torch.device("cpu")
    else:
208
209
210
211
        if xpu_available:
            return torch.device("xpu")
        else:
            return torch.cuda.current_device()
212
213
214
215
216

def get_autocast_device(dev):
    if hasattr(dev, 'type'):
        return dev.type
    return "cuda"
217

218

219
def xformers_enabled():
220
    if vram_state == VRAMState.CPU:
221
        return False
222
    return XFORMERS_IS_AVAILABLE
223

224
225
226
227
228
229
230
231
232
233
234
235
236

def xformers_enabled_vae():
    enabled = xformers_enabled()
    if not enabled:
        return False
    try:
        #0.0.18 has a bug where Nan is returned when inputs are too big (1152x1920 res images and above)
        if xformers.version.__version__ == "0.0.18":
            return False
    except:
        pass
    return enabled

237
238
239
def pytorch_attention_enabled():
    return ENABLE_PYTORCH_ATTENTION

240
def get_free_memory(dev=None, torch_free_too=False):
241
    global xpu_available
242
    if dev is None:
243
        dev = get_torch_device()
244

Yurii Mazurevich's avatar
Yurii Mazurevich committed
245
    if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
246
247
248
        mem_free_total = psutil.virtual_memory().available
        mem_free_torch = mem_free_total
    else:
249
250
251
252
253
254
255
256
257
258
        if xpu_available:
            mem_free_total = torch.xpu.get_device_properties(dev).total_memory - torch.xpu.memory_allocated(dev)
            mem_free_torch = mem_free_total
        else:
            stats = torch.cuda.memory_stats(dev)
            mem_active = stats['active_bytes.all.current']
            mem_reserved = stats['reserved_bytes.all.current']
            mem_free_cuda, _ = torch.cuda.mem_get_info(dev)
            mem_free_torch = mem_reserved - mem_active
            mem_free_total = mem_free_cuda + mem_free_torch
259
260
261
262
263

    if torch_free_too:
        return (mem_free_total, mem_free_torch)
    else:
        return mem_free_total
264
265
266

def maximum_batch_area():
    global vram_state
267
    if vram_state == VRAMState.NO_VRAM:
268
269
270
271
272
        return 0

    memory_free = get_free_memory() / (1024 * 1024)
    area = ((memory_free - 1024) * 0.9) / (0.6)
    return int(max(area, 0))
273
274
275

def cpu_mode():
    global vram_state
276
    return vram_state == VRAMState.CPU
277

Yurii Mazurevich's avatar
Yurii Mazurevich committed
278
279
def mps_mode():
    global vram_state
280
    return vram_state == VRAMState.MPS
Yurii Mazurevich's avatar
Yurii Mazurevich committed
281

282
def should_use_fp16():
283
284
    global xpu_available
    if cpu_mode() or mps_mode() or xpu_available:
285
286
287
288
289
        return False #TODO ?

    if torch.cuda.is_bf16_supported():
        return True

comfyanonymous's avatar
comfyanonymous committed
290
    props = torch.cuda.get_device_properties("cuda")
291
292
293
294
    if props.major < 7:
        return False

    #FP32 is faster on those cards?
295
    nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600"]
296
297
298
299
300
301
    for x in nvidia_16_series:
        if x in props.name:
            return False

    return True

302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
#TODO: might be cleaner to put this somewhere else
import threading

class InterruptProcessingException(Exception):
    pass

interrupt_processing_mutex = threading.RLock()

interrupt_processing = False
def interrupt_current_processing(value=True):
    global interrupt_processing
    global interrupt_processing_mutex
    with interrupt_processing_mutex:
        interrupt_processing = value

def processing_interrupted():
    global interrupt_processing
    global interrupt_processing_mutex
    with interrupt_processing_mutex:
        return interrupt_processing

def throw_exception_if_processing_interrupted():
    global interrupt_processing
    global interrupt_processing_mutex
    with interrupt_processing_mutex:
        if interrupt_processing:
            interrupt_processing = False
            raise InterruptProcessingException()