diff --git a/py/power_puter.py b/py/power_puter.py index da19f83..c47fa37 100644 --- a/py/power_puter.py +++ b/py/power_puter.py @@ -39,6 +39,20 @@ class Function(): args: tuple[int, Optional[int]] +def purge_vram(purge_models=True): + """Purges vram and, optionally, unloads models.""" + import gc + import torch + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + if purge_models: + import comfy + comfy.model_management.unload_all_models() + comfy.model_management.soft_empty_cache() + + _FUNCTIONS = { fn.name: fn for fn in [ Function(name="round", call=round, args=(1, 2)), @@ -51,6 +65,7 @@ _FUNCTIONS = { Function(name="random_choice", call=random.choice, args=(2, None)), Function(name="re", call=re.compile, args=(1, 1)), Function(name="len", call=len, args=(1, 1)), + Function(name="enumerate", call=enumerate, args=(1, 1)), # Casts Function(name="int", call=int, args=(1, 1)), Function(name="float", call=float, args=(1, 1)), @@ -62,6 +77,7 @@ _FUNCTIONS = { Function(name="node", call='_get_node', args=(0, 1)), Function(name="nodes", call='_get_nodes', args=(0, 1)), Function(name="input_node", call='_get_input_node', args=(0, 1)), + Function(name="purge_vram", call=purge_vram, args=(0, 1)), Function(name="dir", call=dir, args=(1, 1)), Function(name="type", call=type, args=(1, 1)), Function(name="print", call=print, args=(0, None)), diff --git a/pyproject.toml b/pyproject.toml index 8677fab..5c229af 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "rgthree-comfy" description = "Making ComfyUI more comfortable." -version = "1.0.2506222155" +version = "1.0.2506222201" license = { file = "LICENSE" } dependencies = []