From 3fc7c14d5bacae92fccaf00947b5aa478f1401ef Mon Sep 17 00:00:00 2001
From: krator18 <latentspacedirective@gmail.com>
Date: Fri, 6 Dec 2024 02:06:02 -0500
Subject: [PATCH] Overhaul of model loader. Added lazy switch node to allow for
 selecting from multiple types dynamically (checkpoint, diffusion, unet)

---
 TEXTURAIZER_load_data.py   | 23 +++++++++++++------
 TEXTURAIZER_other_nodes.py | 46 ++++++++++++++++++++++++++++++++++++++
 __init__.py                |  2 +-
 3 files changed, 63 insertions(+), 8 deletions(-)

diff --git a/TEXTURAIZER_load_data.py b/TEXTURAIZER_load_data.py
index 18ef451..2d11d66 100644
--- a/TEXTURAIZER_load_data.py
+++ b/TEXTURAIZER_load_data.py
@@ -265,6 +265,7 @@ class Texturaizer_GetJsonData:
         return (data_hash,)
 
 checkpoint_names = folder_paths.get_filename_list("checkpoints")
+diffusion_names = folder_paths.get_filename_list("diffusion_models")
 try:
     unet_names = folder_paths.get_filename_list("unet_gguf")
 except:
@@ -277,8 +278,8 @@ class Texturaizer_GetModelName(Texturaizer_GetJsonData):
     Extends the JSON retrieval functionality to return model-specific information.
     """
 
-    RETURN_TYPES = (checkpoint_names, unet_names, "STRING")
-    RETURN_NAMES = ("checkpoint_name", "unet_name", "data_hash")
+    RETURN_TYPES = (checkpoint_names, diffusion_names, unet_names, "INT", "STRING")
+    RETURN_NAMES = ("checkpoint_name", "diffusion_name", "unet_name", "model_type_idx", "data_hash")
     OUTPUT_TOOLTIPS = (
         "Diffusion checkpoint model name (stable diffusion).",
         "Diffusion unet model name (flux).",
@@ -294,9 +295,14 @@ class Texturaizer_GetModelName(Texturaizer_GetJsonData):
         """
         data = get_data(directory_optional, data_optional)
         scene_data = data.get("scene_info", {})
-        model = scene_data["ai_model"]
-        data_hash = calculate_data_hash(model)
-        return (model, model, data_hash)
+        checkpoint_model = scene_data.get("checkpoint_model", scene_data.get("ai_model"))
+        diffusion_model = scene_data.get("diffusion_model", scene_data.get("ai_model"))
+        unet_model = scene_data.get("unet_model", scene_data.get("ai_model"))
+        model_type = scene_data.get("model_type", "CHECKPOINT")
+        model_type_mapping = {"CHECKPOINT": 1, "DIFFUSION": 2, "UNET": 3}
+        model_type_index = model_type_mapping.get(model_type, 1)
+        data_hash = calculate_data_hash([checkpoint_model, diffusion_model, unet_model, model_type])
+        return (checkpoint_model, diffusion_model, unet_model, model_type_index, data_hash)
 
     @staticmethod
     def IS_CHANGED(directory_optional="", data_optional={}):
@@ -305,8 +311,11 @@ class Texturaizer_GetModelName(Texturaizer_GetJsonData):
         """
         data = get_data(directory_optional, data_optional)
         scene_data = data.get("scene_info", {})
-        model = scene_data["ai_model"]
-        data_hash = calculate_data_hash(model)
+        checkpoint_model = scene_data.get("checkpoint_model", scene_data.get("ai_model"))
+        diffusion_model = scene_data.get("diffusion_model", scene_data.get("ai_model"))
+        unet_model = scene_data.get("unet_model", scene_data.get("ai_model"))
+        model_type = scene_data.get("model_type", "CHECKPOINT")
+        data_hash = calculate_data_hash([checkpoint_model, diffusion_model, unet_model, model_type])
         return (data_hash,)
 
 try:
diff --git a/TEXTURAIZER_other_nodes.py b/TEXTURAIZER_other_nodes.py
index b1e2a57..7eab31d 100644
--- a/TEXTURAIZER_other_nodes.py
+++ b/TEXTURAIZER_other_nodes.py
@@ -31,10 +31,56 @@ class Texturaizer_SwitchAny:
         return (on_true,) if boolean else (on_false,)
 
 
+class Texturaizer_SwitchLazy:
+    """
+    Node that switches between three inputs based on an index.
+    Returns the selected input and blocks others.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "index": ("INT", {"default": 1, "min": 1, "max": 3, "tooltip": "Select which input to output (1-3)."}),
+            },
+            "optional": {
+                "input1": (any, {"lazy": True}),
+                "input2": (any, {"lazy": True}),
+                "input3": (any, {"lazy": True}),
+            }
+        }
+
+    CATEGORY = "Texturaizer"
+    RETURN_TYPES = (any, 'INT')
+    RETURN_NAMES = ("selected", 'index')
+    FUNCTION = "execute"
+
+    def check_lazy_status(self, *args, **kwargs):
+        """
+        Determines which input needs to be evaluated based on the index.
+        """
+        selected_index = int(kwargs['index'])
+        selected_input = f"input{selected_index}"
+        return [selected_input]
+
+    @staticmethod
+    def execute(*args, **kwargs):
+        selected_index = int(kwargs['index'])
+        selected_input = f"input{selected_index}"
+
+        if selected_input in kwargs and kwargs[selected_input] is not None:
+            return kwargs[selected_input], selected_index
+        else:
+            print(f"Execution blocked for unselected input: {selected_input}")
+            return None, selected_index
+
+
 NODE_CLASS_MAPPINGS = {
     "Texturaizer_SwitchAny": Texturaizer_SwitchAny,
+    "Texturaizer_SwitchLazy": Texturaizer_SwitchLazy,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "Texturaizer_SwitchAny": "Switch Any (Texturaizer)",
+    "Texturaizer_SwitchLazy": "Switch Lazy (Texturaizer)",
 }
diff --git a/__init__.py b/__init__.py
index 93b2ab6..ca2b481 100644
--- a/__init__.py
+++ b/__init__.py
@@ -9,7 +9,7 @@ Developed by Luke Kratsios and the LatentSpaceDirective.
 More about the creator: www.lukekratsios.com
 """
 
-__version__ = "0.0.1"
+__version__ = "0.0.2"
 
 import importlib