mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-09-13 04:55:53 +00:00
Don't unload/reload model from CPU uselessly.
This commit is contained in:
26
comfy/model_management.py
Normal file
26
comfy/model_management.py
Normal file
@@ -0,0 +1,26 @@
|
||||
|
||||
|
||||
current_loaded_model = None
|
||||
|
||||
|
||||
def unload_model():
|
||||
global current_loaded_model
|
||||
if current_loaded_model is not None:
|
||||
current_loaded_model.model.cpu()
|
||||
current_loaded_model.unpatch_model()
|
||||
current_loaded_model = None
|
||||
|
||||
|
||||
def load_model_gpu(model):
|
||||
global current_loaded_model
|
||||
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
|
||||
real_model.cuda()
|
||||
return current_loaded_model
|
Reference in New Issue
Block a user