From bf9b9e9fcdfecdfd5bea1e77eec9ea2502954a33 Mon Sep 17 00:00:00 2001
From: Blanchon <julien@blanchon.cc>
Date: Fri, 26 Jul 2024 02:31:40 +0200
Subject: [PATCH] Use ./models directory for luts file and add StringToLuts
 utility

---
 image.py | 36 +++++++++++++++++++++++++++++++++---
 1 file changed, 33 insertions(+), 3 deletions(-)

diff --git a/image.py b/image.py
index 1891e23..f62ba0f 100644
--- a/image.py
+++ b/image.py
@@ -15,6 +15,7 @@ import math
 import os
 import numpy as np
 import folder_paths
+from pathlib import Path
 import random
 
 """
@@ -926,7 +927,32 @@ class ImagePosterize:
         return(image,)
 
 
-LUTS_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "luts")
+script_directory = Path(__file__).parent
+folder_paths.add_model_folder_path("luts", (script_directory / "luts").as_posix())
+folder_paths.add_model_folder_path(
+    "luts", (Path(folder_paths.models_dir) / "luts").as_posix()
+)
+
+class StringToLuts:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "lut_file_name": (
+                    "STRING",
+                    {"multiline": False, "default": "lut_file.cube"},
+                ),
+            }
+        }
+
+    RETURN_TYPES = (folder_paths.get_filename_list("luts"),)
+    RETURN_NAMES = ("lut_file",)
+    FUNCTION = "execute"
+    CATEGORY = "essentials/image processing"
+
+    def execute(self, lut_file_name):
+        return (lut_file_name,)
+
 # From https://github.com/yoonsikp/pycubelut/blob/master/pycubelut.py (MIT license)
 class ImageApplyLUT:
     @classmethod
@@ -934,7 +960,7 @@ class ImageApplyLUT:
         return {
             "required": {
                 "image": ("IMAGE",),
-                "lut_file": (sorted([f for f in os.listdir(LUTS_DIR) if f.lower().endswith('.cube')]), ),
+                "lut_file": (folder_paths.get_filename_list("luts"),),
                 "gamma_correction": ("BOOLEAN", { "default": True }),
                 "clip_values": ("BOOLEAN", { "default": True }),
                 "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1 }),
@@ -948,8 +974,12 @@ class ImageApplyLUT:
     def execute(self, image, lut_file, gamma_correction, clip_values, strength):
         from colour.io.luts.iridas_cube import read_LUT_IridasCube
 
+        lut_file_path = folder_paths.get_full_path("luts", lut_file)
+        if not lut_file_path or not Path(lut_file_path).exists():
+            raise FileNotFoundError(f"Could not find LUT file: {lut_file_path}")
+
         device = image.device
-        lut = read_LUT_IridasCube(os.path.join(LUTS_DIR, lut_file))
+        lut = read_LUT_IridasCube(lut_file_path)
         lut.name = lut_file
 
         if clip_values: