diff --git a/comfy/model_management.py b/comfy/model_management.py index 2cf792b5..c72ed247 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -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)