From ac5ffb4ed05546545ce7cf38e7b69b5152714eed Mon Sep 17 00:00:00 2001
From: SeanScripts <64337075+SeanScripts@users.noreply.github.com>
Date: Thu, 12 Jun 2025 21:22:11 -0700
Subject: [PATCH] Fix #6

---
 pyproject.toml | 2 +-
 unloadModel.py | 6 ++----
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index aa994c4..b39794a 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.1"
+version = "1.0.2"
 license = {file = "LICENSE"}
 
 [project.urls]
diff --git a/unloadModel.py b/unloadModel.py
index d5f66e1..021c2cb 100644
--- a/unloadModel.py
+++ b/unloadModel.py
@@ -51,8 +51,7 @@ class UnloadModelNode:
             torch.cuda.ipc_collect()
         except:
             print("   - Unable to clear cache")
-        #time.sleep(2) # why?
-        return (list(kwargs.values()))
+        return (kwargs.get("value"),)
 
 class UnloadAllModelsNode:
     @classmethod
@@ -81,8 +80,7 @@ class UnloadAllModelsNode:
             torch.cuda.ipc_collect()
         except:
             print("   - Unable to clear cache")
-        #time.sleep(2) # why?
-        return (list(kwargs.values()))
+        return (kwargs.get("value"),)
 
 
 NODE_CLASS_MAPPINGS = {