From 4b239d2d7533d285b8ab0996378a1251c2239693 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Sun, 18 Aug 2024 05:02:34 +0200
Subject: [PATCH] Fix name change / regular unet list crosstalk

#39
---
 nodes.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/nodes.py b/nodes.py
index 53bbc52..a4c51db 100644
--- a/nodes.py
+++ b/nodes.py
@@ -13,8 +13,10 @@ import folder_paths
 
 from .ops import GGMLTensor, GGMLOps
 
-# TODO: This causes gguf files to show up in the main unet loader
-folder_paths.folder_names_and_paths["unet"][1].add(".gguf")
+# Add a custom key for files ending in .gguf
+if "unet_gguf" not in folder_paths.folder_names_and_paths:
+    orig = folder_paths.folder_names_and_paths.get("diffusion_models", folder_paths.folder_names_and_paths.get("unet", [[], set()]))
+    folder_paths.folder_names_and_paths["unet_gguf"] = (orig[0], {".gguf"})
 
 def gguf_sd_loader(path):
     """
@@ -43,7 +45,7 @@ def gguf_sd_loader(path):
 class UnetLoaderGGUF:
     @classmethod
     def INPUT_TYPES(s):
-        unet_names = [x for x in folder_paths.get_filename_list("unet") if x.endswith(".gguf")]
+        unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")]
         return {
             "required": {
                 "unet_name": (unet_names,),