model_management.py 10.8 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
25
26
27
28
29
directml_enabled = False
if args.directml:
    import torch_directml
    print("Using directml")
    directml_enabled = True
    # torch_directml.disable_tiled_resources(True)

30
31
try:
    import torch
藍+85CD's avatar
藍+85CD committed
32
33
34
35
36
37
    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:
38
        total_vram = torch.cuda.mem_get_info(torch.cuda.current_device())[1] / (1024 * 1024)
39
    total_ram = psutil.virtual_memory().total / (1024 * 1024)
40
    if not args.normalvram and not args.cpu:
41
42
        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")
43
            set_vram_to = VRAMState.LOW_VRAM
comfyanonymous's avatar
comfyanonymous committed
44
        elif total_vram > total_ram * 1.1 and total_vram > 14336:
45
            print("Enabling highvram mode because your GPU has more vram than your computer has ram. If you don't want this use: --normalvram")
46
            vram_state = VRAMState.HIGH_VRAM
47
48
49
except:
    pass

50
51
52
53
54
try:
    OOM_EXCEPTION = torch.cuda.OutOfMemoryError
except:
    OOM_EXCEPTION = Exception

55
56
XFORMERS_VERSION = ""
XFORMERS_ENABLED_VAE = True
57
58
if args.disable_xformers:
    XFORMERS_IS_AVAILABLE = False
59
60
61
62
else:
    try:
        import xformers
        import xformers.ops
63
        XFORMERS_IS_AVAILABLE = True
64
65
66
67
68
69
70
71
72
73
74
        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
75
    except:
76
        XFORMERS_IS_AVAILABLE = False
77

78
79
ENABLE_PYTORCH_ATTENTION = args.use_pytorch_cross_attention
if ENABLE_PYTORCH_ATTENTION:
80
81
82
    torch.backends.cuda.enable_math_sdp(True)
    torch.backends.cuda.enable_flash_sdp(True)
    torch.backends.cuda.enable_mem_efficient_sdp(True)
83
    XFORMERS_IS_AVAILABLE = False
84

85
86
87
88
89
90
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
91

92
93
94
95
96
FORCE_FP32 = False
if args.force_fp32:
    print("Forcing FP32, if this improves things please report it.")
    FORCE_FP32 = True

97

98
if set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM):
99
100
101
102
103
104
105
106
    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.")
107
108

    total_vram_available_mb = (total_vram - 1024) // 2
109
    total_vram_available_mb = int(max(256, total_vram_available_mb))
110

111
112
try:
    if torch.backends.mps.is_available():
113
        vram_state = VRAMState.MPS
114
115
116
except:
    pass

117
118
if args.cpu:
    vram_state = VRAMState.CPU
119

120
print(f"Set vram state to: {vram_state.name}")
121

122
123

current_loaded_model = None
comfyanonymous's avatar
comfyanonymous committed
124
current_gpu_controlnets = []
125

126
127
128
model_accelerated = False


129
130
def unload_model():
    global current_loaded_model
131
    global model_accelerated
comfyanonymous's avatar
comfyanonymous committed
132
    global current_gpu_controlnets
133
134
    global vram_state

135
    if current_loaded_model is not None:
136
137
138
139
        if model_accelerated:
            accelerate.hooks.remove_hook_from_submodules(current_loaded_model.model)
            model_accelerated = False

140
        #never unload models from GPU on high vram
141
        if vram_state != VRAMState.HIGH_VRAM:
142
            current_loaded_model.model.cpu()
143
            current_loaded_model.model_patches_to("cpu")
144
145
        current_loaded_model.unpatch_model()
        current_loaded_model = None
146

147
    if vram_state != VRAMState.HIGH_VRAM:
148
149
150
151
        if len(current_gpu_controlnets) > 0:
            for n in current_gpu_controlnets:
                n.cpu()
            current_gpu_controlnets = []
152
153
154
155


def load_model_gpu(model):
    global current_loaded_model
156
157
158
    global vram_state
    global model_accelerated

159
160
161
162
163
164
165
166
    if model is current_loaded_model:
        return
    unload_model()
    try:
        real_model = model.patch_model()
    except Exception as e:
        model.unpatch_model()
        raise e
167
168

    model.model_patches_to(get_torch_device())
169
    current_loaded_model = model
170
    if vram_state == VRAMState.CPU:
171
        pass
172
    elif vram_state == VRAMState.MPS:
Yurii Mazurevich's avatar
Yurii Mazurevich committed
173
174
175
        mps_device = torch.device("mps")
        real_model.to(mps_device)
        pass
176
    elif vram_state == VRAMState.NORMAL_VRAM or vram_state == VRAMState.HIGH_VRAM:
177
        model_accelerated = False
comfyanonymous's avatar
comfyanonymous committed
178
        real_model.to(get_torch_device())
179
    else:
180
        if vram_state == VRAMState.NO_VRAM:
181
            device_map = accelerate.infer_auto_device_map(real_model, max_memory={0: "256MiB", "cpu": "16GiB"})
182
        elif vram_state == VRAMState.LOW_VRAM:
183
            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
184

comfyanonymous's avatar
comfyanonymous committed
185
        accelerate.dispatch_model(real_model, device_map=device_map, main_device=get_torch_device())
186
        model_accelerated = True
187
    return current_loaded_model
188

189
def load_controlnet_gpu(control_models):
comfyanonymous's avatar
comfyanonymous committed
190
    global current_gpu_controlnets
191
    global vram_state
192
    if vram_state == VRAMState.CPU:
193
        return
194

195
    if vram_state == VRAMState.LOW_VRAM or vram_state == VRAMState.NO_VRAM:
196
197
198
        #don't load controlnets like this if low vram because they will be loaded right before running and unloaded right after
        return

199
200
201
202
    models = []
    for m in control_models:
        models += m.get_models()

comfyanonymous's avatar
comfyanonymous committed
203
204
205
206
    for m in current_gpu_controlnets:
        if m not in models:
            m.cpu()

207
    device = get_torch_device()
comfyanonymous's avatar
comfyanonymous committed
208
209
    current_gpu_controlnets = []
    for m in models:
210
        current_gpu_controlnets.append(m.to(device))
comfyanonymous's avatar
comfyanonymous committed
211

212

213
214
def load_if_low_vram(model):
    global vram_state
215
    if vram_state == VRAMState.LOW_VRAM or vram_state == VRAMState.NO_VRAM:
comfyanonymous's avatar
comfyanonymous committed
216
        return model.to(get_torch_device())
217
218
219
220
    return model

def unload_if_low_vram(model):
    global vram_state
221
    if vram_state == VRAMState.LOW_VRAM or vram_state == VRAMState.NO_VRAM:
222
223
224
        return model.cpu()
    return model

225
def get_torch_device():
226
    global xpu_available
227
228
229
    global directml_enabled
    if directml_enabled:
        return torch_directml.device()
230
    if vram_state == VRAMState.MPS:
Yurii Mazurevich's avatar
Yurii Mazurevich committed
231
        return torch.device("mps")
232
    if vram_state == VRAMState.CPU:
233
234
        return torch.device("cpu")
    else:
235
236
237
238
        if xpu_available:
            return torch.device("xpu")
        else:
            return torch.cuda.current_device()
239
240
241
242
243

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

245

246
def xformers_enabled():
247
248
    global xpu_available
    global directml_enabled
249
    if vram_state == VRAMState.CPU:
250
        return False
251
252
253
254
    if xpu_available:
        return False
    if directml_enabled:
        return False
255
    return XFORMERS_IS_AVAILABLE
256

257
258
259
260
261

def xformers_enabled_vae():
    enabled = xformers_enabled()
    if not enabled:
        return False
262
263

    return XFORMERS_ENABLED_VAE
264

265
266
267
def pytorch_attention_enabled():
    return ENABLE_PYTORCH_ATTENTION

268
def get_free_memory(dev=None, torch_free_too=False):
269
    global xpu_available
270
    global directml_enabled
271
    if dev is None:
272
        dev = get_torch_device()
273

Yurii Mazurevich's avatar
Yurii Mazurevich committed
274
    if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
275
276
277
        mem_free_total = psutil.virtual_memory().available
        mem_free_torch = mem_free_total
    else:
278
279
280
281
        if directml_enabled:
            mem_free_total = 1024 * 1024 * 1024 #TODO
            mem_free_torch = mem_free_total
        elif xpu_available:
282
283
284
285
286
287
288
289
290
            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
291
292
293
294
295

    if torch_free_too:
        return (mem_free_total, mem_free_torch)
    else:
        return mem_free_total
296
297
298

def maximum_batch_area():
    global vram_state
299
    if vram_state == VRAMState.NO_VRAM:
300
301
302
303
304
        return 0

    memory_free = get_free_memory() / (1024 * 1024)
    area = ((memory_free - 1024) * 0.9) / (0.6)
    return int(max(area, 0))
305
306
307

def cpu_mode():
    global vram_state
308
    return vram_state == VRAMState.CPU
309

Yurii Mazurevich's avatar
Yurii Mazurevich committed
310
311
def mps_mode():
    global vram_state
312
    return vram_state == VRAMState.MPS
Yurii Mazurevich's avatar
Yurii Mazurevich committed
313

314
def should_use_fp16():
315
    global xpu_available
316
317
    global directml_enabled

318
319
320
    if FORCE_FP32:
        return False

321
322
323
    if directml_enabled:
        return False

324
    if cpu_mode() or mps_mode() or xpu_available:
325
326
327
328
329
        return False #TODO ?

    if torch.cuda.is_bf16_supported():
        return True

comfyanonymous's avatar
comfyanonymous committed
330
    props = torch.cuda.get_device_properties("cuda")
331
332
333
334
    if props.major < 7:
        return False

    #FP32 is faster on those cards?
335
    nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600"]
336
337
338
339
340
341
    for x in nvidia_16_series:
        if x in props.name:
            return False

    return True

342
343
344
345
346
347
348
349
350
def soft_empty_cache():
    global xpu_available
    if xpu_available:
        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()

351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
#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()