From 7a1bb513976b16d56563690eada4bbb267353a8a Mon Sep 17 00:00:00 2001
From: Zuellni <123005779+Zuellni@users.noreply.github.com>
Date: Sat, 7 Oct 2023 11:12:31 +0200
Subject: [PATCH] Allow loading nodes on their own without using the repo

---
 README.md                |  2 +-
 __init__.py              | 16 +---------------
 nodes.py => pickscore.py | 23 +++++++++++++++++++----
 3 files changed, 21 insertions(+), 20 deletions(-)
 rename nodes.py => pickscore.py (87%)

diff --git a/README.md b/README.md
index 622aaf5..ac5d042 100644
--- a/README.md
+++ b/README.md
@@ -11,7 +11,7 @@ Name | Description
 :--- | :---
 Loader | Loads scoring models from [Hugging Face](https://huggingface.co) or a given directory. The [default](https://huggingface.co/yuvalkirstain/PickScore_v1) one is around 4GB.
 Processor | Takes images/text and converts them to inputs for the `Selector` node.
-Selector | Selects up to `count` best images/latents/masks. Interrupts processing if the `threshold` isn't reached.
+Selector | Selects up to `limit` best images/latents/masks. Interrupts processing if the `threshold` isn't reached.
 
 ## Workflow
 The image below can be opened in ComfyUI.
diff --git a/__init__.py b/__init__.py
index 55269b4..89ab8ea 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,15 +1 @@
-from .nodes import ImageProcessor, Loader, Selector, TextProcessor
-
-NODE_CLASS_MAPPINGS = {
-    "ZuellniPickScoreLoader": Loader,
-    "ZuellniPickScoreImageProcessor": ImageProcessor,
-    "ZuellniPickScoreTextProcessor": TextProcessor,
-    "ZuellniPickScoreSelector": Selector,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "ZuellniPickScoreLoader": "PickScore Loader",
-    "ZuellniPickScoreImageProcessor": "PickScore Image Processor",
-    "ZuellniPickScoreTextProcessor": "PickScore Text Processor",
-    "ZuellniPickScoreSelector": "PickScore Selector",
-}
+from .pickscore import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
diff --git a/nodes.py b/pickscore.py
similarity index 87%
rename from nodes.py
rename to pickscore.py
index bffef7a..ada2e1a 100644
--- a/nodes.py
+++ b/pickscore.py
@@ -88,8 +88,8 @@ class Selector:
                 "model": ("PS_MODEL",),
                 "image_inputs": ("IMAGE_INPUTS",),
                 "text_inputs": ("TEXT_INPUTS",),
-                "threshold": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.001}),
-                "count": ("INT", {"default": 1, "min": 1, "max": 1024}),
+                "threshold": ("FLOAT", {"default": 0, "max": 1, "step": 0.001}),
+                "limit": ("INT", {"default": 1, "min": 1, "max": 1024}),
             },
             "optional": {
                 "images": ("IMAGE",),
@@ -109,7 +109,7 @@ class Selector:
         image_inputs,
         text_inputs,
         threshold,
-        count,
+        limit,
         images=None,
         latents=None,
         masks=None,
@@ -131,7 +131,7 @@ class Selector:
 
         scores = scores.cpu().tolist()
         scores = {k: v for k, v in enumerate(scores) if v >= threshold}
-        scores = sorted(scores.items(), key=lambda k: k[1], reverse=True)[:count]
+        scores = sorted(scores.items(), key=lambda k: k[1], reverse=True)[:limit]
         scores_str = ", ".join([str(round(v, 3)) for k, v in scores])
 
         if images is not None:
@@ -151,3 +151,18 @@ class Selector:
             raise InterruptProcessingException()
 
         return (scores_str, images, latents, masks)
+
+
+NODE_CLASS_MAPPINGS = {
+    "ZuellniPickScoreLoader": Loader,
+    "ZuellniPickScoreImageProcessor": ImageProcessor,
+    "ZuellniPickScoreTextProcessor": TextProcessor,
+    "ZuellniPickScoreSelector": Selector,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ZuellniPickScoreLoader": "Loader",
+    "ZuellniPickScoreImageProcessor": "Image Processor",
+    "ZuellniPickScoreTextProcessor": "Text Processor",
+    "ZuellniPickScoreSelector": "Selector",
+}