From aa63c612dea93925d9ad7ffb8e3094434a41d1af Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Thu, 27 Nov 2025 13:46:54 -0800
Subject: [PATCH] Progress on model loading

---
 modules/shared.py |  1 +
 modules/utils.py  | 25 +++++++++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/modules/shared.py b/modules/shared.py
index 2fbc205f..fda4ece6 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -22,6 +22,7 @@ lora_names = []
 
 # Image model variables
 image_model = None
+image_model_name = 'None'
 
 # Generation variables
 stop_everything = False
diff --git a/modules/utils.py b/modules/utils.py
index e8d23a02..5315d0f8 100644
--- a/modules/utils.py
+++ b/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)