model_management.py 12.9 KB
Newer Older
1
2
import psutil
from enum import Enum
comfyanonymous's avatar
comfyanonymous committed
3
from comfy.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
lowvram_available = True
藍+85CD's avatar
藍+85CD committed
20
xpu_available = False
21

22
directml_enabled = False
23
if args.directml is not None:
24
25
    import torch_directml
    directml_enabled = True
26
27
28
29
30
31
    device_index = args.directml
    if device_index < 0:
        directml_device = torch_directml.device()
    else:
        directml_device = torch_directml.device(device_index)
    print("Using directml with device:", torch_directml.device_name(device_index))
32
    # torch_directml.disable_tiled_resources(True)
33
    lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default.
34

35
36
try:
    import torch
37
    if directml_enabled:
38
        pass #TODO
39
40
41
42
43
44
45
46
    else:
        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:
            total_vram = torch.cuda.mem_get_info(torch.cuda.current_device())[1] / (1024 * 1024)
47
    total_ram = psutil.virtual_memory().total / (1024 * 1024)
48
    if not args.normalvram and not args.cpu:
49
        if lowvram_available and total_vram <= 4096:
50
            print("Trying to enable lowvram mode because your GPU seems to have 4GB or less. If you don't want this use: --normalvram")
51
            set_vram_to = VRAMState.LOW_VRAM
comfyanonymous's avatar
comfyanonymous committed
52
        elif total_vram > total_ram * 1.1 and total_vram > 14336:
53
            print("Enabling highvram mode because your GPU has more vram than your computer has ram. If you don't want this use: --normalvram")
54
            vram_state = VRAMState.HIGH_VRAM
55
56
57
except:
    pass

58
59
60
61
62
try:
    OOM_EXCEPTION = torch.cuda.OutOfMemoryError
except:
    OOM_EXCEPTION = Exception

63
64
XFORMERS_VERSION = ""
XFORMERS_ENABLED_VAE = True
65
66
if args.disable_xformers:
    XFORMERS_IS_AVAILABLE = False
67
68
69
70
else:
    try:
        import xformers
        import xformers.ops
71
        XFORMERS_IS_AVAILABLE = True
72
73
74
75
76
77
78
79
80
81
82
        try:
            XFORMERS_VERSION = xformers.version.__version__
            print("xformers version:", XFORMERS_VERSION)
            if XFORMERS_VERSION.startswith("0.0.18"):
                print()
                print("WARNING: This version of xformers has a major bug where you will get black images when generating high resolution images.")
                print("Please downgrade or upgrade xformers to a different version.")
                print()
                XFORMERS_ENABLED_VAE = False
        except:
            pass
83
    except:
84
        XFORMERS_IS_AVAILABLE = False
85

86
87
ENABLE_PYTORCH_ATTENTION = args.use_pytorch_cross_attention
if ENABLE_PYTORCH_ATTENTION:
88
89
90
    torch.backends.cuda.enable_math_sdp(True)
    torch.backends.cuda.enable_flash_sdp(True)
    torch.backends.cuda.enable_mem_efficient_sdp(True)
91
    XFORMERS_IS_AVAILABLE = False
92

93
94
if args.lowvram:
    set_vram_to = VRAMState.LOW_VRAM
95
    lowvram_available = True
96
97
98
99
elif args.novram:
    set_vram_to = VRAMState.NO_VRAM
elif args.highvram:
    vram_state = VRAMState.HIGH_VRAM
100

101
102
103
104
105
FORCE_FP32 = False
if args.force_fp32:
    print("Forcing FP32, if this improves things please report it.")
    FORCE_FP32 = True

106

107
108

if lowvram_available:
109
110
    try:
        import accelerate
111
112
        if set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM):
            vram_state = set_vram_to
113
114
115
    except Exception as e:
        import traceback
        print(traceback.format_exc())
116
117
        print("ERROR: LOW VRAM MODE NEEDS accelerate.")
        lowvram_available = False
118

119

120
121
try:
    if torch.backends.mps.is_available():
122
        vram_state = VRAMState.MPS
123
124
125
except:
    pass

126
127
if args.cpu:
    vram_state = VRAMState.CPU
128

129
print(f"Set vram state to: {vram_state.name}")
130

131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
def get_torch_device():
    global xpu_available
    global directml_enabled
    if directml_enabled:
        global directml_device
        return directml_device
    if vram_state == VRAMState.MPS:
        return torch.device("mps")
    if vram_state == VRAMState.CPU:
        return torch.device("cpu")
    else:
        if xpu_available:
            return torch.device("xpu")
        else:
            return torch.cuda.current_device()

def get_torch_device_name(device):
    if hasattr(device, 'type'):
        return "{}".format(device.type)
    return "CUDA {}: {}".format(device, torch.cuda.get_device_name(device))

try:
    print("Using device:", get_torch_device_name(get_torch_device()))
except:
    print("Could not pick default device.")

157
158

current_loaded_model = None
comfyanonymous's avatar
comfyanonymous committed
159
current_gpu_controlnets = []
160

161
162
163
model_accelerated = False


164
165
def unload_model():
    global current_loaded_model
166
    global model_accelerated
comfyanonymous's avatar
comfyanonymous committed
167
    global current_gpu_controlnets
168
169
    global vram_state

170
    if current_loaded_model is not None:
171
172
173
174
        if model_accelerated:
            accelerate.hooks.remove_hook_from_submodules(current_loaded_model.model)
            model_accelerated = False

175
        #never unload models from GPU on high vram
176
        if vram_state != VRAMState.HIGH_VRAM:
177
            current_loaded_model.model.cpu()
178
            current_loaded_model.model_patches_to("cpu")
179
180
        current_loaded_model.unpatch_model()
        current_loaded_model = None
181

182
    if vram_state != VRAMState.HIGH_VRAM:
183
184
185
186
        if len(current_gpu_controlnets) > 0:
            for n in current_gpu_controlnets:
                n.cpu()
            current_gpu_controlnets = []
187
188
189
190


def load_model_gpu(model):
    global current_loaded_model
191
192
193
    global vram_state
    global model_accelerated

194
195
196
197
198
199
200
201
    if model is current_loaded_model:
        return
    unload_model()
    try:
        real_model = model.patch_model()
    except Exception as e:
        model.unpatch_model()
        raise e
202

203
204
205
206
207
208
209
    torch_dev = get_torch_device()
    model.model_patches_to(torch_dev)

    vram_set_state = vram_state
    if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM):
        model_size = model.model_size()
        current_free_mem = get_free_memory(torch_dev)
210
        lowvram_model_memory = int(max(256 * (1024 * 1024), (current_free_mem - 1024 * (1024 * 1024)) / 1.3 ))
211
212
213
        if model_size > (current_free_mem - (512 * 1024 * 1024)): #only switch to lowvram if really necessary
            vram_set_state = VRAMState.LOW_VRAM

214
    current_loaded_model = model
215
216

    if vram_set_state == VRAMState.CPU:
217
        pass
218
    elif vram_set_state == VRAMState.MPS:
Yurii Mazurevich's avatar
Yurii Mazurevich committed
219
220
221
        mps_device = torch.device("mps")
        real_model.to(mps_device)
        pass
222
    elif vram_set_state == VRAMState.NORMAL_VRAM or vram_set_state == VRAMState.HIGH_VRAM:
223
        model_accelerated = False
comfyanonymous's avatar
comfyanonymous committed
224
        real_model.to(get_torch_device())
225
    else:
226
        if vram_set_state == VRAMState.NO_VRAM:
227
            device_map = accelerate.infer_auto_device_map(real_model, max_memory={0: "256MiB", "cpu": "16GiB"})
228
229
        elif vram_set_state == VRAMState.LOW_VRAM:
            device_map = accelerate.infer_auto_device_map(real_model, max_memory={0: "{}MiB".format(lowvram_model_memory // (1024 * 1024)), "cpu": "16GiB"})
comfyanonymous's avatar
comfyanonymous committed
230

comfyanonymous's avatar
comfyanonymous committed
231
        accelerate.dispatch_model(real_model, device_map=device_map, main_device=get_torch_device())
232
        model_accelerated = True
233
    return current_loaded_model
234

235
def load_controlnet_gpu(control_models):
comfyanonymous's avatar
comfyanonymous committed
236
    global current_gpu_controlnets
237
    global vram_state
238
    if vram_state == VRAMState.CPU:
239
        return
240

241
    if vram_state == VRAMState.LOW_VRAM or vram_state == VRAMState.NO_VRAM:
242
243
244
        for m in control_models:
            if hasattr(m, 'set_lowvram'):
                m.set_lowvram(True)
245
246
247
        #don't load controlnets like this if low vram because they will be loaded right before running and unloaded right after
        return

248
249
250
251
    models = []
    for m in control_models:
        models += m.get_models()

comfyanonymous's avatar
comfyanonymous committed
252
253
254
255
    for m in current_gpu_controlnets:
        if m not in models:
            m.cpu()

256
    device = get_torch_device()
comfyanonymous's avatar
comfyanonymous committed
257
258
    current_gpu_controlnets = []
    for m in models:
259
        current_gpu_controlnets.append(m.to(device))
comfyanonymous's avatar
comfyanonymous committed
260

261

262
263
def load_if_low_vram(model):
    global vram_state
264
    if vram_state == VRAMState.LOW_VRAM or vram_state == VRAMState.NO_VRAM:
comfyanonymous's avatar
comfyanonymous committed
265
        return model.to(get_torch_device())
266
267
268
269
    return model

def unload_if_low_vram(model):
    global vram_state
270
    if vram_state == VRAMState.LOW_VRAM or vram_state == VRAMState.NO_VRAM:
271
272
273
        return model.cpu()
    return model

274
275
276
277
def get_autocast_device(dev):
    if hasattr(dev, 'type'):
        return dev.type
    return "cuda"
278

279

280
def xformers_enabled():
281
282
    global xpu_available
    global directml_enabled
283
    if vram_state == VRAMState.CPU:
284
        return False
285
286
287
288
    if xpu_available:
        return False
    if directml_enabled:
        return False
289
    return XFORMERS_IS_AVAILABLE
290

291
292
293
294
295

def xformers_enabled_vae():
    enabled = xformers_enabled()
    if not enabled:
        return False
296
297

    return XFORMERS_ENABLED_VAE
298

299
def pytorch_attention_enabled():
300
    global ENABLE_PYTORCH_ATTENTION
301
302
    return ENABLE_PYTORCH_ATTENTION

303
304
305
306
307
308
309
310
def pytorch_attention_flash_attention():
    global ENABLE_PYTORCH_ATTENTION
    if ENABLE_PYTORCH_ATTENTION:
        #TODO: more reliable way of checking for flash attention?
        if torch.version.cuda: #pytorch flash attention only works on Nvidia
            return True
    return False

311
def get_free_memory(dev=None, torch_free_too=False):
312
    global xpu_available
313
    global directml_enabled
314
    if dev is None:
315
        dev = get_torch_device()
316

Yurii Mazurevich's avatar
Yurii Mazurevich committed
317
    if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
318
319
320
        mem_free_total = psutil.virtual_memory().available
        mem_free_torch = mem_free_total
    else:
321
322
323
324
        if directml_enabled:
            mem_free_total = 1024 * 1024 * 1024 #TODO
            mem_free_torch = mem_free_total
        elif xpu_available:
325
326
327
328
329
330
331
332
333
            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
334
335
336
337
338

    if torch_free_too:
        return (mem_free_total, mem_free_torch)
    else:
        return mem_free_total
339
340
341

def maximum_batch_area():
    global vram_state
342
    if vram_state == VRAMState.NO_VRAM:
343
344
345
        return 0

    memory_free = get_free_memory() / (1024 * 1024)
346
    if xformers_enabled() or pytorch_attention_flash_attention():
347
        #TODO: this needs to be tweaked
348
        area = 20 * memory_free
349
350
351
    else:
        #TODO: this formula is because AMD sucks and has memory management issues which might be fixed in the future
        area = ((memory_free - 1024) * 0.9) / (0.6)
352
    return int(max(area, 0))
353
354
355

def cpu_mode():
    global vram_state
356
    return vram_state == VRAMState.CPU
357

Yurii Mazurevich's avatar
Yurii Mazurevich committed
358
359
def mps_mode():
    global vram_state
360
    return vram_state == VRAMState.MPS
Yurii Mazurevich's avatar
Yurii Mazurevich committed
361

362
def should_use_fp16():
363
    global xpu_available
364
365
    global directml_enabled

366
367
368
    if FORCE_FP32:
        return False

369
370
371
    if directml_enabled:
        return False

372
    if cpu_mode() or mps_mode() or xpu_available:
373
374
375
376
377
        return False #TODO ?

    if torch.cuda.is_bf16_supported():
        return True

comfyanonymous's avatar
comfyanonymous committed
378
    props = torch.cuda.get_device_properties("cuda")
379
380
381
382
    if props.major < 7:
        return False

    #FP32 is faster on those cards?
383
    nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600"]
384
385
386
387
388
389
    for x in nvidia_16_series:
        if x in props.name:
            return False

    return True

390
391
def soft_empty_cache():
    global xpu_available
comfyanonymous's avatar
comfyanonymous committed
392
393
394
395
    global vram_state
    if vram_state == VRAMState.MPS:
        torch.mps.empty_cache()
    elif xpu_available:
396
397
398
399
400
401
        torch.xpu.empty_cache()
    elif torch.cuda.is_available():
        if torch.version.cuda: #This seems to make things worse on ROCm so I only do it for cuda
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()

402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
#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()