From 90f6043f46b3164a35fcfeb72dae4beb4d98577a Mon Sep 17 00:00:00 2001 From: SeanScripts Date: Fri, 4 Oct 2024 21:28:03 -0700 Subject: [PATCH] Support unloading vision models You can also use the built-in "unload after generate", but this gives a bit more control. --- pyproject.toml | 2 +- unloadModel.py | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 9656579..aa994c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-unload-model" description = "For unloading a model or all models, using the memory management that is already present in ComfyUI. Copied from [a/https://github.com/willblaschko/ComfyUI-Unload-Models](https://github.com/willblaschko/ComfyUI-Unload-Models) but without the unnecessary extra stuff." -version = "1.0.0" +version = "1.0.1" license = {file = "LICENSE"} [project.urls] diff --git a/unloadModel.py b/unloadModel.py index 53c2118..d5f66e1 100644 --- a/unloadModel.py +++ b/unloadModel.py @@ -32,6 +32,16 @@ class UnloadModelNode: if kwargs.get("model") in loaded_models: print(" - Model found in memory, unloading...") loaded_models.remove(kwargs.get("model")) + else: + # Just delete it, I need the VRAM! + model = kwargs.get("model") + if type(model) == dict: + keys = [(key, type(value).__name__) for key, value in model.items()] + for key, model_type in keys: + if key == 'model': + print(f"Unloading model of type {model_type}") + del model[key] + # Emptying the cache after this should free the memory. model_management.free_memory(1e30, model_management.get_torch_device(), loaded_models) model_management.soft_empty_cache(True) try: