mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-01-11 10:25:16 +00:00
27 lines
590 B
Python
27 lines
590 B
Python
|
|
||
|
|
|
||
|
|
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
|