model_management.py 5.88 KB
Newer Older
1

2
3
4
5
CPU = 0
NO_VRAM = 1
LOW_VRAM = 2
NORMAL_VRAM = 3
6
HIGH_VRAM = 4
7
8
9
10

accelerate_enabled = False
vram_state = NORMAL_VRAM

11
total_vram = 0
12
13
total_vram_available_mb = -1

14
import sys
15
import psutil
16
17
18

set_vram_to = NORMAL_VRAM

19
20
21
try:
    import torch
    total_vram = torch.cuda.mem_get_info(torch.cuda.current_device())[1] / (1024 * 1024)
22
23
24
25
26
27
    total_ram = psutil.virtual_memory().total / (1024 * 1024)
    forced_normal_vram = "--normalvram" in sys.argv
    if not forced_normal_vram:
        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")
            set_vram_to = LOW_VRAM
comfyanonymous's avatar
comfyanonymous committed
28
        elif total_vram > total_ram * 1.1 and total_vram > 14336:
29
30
            print("Enabling highvram mode because your GPU has more vram than your computer has ram. If you don't want this use: --normalvram")
            vram_state = HIGH_VRAM
31
32
33
except:
    pass

34
35
36
37
if "--lowvram" in sys.argv:
    set_vram_to = LOW_VRAM
if "--novram" in sys.argv:
    set_vram_to = NO_VRAM
38
39
if "--highvram" in sys.argv:
    vram_state = HIGH_VRAM
40
41


42
if set_vram_to == LOW_VRAM or set_vram_to == NO_VRAM:
43
44
45
46
47
48
49
50
    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.")
51
52

    total_vram_available_mb = (total_vram - 1024) // 2
53
    total_vram_available_mb = int(max(256, total_vram_available_mb))
54
55


56
print("Set vram state to:", ["CPU", "NO VRAM", "LOW VRAM", "NORMAL VRAM", "HIGH VRAM"][vram_state])
57

58
59

current_loaded_model = None
comfyanonymous's avatar
comfyanonymous committed
60
current_gpu_controlnets = []
61

62
63
64
model_accelerated = False


65
66
def unload_model():
    global current_loaded_model
67
    global model_accelerated
comfyanonymous's avatar
comfyanonymous committed
68
    global current_gpu_controlnets
69
70
    global vram_state

71
    if current_loaded_model is not None:
72
73
74
75
        if model_accelerated:
            accelerate.hooks.remove_hook_from_submodules(current_loaded_model.model)
            model_accelerated = False

76
77
78
        #never unload models from GPU on high vram
        if vram_state != HIGH_VRAM:
            current_loaded_model.model.cpu()
79
80
        current_loaded_model.unpatch_model()
        current_loaded_model = None
81
82
83
84
85
86

    if vram_state != HIGH_VRAM:
        if len(current_gpu_controlnets) > 0:
            for n in current_gpu_controlnets:
                n.cpu()
            current_gpu_controlnets = []
87
88
89
90


def load_model_gpu(model):
    global current_loaded_model
91
92
93
    global vram_state
    global model_accelerated

94
95
96
97
98
99
100
101
102
    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
103
104
    if vram_state == CPU:
        pass
105
    elif vram_state == NORMAL_VRAM or vram_state == HIGH_VRAM:
106
107
108
109
110
111
        model_accelerated = False
        real_model.cuda()
    else:
        if vram_state == NO_VRAM:
            device_map = accelerate.infer_auto_device_map(real_model, max_memory={0: "256MiB", "cpu": "16GiB"})
        elif vram_state == LOW_VRAM:
112
            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
113

114
115
        accelerate.dispatch_model(real_model, device_map=device_map, main_device="cuda")
        model_accelerated = True
116
    return current_loaded_model
117

comfyanonymous's avatar
comfyanonymous committed
118
119
def load_controlnet_gpu(models):
    global current_gpu_controlnets
120
121
122
123
124
125
    global vram_state

    if vram_state == LOW_VRAM or vram_state == NO_VRAM:
        #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
126
127
128
129
130
131
132
133
    for m in current_gpu_controlnets:
        if m not in models:
            m.cpu()

    current_gpu_controlnets = []
    for m in models:
        current_gpu_controlnets.append(m.cuda())

134

135
136
137
138
139
140
141
142
143
144
145
146
147
def load_if_low_vram(model):
    global vram_state
    if vram_state == LOW_VRAM or vram_state == NO_VRAM:
        return model.cuda()
    return model

def unload_if_low_vram(model):
    global vram_state
    if vram_state == LOW_VRAM or vram_state == NO_VRAM:
        return model.cpu()
    return model


148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
def get_free_memory(dev=None, torch_free_too=False):
    if dev is None:
        dev = torch.cuda.current_device()

    if hasattr(dev, 'type') and dev.type == 'cpu':
        mem_free_total = psutil.virtual_memory().available
        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

    if torch_free_too:
        return (mem_free_total, mem_free_torch)
    else:
        return mem_free_total
167
168
169
170
171
172
173
174
175

def maximum_batch_area():
    global vram_state
    if vram_state == NO_VRAM:
        return 0

    memory_free = get_free_memory() / (1024 * 1024)
    area = ((memory_free - 1024) * 0.9) / (0.6)
    return int(max(area, 0))
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
#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()