Make sure additional_models are unloaded as well when perform

This commit is contained in:
Jedrzej Kosinski 2025-01-23 19:06:05 -06:00
parent 02a4d0ad7d
commit 5db4277449

View File

@ -1146,14 +1146,25 @@ def soft_empty_cache(force=False):
def unload_all_models():
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.'
initial_keep_loaded: list[LoadedModel] = current_loaded_models.copy()
additional_models = []
if unload_additional_models:
additional_models = model.get_nested_additional_models()
keep_loaded = []
for loaded_model in initial_keep_loaded:
if loaded_model.model is not None:
if model.clone_base_uuid == loaded_model.model.clone_base_uuid:
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)
free_memory(1e30, get_torch_device(), keep_loaded)