model_management.py 20.4 KB
Newer Older
1
2
import psutil
from enum import Enum
comfyanonymous's avatar
comfyanonymous committed
3
from comfy.cli_args import args
4
import torch
comfyanonymous's avatar
comfyanonymous committed
5
import sys
6

7
class VRAMState(Enum):
8
9
    DISABLED = 0    #No vram present: no need to move models to vram
    NO_VRAM = 1     #Very low vram: enable all the options to save vram
10
11
12
    LOW_VRAM = 2
    NORMAL_VRAM = 3
    HIGH_VRAM = 4
13
    SHARED = 5      #No dedicated vram: memory shared between CPU and GPU but models still need to be moved between both.
14
15
16
17
18

class CPUState(Enum):
    GPU = 0
    CPU = 1
    MPS = 2
19

20
21
22
# Determine VRAM State
vram_state = VRAMState.NORMAL_VRAM
set_vram_to = VRAMState.NORMAL_VRAM
23
cpu_state = CPUState.GPU
24

25
total_vram = 0
26

27
lowvram_available = True
藍+85CD's avatar
藍+85CD committed
28
xpu_available = False
29

30
directml_enabled = False
31
if args.directml is not None:
32
33
    import torch_directml
    directml_enabled = True
34
35
36
37
38
39
    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))
40
    # torch_directml.disable_tiled_resources(True)
41
    lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default.
42

43
try:
44
45
46
    import intel_extension_for_pytorch as ipex
    if torch.xpu.is_available():
        xpu_available = True
47
48
49
except:
    pass

50
51
52
try:
    if torch.backends.mps.is_available():
        cpu_state = CPUState.MPS
KarryCharon's avatar
KarryCharon committed
53
        import torch.mps
54
55
56
57
58
59
except:
    pass

if args.cpu:
    cpu_state = CPUState.CPU

60
61
62
def get_torch_device():
    global xpu_available
    global directml_enabled
63
    global cpu_state
64
65
66
    if directml_enabled:
        global directml_device
        return directml_device
67
    if cpu_state == CPUState.MPS:
68
        return torch.device("mps")
69
    if cpu_state == CPUState.CPU:
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
        return torch.device("cpu")
    else:
        if xpu_available:
            return torch.device("xpu")
        else:
            return torch.device(torch.cuda.current_device())

def get_total_memory(dev=None, torch_total_too=False):
    global xpu_available
    global directml_enabled
    if dev is None:
        dev = get_torch_device()

    if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
        mem_total = psutil.virtual_memory().total
        mem_total_torch = mem_total
    else:
        if directml_enabled:
            mem_total = 1024 * 1024 * 1024 #TODO
            mem_total_torch = mem_total
        elif xpu_available:
91
92
            stats = torch.xpu.memory_stats(dev)
            mem_reserved = stats['reserved_bytes.all.current']
93
            mem_total = torch.xpu.get_device_properties(dev).total_memory
94
            mem_total_torch = mem_reserved
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
        else:
            stats = torch.cuda.memory_stats(dev)
            mem_reserved = stats['reserved_bytes.all.current']
            _, mem_total_cuda = torch.cuda.mem_get_info(dev)
            mem_total_torch = mem_reserved
            mem_total = mem_total_cuda

    if torch_total_too:
        return (mem_total, mem_total_torch)
    else:
        return mem_total

total_vram = get_total_memory(get_torch_device()) / (1024 * 1024)
total_ram = psutil.virtual_memory().total / (1024 * 1024)
print("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))
if not args.normalvram and not args.cpu:
    if lowvram_available and 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")
        set_vram_to = VRAMState.LOW_VRAM

115
116
117
118
119
try:
    OOM_EXCEPTION = torch.cuda.OutOfMemoryError
except:
    OOM_EXCEPTION = Exception

120
121
XFORMERS_VERSION = ""
XFORMERS_ENABLED_VAE = True
122
123
if args.disable_xformers:
    XFORMERS_IS_AVAILABLE = False
124
125
126
127
else:
    try:
        import xformers
        import xformers.ops
128
        XFORMERS_IS_AVAILABLE = True
129
130
131
132
133
134
135
136
137
138
139
        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
140
    except:
141
        XFORMERS_IS_AVAILABLE = False
142

143
144
145
146
147
148
def is_nvidia():
    global cpu_state
    if cpu_state == CPUState.GPU:
        if torch.version.cuda:
            return True

149
ENABLE_PYTORCH_ATTENTION = args.use_pytorch_cross_attention
150
151
152
153
154
155
156
157
158
159

if ENABLE_PYTORCH_ATTENTION == False and XFORMERS_IS_AVAILABLE == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
    try:
        if is_nvidia():
            torch_version = torch.version.__version__
            if int(torch_version[0]) >= 2:
                ENABLE_PYTORCH_ATTENTION = True
    except:
        pass

160
if ENABLE_PYTORCH_ATTENTION:
161
162
163
    torch.backends.cuda.enable_math_sdp(True)
    torch.backends.cuda.enable_flash_sdp(True)
    torch.backends.cuda.enable_mem_efficient_sdp(True)
164
    XFORMERS_IS_AVAILABLE = False
165

166
167
if args.lowvram:
    set_vram_to = VRAMState.LOW_VRAM
168
    lowvram_available = True
169
170
elif args.novram:
    set_vram_to = VRAMState.NO_VRAM
171
elif args.highvram or args.gpu_only:
172
    vram_state = VRAMState.HIGH_VRAM
173

174
FORCE_FP32 = False
175
FORCE_FP16 = False
176
177
178
179
if args.force_fp32:
    print("Forcing FP32, if this improves things please report it.")
    FORCE_FP32 = True

180
181
182
183
if args.force_fp16:
    print("Forcing FP16.")
    FORCE_FP16 = True

184
if lowvram_available:
185
186
    try:
        import accelerate
187
188
        if set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM):
            vram_state = set_vram_to
189
190
191
    except Exception as e:
        import traceback
        print(traceback.format_exc())
192
193
        print("ERROR: LOW VRAM MODE NEEDS accelerate.")
        lowvram_available = False
194

195

196
197
if cpu_state != CPUState.GPU:
    vram_state = VRAMState.DISABLED
198

199
200
if cpu_state == CPUState.MPS:
    vram_state = VRAMState.SHARED
201

202
print(f"Set vram state to: {vram_state.name}")
203

204
205
206
207
DISABLE_SMART_MEMORY = args.disable_smart_memory

if DISABLE_SMART_MEMORY:
    print("Disabling smart memory management")
208

209
def get_torch_device_name(device):
210
    global xpu_available
211
    if hasattr(device, 'type'):
212
        if device.type == "cuda":
213
214
215
216
217
            try:
                allocator_backend = torch.cuda.get_allocator_backend()
            except:
                allocator_backend = ""
            return "{} {} : {}".format(device, torch.cuda.get_device_name(device), allocator_backend)
218
219
        else:
            return "{}".format(device.type)
220
221
    elif xpu_available:
        return "{} {}".format(device, torch.xpu.get_device_name(device))
222
223
    else:
        return "CUDA {}: {}".format(device, torch.cuda.get_device_name(device))
224
225

try:
226
    print("Device:", get_torch_device_name(get_torch_device()))
227
228
229
except:
    print("Could not pick default device.")

230

comfyanonymous's avatar
comfyanonymous committed
231
current_loaded_models = []
232

comfyanonymous's avatar
comfyanonymous committed
233
234
235
236
237
class LoadedModel:
    def __init__(self, model):
        self.model = model
        self.model_accelerated = False
        self.device = model.load_device
238

comfyanonymous's avatar
comfyanonymous committed
239
240
    def model_memory(self):
        return self.model.model_size()
241

comfyanonymous's avatar
comfyanonymous committed
242
243
244
245
246
    def model_memory_required(self, device):
        if device == self.model.current_device:
            return 0
        else:
            return self.model_memory()
247

comfyanonymous's avatar
comfyanonymous committed
248
    def model_load(self, lowvram_model_memory=0):
249
        global xpu_available
comfyanonymous's avatar
comfyanonymous committed
250
251
252
        patch_model_to = None
        if lowvram_model_memory == 0:
            patch_model_to = self.device
253

comfyanonymous's avatar
comfyanonymous committed
254
255
        self.model.model_patches_to(self.device)
        self.model.model_patches_to(self.model.model_dtype())
256

comfyanonymous's avatar
comfyanonymous committed
257
258
259
260
261
262
        try:
            self.real_model = self.model.patch_model(device_to=patch_model_to) #TODO: do something with loras and offloading to CPU
        except Exception as e:
            self.model.unpatch_model(self.model.offload_device)
            self.model_unload()
            raise e
263

comfyanonymous's avatar
comfyanonymous committed
264
265
266
267
268
        if lowvram_model_memory > 0:
            print("loading in lowvram mode", lowvram_model_memory/(1024 * 1024))
            device_map = accelerate.infer_auto_device_map(self.real_model, max_memory={0: "{}MiB".format(lowvram_model_memory // (1024 * 1024)), "cpu": "16GiB"})
            accelerate.dispatch_model(self.real_model, device_map=device_map, main_device=self.device)
            self.model_accelerated = True
269

270
        if xpu_available and not args.disable_ipex_optimize:
271
            self.real_model = torch.xpu.optimize(self.real_model.eval(), inplace=True, auto_kernel_selection=True, graph_mode=True)
272

comfyanonymous's avatar
comfyanonymous committed
273
        return self.real_model
274

comfyanonymous's avatar
comfyanonymous committed
275
276
277
278
    def model_unload(self):
        if self.model_accelerated:
            accelerate.hooks.remove_hook_from_submodules(self.real_model)
            self.model_accelerated = False
279

comfyanonymous's avatar
comfyanonymous committed
280
281
        self.model.unpatch_model(self.model.offload_device)
        self.model.model_patches_to(self.model.offload_device)
282

comfyanonymous's avatar
comfyanonymous committed
283
284
    def __eq__(self, other):
        return self.model is other.model
comfyanonymous's avatar
comfyanonymous committed
285

comfyanonymous's avatar
comfyanonymous committed
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
def minimum_inference_memory():
    return (1024 * 1024 * 1024)

def unload_model_clones(model):
    to_unload = []
    for i in range(len(current_loaded_models)):
        if model.is_clone(current_loaded_models[i].model):
            to_unload = [i] + to_unload

    for i in to_unload:
        print("unload clone", i)
        current_loaded_models.pop(i).model_unload()

def free_memory(memory_required, device, keep_loaded=[]):
    unloaded_model = False
    for i in range(len(current_loaded_models) -1, -1, -1):
comfyanonymous's avatar
comfyanonymous committed
302
303
304
        if not DISABLE_SMART_MEMORY:
            if get_free_memory(device) > memory_required:
                break
comfyanonymous's avatar
comfyanonymous committed
305
306
307
        shift_model = current_loaded_models[i]
        if shift_model.device == device:
            if shift_model not in keep_loaded:
comfyanonymous's avatar
comfyanonymous committed
308
309
310
                m = current_loaded_models.pop(i)
                m.model_unload()
                del m
comfyanonymous's avatar
comfyanonymous committed
311
312
313
314
315
316
317
                unloaded_model = True

    if unloaded_model:
        soft_empty_cache()


def load_models_gpu(models, memory_required=0):
318
319
    global vram_state

comfyanonymous's avatar
comfyanonymous committed
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
    inference_memory = minimum_inference_memory()
    extra_mem = max(inference_memory, memory_required)

    models_to_load = []
    models_already_loaded = []
    for x in models:
        loaded_model = LoadedModel(x)

        if loaded_model in current_loaded_models:
            index = current_loaded_models.index(loaded_model)
            current_loaded_models.insert(0, current_loaded_models.pop(index))
            models_already_loaded.append(loaded_model)
        else:
            models_to_load.append(loaded_model)

    if len(models_to_load) == 0:
        devs = set(map(lambda a: a.device, models_already_loaded))
        for d in devs:
            if d != torch.device("cpu"):
                free_memory(extra_mem, d, models_already_loaded)
340
341
        return

comfyanonymous's avatar
comfyanonymous committed
342
    print("loading new")
343

comfyanonymous's avatar
comfyanonymous committed
344
345
346
347
    total_memory_required = {}
    for loaded_model in models_to_load:
        unload_model_clones(loaded_model.model)
        total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device)
comfyanonymous's avatar
comfyanonymous committed
348

comfyanonymous's avatar
comfyanonymous committed
349
350
351
    for device in total_memory_required:
        if device != torch.device("cpu"):
            free_memory(total_memory_required[device] * 1.3 + extra_mem, device, models_already_loaded)
comfyanonymous's avatar
comfyanonymous committed
352

comfyanonymous's avatar
comfyanonymous committed
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
    for loaded_model in models_to_load:
        model = loaded_model.model
        torch_dev = model.load_device
        if is_device_cpu(torch_dev):
            vram_set_state = VRAMState.DISABLED
        else:
            vram_set_state = vram_state
        lowvram_model_memory = 0
        if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM):
            model_size = loaded_model.model_memory_required(torch_dev)
            current_free_mem = get_free_memory(torch_dev)
            lowvram_model_memory = int(max(256 * (1024 * 1024), (current_free_mem - 1024 * (1024 * 1024)) / 1.3 ))
            if model_size > (current_free_mem - inference_memory): #only switch to lowvram if really necessary
                vram_set_state = VRAMState.LOW_VRAM
            else:
                lowvram_model_memory = 0
369

comfyanonymous's avatar
comfyanonymous committed
370
371
        if vram_set_state == VRAMState.NO_VRAM:
            lowvram_model_memory = 256 * 1024 * 1024
372

comfyanonymous's avatar
comfyanonymous committed
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
        cur_loaded_model = loaded_model.model_load(lowvram_model_memory)
        current_loaded_models.insert(0, loaded_model)
    return


def load_model_gpu(model):
    return load_models_gpu([model])

def cleanup_models():
    to_delete = []
    for i in range(len(current_loaded_models)):
        print(sys.getrefcount(current_loaded_models[i].model))
        if sys.getrefcount(current_loaded_models[i].model) <= 2:
            to_delete = [i] + to_delete

    for i in to_delete:
        x = current_loaded_models.pop(i)
        x.model_unload()
        del x
392

393
394
395
396
397
398
def dtype_size(dtype):
    dtype_size = 4
    if dtype == torch.float16 or dtype == torch.bfloat16:
        dtype_size = 2
    return dtype_size

399
def unet_offload_device():
comfyanonymous's avatar
comfyanonymous committed
400
    if vram_state == VRAMState.HIGH_VRAM:
401
402
403
404
        return get_torch_device()
    else:
        return torch.device("cpu")

comfyanonymous's avatar
comfyanonymous committed
405
406
407
408
409
410
def unet_inital_load_device(parameters, dtype):
    torch_dev = get_torch_device()
    if vram_state == VRAMState.HIGH_VRAM:
        return torch_dev

    cpu_dev = torch.device("cpu")
411
412
413
    if DISABLE_SMART_MEMORY:
        return cpu_dev

414
    model_size = dtype_size(dtype) * parameters
comfyanonymous's avatar
comfyanonymous committed
415
416
417
418
419
420
421
422

    mem_dev = get_free_memory(torch_dev)
    mem_cpu = get_free_memory(cpu_dev)
    if mem_dev > mem_cpu and model_size < mem_dev:
        return torch_dev
    else:
        return cpu_dev

423
def text_encoder_offload_device():
comfyanonymous's avatar
comfyanonymous committed
424
    if args.gpu_only:
425
426
427
428
        return get_torch_device()
    else:
        return torch.device("cpu")

429
def text_encoder_device():
comfyanonymous's avatar
comfyanonymous committed
430
    if args.gpu_only:
431
        return get_torch_device()
432
    elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM:
433
        if should_use_fp16(prioritize_performance=False):
434
435
436
            return get_torch_device()
        else:
            return torch.device("cpu")
437
438
439
    else:
        return torch.device("cpu")

440
441
442
443
def vae_device():
    return get_torch_device()

def vae_offload_device():
comfyanonymous's avatar
comfyanonymous committed
444
    if args.gpu_only:
445
446
447
448
        return get_torch_device()
    else:
        return torch.device("cpu")

449
450
451
452
453
454
455
456
def vae_dtype():
    if args.fp16_vae:
        return torch.float16
    elif args.bf16_vae:
        return torch.bfloat16
    else:
        return torch.float32

457
458
459
460
def get_autocast_device(dev):
    if hasattr(dev, 'type'):
        return dev.type
    return "cuda"
461

462

463
def xformers_enabled():
464
465
    global xpu_available
    global directml_enabled
466
467
    global cpu_state
    if cpu_state != CPUState.GPU:
468
        return False
469
470
471
472
    if xpu_available:
        return False
    if directml_enabled:
        return False
473
    return XFORMERS_IS_AVAILABLE
474

475
476
477
478
479

def xformers_enabled_vae():
    enabled = xformers_enabled()
    if not enabled:
        return False
480
481

    return XFORMERS_ENABLED_VAE
482

483
def pytorch_attention_enabled():
484
    global ENABLE_PYTORCH_ATTENTION
485
486
    return ENABLE_PYTORCH_ATTENTION

487
488
489
490
def pytorch_attention_flash_attention():
    global ENABLE_PYTORCH_ATTENTION
    if ENABLE_PYTORCH_ATTENTION:
        #TODO: more reliable way of checking for flash attention?
491
        if is_nvidia(): #pytorch flash attention only works on Nvidia
492
493
494
            return True
    return False

495
def get_free_memory(dev=None, torch_free_too=False):
496
    global xpu_available
497
    global directml_enabled
498
    if dev is None:
499
        dev = get_torch_device()
500

Yurii Mazurevich's avatar
Yurii Mazurevich committed
501
    if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
502
503
504
        mem_free_total = psutil.virtual_memory().available
        mem_free_torch = mem_free_total
    else:
505
506
507
508
        if directml_enabled:
            mem_free_total = 1024 * 1024 * 1024 #TODO
            mem_free_torch = mem_free_total
        elif xpu_available:
509
510
511
512
513
            stats = torch.xpu.memory_stats(dev)
            mem_active = stats['active_bytes.all.current']
            mem_allocated = stats['allocated_bytes.all.current']
            mem_reserved = stats['reserved_bytes.all.current']
            mem_free_torch = mem_reserved - mem_active
514
            mem_free_total = torch.xpu.get_device_properties(dev).total_memory - mem_allocated
515
516
517
518
519
520
521
        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
522
523
524
525
526

    if torch_free_too:
        return (mem_free_total, mem_free_torch)
    else:
        return mem_free_total
527

comfyanonymous's avatar
comfyanonymous committed
528
529
530
531
532
533
534
def batch_area_memory(area):
    if xformers_enabled() or pytorch_attention_flash_attention():
        #TODO: these formulas are copied from maximum_batch_area below
        return (area / 20) * (1024 * 1024)
    else:
        return (((area * 0.6) / 0.9) + 1024) * (1024 * 1024)

535
536
def maximum_batch_area():
    global vram_state
537
    if vram_state == VRAMState.NO_VRAM:
538
539
540
        return 0

    memory_free = get_free_memory() / (1024 * 1024)
541
    if xformers_enabled() or pytorch_attention_flash_attention():
542
        #TODO: this needs to be tweaked
543
        area = 20 * memory_free
544
545
546
    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)
547
    return int(max(area, 0))
548
549

def cpu_mode():
550
551
    global cpu_state
    return cpu_state == CPUState.CPU
552

Yurii Mazurevich's avatar
Yurii Mazurevich committed
553
def mps_mode():
554
555
    global cpu_state
    return cpu_state == CPUState.MPS
Yurii Mazurevich's avatar
Yurii Mazurevich committed
556

557
558
def is_device_cpu(device):
    if hasattr(device, 'type'):
comfyanonymous's avatar
comfyanonymous committed
559
560
561
562
563
564
565
        if (device.type == 'cpu'):
            return True
    return False

def is_device_mps(device):
    if hasattr(device, 'type'):
        if (device.type == 'mps'):
566
567
568
            return True
    return False

569
def should_use_fp16(device=None, model_params=0, prioritize_performance=True):
570
    global xpu_available
571
572
    global directml_enabled

573
574
575
576
    if device is not None:
        if is_device_cpu(device):
            return False

577
578
579
    if FORCE_FP16:
        return True

580
    if device is not None: #TODO
581
        if is_device_mps(device):
582
            return False
583

584
585
586
    if FORCE_FP32:
        return False

587
588
589
    if directml_enabled:
        return False

590
    if cpu_mode() or mps_mode():
591
592
        return False #TODO ?

comfyanonymous's avatar
comfyanonymous committed
593
594
595
596
    if xpu_available:
        return True

    if torch.cuda.is_bf16_supported():
597
598
        return True

comfyanonymous's avatar
comfyanonymous committed
599
    props = torch.cuda.get_device_properties("cuda")
600
601
602
603
604
605
606
607
608
609
610
611
612
613
    if props.major < 6:
        return False

    fp16_works = False
    #FP16 is confirmed working on a 1080 (GP104) but it's a bit slower than FP32 so it should only be enabled
    #when the model doesn't actually fit on the card
    #TODO: actually test if GP106 and others have the same type of behavior
    nvidia_10_series = ["1080", "1070", "titan x", "p3000", "p3200", "p4000", "p4200", "p5000", "p5200", "p6000", "1060", "1050"]
    for x in nvidia_10_series:
        if x in props.name.lower():
            fp16_works = True

    if fp16_works:
        free_model_memory = (get_free_memory() * 0.9 - minimum_inference_memory())
614
        if (not prioritize_performance) or model_params * 4 > free_model_memory:
615
616
            return True

617
618
619
    if props.major < 7:
        return False

620
    #FP16 is just broken on these cards
621
    nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600", "MX550", "MX450", "CMP 30HX"]
622
623
624
625
626
627
    for x in nvidia_16_series:
        if x in props.name:
            return False

    return True

628
629
def soft_empty_cache():
    global xpu_available
630
631
    global cpu_state
    if cpu_state == CPUState.MPS:
comfyanonymous's avatar
comfyanonymous committed
632
633
        torch.mps.empty_cache()
    elif xpu_available:
634
635
        torch.xpu.empty_cache()
    elif torch.cuda.is_available():
636
        if is_nvidia(): #This seems to make things worse on ROCm so I only do it for cuda
637
638
639
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()

640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
#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()