Support unloading vision models
You can also use the built-in "unload after generate", but this gives a bit more control.
This commit is contained in:
+1
-1
@@ -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]
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user