298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-04-19 19:03:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

Make sure additional_models are unloaded as well when perform

Browse Source
This commit is contained in:
Jedrzej Kosinski 2025-01-23 19:06:05 -06:00
parent 02a4d0ad7d
commit 5db4277449
1 changed files with 12 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
comfy/model_management.py
View File

@ -1146,14 +1146,25 @@ def soft_empty_cache(force=False):
def unload_all_models(): def unload_all_models():
free_memory(1e30, get_torch_device()) free_memory(1e30, get_torch_device())
def unload_model_and_clones(model: ModelPatcher): def unload_model_and_clones(model: ModelPatcher, unload_additional_models=True):
'Unload only model and its clones - primarily for multigpu cloning purposes.' 'Unload only model and its clones - primarily for multigpu cloning purposes.'
initial_keep_loaded: list[LoadedModel] = current_loaded_models.copy() initial_keep_loaded: list[LoadedModel] = current_loaded_models.copy()
additional_models = []
if unload_additional_models:
additional_models = model.get_nested_additional_models()
keep_loaded = [] keep_loaded = []
for loaded_model in initial_keep_loaded: for loaded_model in initial_keep_loaded:
if loaded_model.model is not None: if loaded_model.model is not None:
if model.clone_base_uuid == loaded_model.model.clone_base_uuid: if model.clone_base_uuid == loaded_model.model.clone_base_uuid:
continue continue
# check additional models if they are a match
skip = False
for add_model in additional_models:
if add_model.clone_base_uuid == loaded_model.model.clone_base_uuid:
skip = True
break
if skip:
continue
keep_loaded.append(loaded_model) keep_loaded.append(loaded_model)
free_memory(1e30, get_torch_device(), keep_loaded) free_memory(1e30, get_torch_device(), keep_loaded)