Don't unload/reload model from CPU uselessly.

This commit is contained in:
comfyanonymous
2023-02-08 03:17:54 -05:00
parent e3e65947f2
commit a84cd0d1ad
3 changed files with 58 additions and 32 deletions

26
comfy/model_management.py Normal file
View 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