95 lines
2.7 KiB
Python
95 lines
2.7 KiB
Python
import comfy.model_management as model_management
|
|
import gc
|
|
import torch
|
|
import time
|
|
|
|
# Note: This doesn't work with reroute for some reason?
|
|
class AnyType(str):
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
any = AnyType("*")
|
|
|
|
class UnloadModelNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {"value": (any, )}, # For passthrough
|
|
"optional": {"model": (any, )},
|
|
}
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(s, **kwargs):
|
|
return True
|
|
|
|
RETURN_TYPES = (any, )
|
|
FUNCTION = "route"
|
|
CATEGORY = "Unload Model"
|
|
|
|
def route(self, **kwargs):
|
|
print("Unload Model:")
|
|
loaded_models = model_management.loaded_models()
|
|
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:
|
|
print(" - Clearing Cache...")
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
except:
|
|
print(" - Unable to clear cache")
|
|
return (kwargs.get("value"),)
|
|
|
|
class UnloadAllModelsNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {"value": (any, )},
|
|
}
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(s, **kwargs):
|
|
return True
|
|
|
|
RETURN_TYPES = (any, )
|
|
FUNCTION = "route"
|
|
CATEGORY = "Unload Model"
|
|
|
|
def route(self, **kwargs):
|
|
print("Unload Model:")
|
|
print(" - Unloading all models...")
|
|
model_management.unload_all_models()
|
|
model_management.soft_empty_cache(True)
|
|
try:
|
|
print(" - Clearing Cache...")
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
except:
|
|
print(" - Unable to clear cache")
|
|
return (kwargs.get("value"),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"UnloadModel": UnloadModelNode,
|
|
"UnloadAllModels": UnloadAllModelsNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"UnloadModel": "Unload Model",
|
|
"UnloadAllModels": "Unload All Models",
|
|
}
|