From 18c7854d764bcd1aa33b96eb63dab8ca3a160e71 Mon Sep 17 00:00:00 2001
From: Zuellni <123005779+Zuellni@users.noreply.github.com>
Date: Sat, 3 Feb 2024 15:13:30 +0100
Subject: [PATCH] Add missing import, look for paths recursively

---
 exllama.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/exllama.py b/exllama.py
index 6901bd5..8046e14 100644
--- a/exllama.py
+++ b/exllama.py
@@ -14,18 +14,19 @@ from exllamav2 import (
     ExLlamaV2Tokenizer,
 )
 from exllamav2.generator import ExLlamaV2Sampler, ExLlamaV2StreamingGenerator
-from folder_paths import folder_names_and_paths, get_folder_paths, models_dir
+from folder_paths import add_model_folder_path, get_folder_paths, models_dir
 
 
 class Loader:
     @classmethod
     def INPUT_TYPES(cls):
-        add_model_folder_path("llm", str(Path(models_dir).joinpath("llm")))
+        add_model_folder_path("llm", str(Path(models_dir) / "llm"))
+
         for folder in get_folder_paths("llm"):
-            if Path(folder).is_dir():
-                for path in Path(folder).glob("*/"):
-                    if (path / "config.json").is_file():
-                        cls._MODELS[path.name] = path
+            for path in Path(folder).rglob("*/"):
+                if (path / "config.json").is_file():
+                    parent = path.relative_to(folder).parent
+                    cls._MODELS[str(parent / path.name)] = path
 
         models = list(cls._MODELS.keys())
         default = models[0] if models else None