From 001a8a596febc5c08fa1388b6bb5ad07bdf9c279 Mon Sep 17 00:00:00 2001
From: Frost Ming <me@frostming.com>
Date: Wed, 6 Nov 2024 13:51:38 +0800
Subject: [PATCH] feat: idl

Signed-off-by: Frost Ming <me@frostming.com>
---
 monkeypatch.py | 31 +++++++++++++++++++++++++++++++
 nodes.py       | 26 +++++++++++++++++++++++---
 2 files changed, 54 insertions(+), 3 deletions(-)
 create mode 100644 monkeypatch.py

diff --git a/monkeypatch.py b/monkeypatch.py
new file mode 100644
index 0000000..f75b731
--- /dev/null
+++ b/monkeypatch.py
@@ -0,0 +1,31 @@
+from threading import Lock
+
+import execution
+
+BENTOML_LAST_ID = None
+BENTOML_OUTPUT_CACHE = {}
+
+_lock = Lock()
+
+
+def store_bentoml_value(func):
+    def wrapped(
+        inputs, class_def, unique_id, outputs=None, dynprompt=None, extra_data={}
+    ):
+        global BENTOML_LAST_ID
+        if hasattr(class_def, "BENTOML_NODE"):
+            with _lock:
+                BENTOML_LAST_ID = unique_id
+        if outputs is None:
+            outputs = BENTOML_OUTPUT_CACHE
+        return func(inputs, class_def, unique_id, outputs, dynprompt, extra_data)
+
+    return wrapped
+
+
+execution.get_input_data = store_bentoml_value(execution.get_input_data)
+
+
+def set_bentoml_output(output):
+    with _lock:
+        BENTOML_OUTPUT_CACHE[BENTOML_LAST_ID] = output
diff --git a/nodes.py b/nodes.py
index 41d013e..9aa35af 100644
--- a/nodes.py
+++ b/nodes.py
@@ -10,6 +10,8 @@ import torch
 from PIL import Image, ImageOps, ImageSequence
 from PIL.PngImagePlugin import PngInfo
 
+from .monkeypatch import set_bentoml_output
+
 
 # AnyType class hijacks the isinstance, issubclass, bool, str, jsonserializable, eq, ne methods to always return True
 class AnyType(str):
@@ -28,8 +30,6 @@ class OutputPath:
         return {
             "required": {
                 "filename": ("STRING", {"default": ""}),
-            },
-            "hidden": {
                 "filename_prefix": ("STRING", {"default": "BentoML"}),
             },
         }
@@ -65,11 +65,11 @@ class OutputImage:
         return {
             "required": {
                 "images": ("IMAGE", {"tooltip": "The images to save."}),
+                "filename_prefix": ("STRING", {"default": "BentoML"}),
             },
             "hidden": {
                 "prompt": "PROMPT",
                 "extra_pnginfo": "EXTRA_PNGINFO",
-                "filename_prefix": ("STRING", {"default": "BentoML"}),
             },
         }
 
@@ -212,6 +212,10 @@ class StringInput:
     def string_input(self, value):
         return (value,)
 
+    @classmethod
+    def VALIDATE_INPUTS(s, value):
+        set_bentoml_output([(value,)])
+
 
 class IntegerInput:
     @classmethod
@@ -229,6 +233,10 @@ class IntegerInput:
     def identity(self, value):
         return (value,)
 
+    @classmethod
+    def VALIDATE_INPUTS(s, value):
+        set_bentoml_output([(value,)])
+
 
 class FloatInput:
     @classmethod
@@ -246,6 +254,10 @@ class FloatInput:
     def identity(self, value):
         return (value,)
 
+    @classmethod
+    def VALIDATE_INPUTS(s, value):
+        set_bentoml_output([(value,)])
+
 
 class BooleanInput:
     @classmethod
@@ -263,6 +275,10 @@ class BooleanInput:
     def identity(self, value):
         return (value,)
 
+    @classmethod
+    def VALIDATE_INPUTS(s, value):
+        set_bentoml_output([(value,)])
+
 
 class PathInput:
     @classmethod
@@ -280,6 +296,10 @@ class PathInput:
     def identity(self, path):
         return (path,)
 
+    @classmethod
+    def VALIDATE_INPUTS(s, path):
+        set_bentoml_output([(path,)])
+
 
 NODE_CLASS_MAPPINGS = {
     "BentoOutputPath": OutputPath,