From 8ce6526562fd9c2d9b4e677d4afeb19c29463ca4 Mon Sep 17 00:00:00 2001
From: ComfyUI_FindModels <comfyui-findmodels@users.noreply.github.com>
Date: Fri, 12 Jun 2026 16:47:16 +0800
Subject: [PATCH] Stop guessing ambiguous model categories

---
 README.md                      |  2 ++
 find_models.py                 |  4 +---
 model_finder.py                | 20 ++++++++++++++++++++
 pyproject.toml                 |  2 +-
 tests/test_download_helpers.py |  2 ++
 tests/test_model_finder.py     | 21 +++++++++++++++++++++
 6 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index eeb07b8..e4a21cc 100644
--- a/README.md
+++ b/README.md
@@ -22,6 +22,8 @@ candidate direct-download URLs from Civitai and Hugging Face.
 - Searches the configured Quark Netdisk libraries directly for missing model filenames.
 - Previews and organizes existing model files into official ComfyUI model directories
   without overwriting files or moving anything outside the configured models root.
+- Leaves ambiguous checkpoint/diffusion model files in place instead of guessing their
+  architecture from names such as `flux`, `wan`, or `z_image`.
 - Does not move, delete, or overwrite model files.
 
 ## Installation
diff --git a/find_models.py b/find_models.py
index 6853c7c..bba8844 100644
--- a/find_models.py
+++ b/find_models.py
@@ -165,10 +165,8 @@ def _classify_existing_file(path: Path) -> str | None:
         return "upscale_models"
     if "/vae" in text or name.startswith("vae") or "_vae" in name:
         return "vae"
-    if any(term in text for term in ("diffusion_model", "/unet/", "flux", "wan")):
+    if any(term in text for term in ("diffusion_model", "/unet/")):
         return "diffusion_models"
-    if path.suffix.lower() in MODEL_EXTENSIONS:
-        return "checkpoints"
     return None
 
 
diff --git a/model_finder.py b/model_finder.py
index 2281305..60f8e61 100644
--- a/model_finder.py
+++ b/model_finder.py
@@ -29,6 +29,20 @@ STANDARD_LOADER_HINTS = (
     "upscale",
     "embedding",
 )
+MODEL_WIDGET_HINTS = (
+    "ckpt",
+    "checkpoint",
+    "model",
+    "lora",
+    "vae",
+    "control",
+    "clip",
+    "encoder",
+    "unet",
+    "diffusion",
+    "upscale",
+    "embedding",
+)
 
 CATEGORY_HINTS = (
     ("upscale_models", ("rife", "frame_interpolation", "film", "upscale", "esrgan", "super_resolution", "superresolution")),
@@ -97,6 +111,8 @@ def basename(value: str) -> str:
 
 
 def is_model_name(value: str) -> bool:
+    if "://" in value or value.startswith(("http:", "https:")):
+        return False
     clean = normalize_path(value).split("?", 1)[0].split("#", 1)[0]
     return PurePosixPath(clean).suffix.lower() in MODEL_EXTENSIONS
 
@@ -122,6 +138,10 @@ def extract_references(payload: Any) -> list[Reference]:
     def add(value: str, hint: str, node_id: Any = None, widget: Any = None, node_type: Any = None) -> None:
         if not is_model_name(value):
             return
+        widget_hint = re.sub(r"[^a-z0-9_]+", "", str(widget or "").lower())
+        node_hint = re.sub(r"[^a-z0-9_]+", "", str(node_type or "").lower())
+        if node_id is not None and not any(term in f"{widget_hint} {node_hint}" for term in MODEL_WIDGET_HINTS):
+            return
         ref = Reference(
             name=normalize_path(value),
             category=classify(hint),
diff --git a/pyproject.toml b/pyproject.toml
index ad4f279..d3c8806 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "comfyui-findmodels"
-version = "1.6.0"
+version = "1.6.1"
 description = "Find and adapt missing models in the current ComfyUI workflow"
 readme = "README.md"
 requires-python = ">=3.10"
diff --git a/tests/test_download_helpers.py b/tests/test_download_helpers.py
index 562bba2..4f80f51 100644
--- a/tests/test_download_helpers.py
+++ b/tests/test_download_helpers.py
@@ -103,6 +103,8 @@ class DownloadHelperTests(unittest.TestCase):
         self.assertEqual(_classify_existing_file(Path("misc/loras/hero.safetensors")), "loras")
         self.assertEqual(_classify_existing_file(Path("misc/rife47.pth")), "upscale_models")
         self.assertEqual(_classify_existing_file(Path("misc/vae_model.safetensors")), "vae")
+        self.assertIsNone(_classify_existing_file(Path("checkpoints/z_image_turbo_bf16.safetensors")))
+        self.assertIsNone(_classify_existing_file(Path("misc/unknown_model.safetensors")))
 
     def test_organize_plan_moves_only_misplaced_files(self):
         root = Path(folder_paths.models_dir)
diff --git a/tests/test_model_finder.py b/tests/test_model_finder.py
index 332dcbc..afb7cee 100644
--- a/tests/test_model_finder.py
+++ b/tests/test_model_finder.py
@@ -98,6 +98,27 @@ class ModelFinderTests(unittest.TestCase):
         self.assertEqual(result["summary"]["unresolved"], 1)
         self.assertEqual(len(result["models"]), 1)
 
+    def test_ignores_model_urls_in_text_widgets(self):
+        payload = {
+            "nodes": [{
+                "id": 12,
+                "type": "TextNode",
+                "widgets": [{
+                    "name": "text",
+                    "value": "https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/model.safetensors",
+                }],
+            }]
+        }
+        result = analyze(payload, lambda category: [])
+        self.assertEqual(result["models"], [])
+
+    def test_ignores_filename_in_non_model_text_widget(self):
+        payload = {
+            "nodes": [{"id": 13, "type": "TextNode", "widgets": [{"name": "text", "value": "model.safetensors"}]}]
+        }
+        result = analyze(payload, lambda category: [])
+        self.assertEqual(result["models"], [])
+
 
 if __name__ == "__main__":
     unittest.main()