diff --git a/backend/memory_management.py b/backend/memory_management.py index e14448ad..5490cae0 100644 --- a/backend/memory_management.py +++ b/backend/memory_management.py @@ -1313,6 +1313,21 @@ def soft_empty_cache(force=False): signal_empty_cache = False +def unload_model(model: "ModelPatcher") -> bool: + index = None + for i, p in enumerate(current_loaded_models): + if p.model == model: + index = i + break + + if index is not None: + mdl = current_loaded_models.pop(index) + del mdl + return True + + return False + + def unload_all_models(): free_memory(1e30, get_torch_device()) diff --git a/backend/patcher/controlnet.py b/backend/patcher/controlnet.py index 772a393d..c966de2e 100644 --- a/backend/patcher/controlnet.py +++ b/backend/patcher/controlnet.py @@ -344,6 +344,7 @@ class ControlNet(ControlBase): def cleanup(self): self.model_sampling_current = None + memory_management.unload_model(self.control_model_wrapped) super().cleanup()