Fix model paths

Author: ionite34

src/inference_core_nodes/prompt_expansion/__init__.py

@@ -10,11 +10,16 @@
 def _add_folder_paths():
     from pathlib import Path
     import folder_paths
-
-    folder_paths.add_model_folder_path(
-        MODEL_FOLDER_NAME,
-        str(Path(folder_paths.models_dir).joinpath(MODEL_FOLDER_NAME)),
-    )
+    
+    folder_names_and_paths = folder_paths.folder_names_and_paths
+    
+    model_folder_path = str(Path(folder_paths.models_dir).joinpath(MODEL_FOLDER_NAME))
+    supported_extensions = {'.ckpt', '.pt', '.bin', '.pth', '.safetensors'}
+    
+    if MODEL_FOLDER_NAME in folder_names_and_paths:
+        folder_names_and_paths[MODEL_FOLDER_NAME][0].append(model_folder_path)
+    else:
+        folder_names_and_paths[MODEL_FOLDER_NAME] = ([model_folder_path], supported_extensions)
 
 
 _add_folder_paths()

src/inference_core_nodes/prompt_expansion/prompt_expansion.py

@@ -101,6 +101,11 @@ def expand_and_join(self, prompt: str, seed: int) -> str:
 @cache
 def load_expansion_runner(model_name: str):
     model_path = folder_paths.get_full_path(MODEL_FOLDER_NAME, model_name)
+    
+    # If model is a file, use the parent directory
+    if Path(model_path).is_file():
+        model_path = str(Path(model_path).parent)
+    
     return FooocusExpansion(model_path)