Progress on model loading

modules/shared.py

@@ -22,6 +22,7 @@ lora_names = []
 
 # Image model variables
 image_model = None
+image_model_name = 'None'
 
 # Generation variables
 stop_everything = False

modules/utils.py

@@ -153,6 +153,31 @@ def get_available_models():
     return filtered_gguf_files + model_dirs
 
 
+def get_available_image_models():
+    model_dir = Path(shared.args.image_model_dir)
+
+    # Find directories with safetensors files
+    dirs_with_safetensors = set()
+    for item in os.listdir(model_dir):
+        item_path = model_dir / item
+        if item_path.is_dir():
+            if any(file.lower().endswith(('.safetensors', '.pt')) for file in os.listdir(item_path) if (item_path / file).is_file()):
+                dirs_with_safetensors.add(item)
+
+    # Find valid model directories
+    model_dirs = []
+    for item in os.listdir(model_dir):
+        item_path = model_dir / item
+        if not item_path.is_dir():
+            continue
+
+        model_dirs.append(item)
+
+    model_dirs = sorted(model_dirs, key=natural_keys)
+
+    return model_dirs
+
+
 def get_available_ggufs():
     model_list = []
     model_dir = Path(shared.args.model_dir)