From 10a6ed1e0966b816af7c2fb5c65d6a42a3d94c33 Mon Sep 17 00:00:00 2001
From: majorsauce <52931269+majorsauce@users.noreply.github.com>
Date: Wed, 27 Dec 2023 01:04:12 +0100
Subject: [PATCH] Allow YOLO models in same directory

To increase sharing of models with other UIs the YOLO models should not have to be split out into the bbox/ and segm/ folder.
This fix only comes into action if the initial lookup in bbox or segm folder failed, allowing YOLO model sharing with other UIs like Automatic1111
---
 impact/subpack_nodes.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/impact/subpack_nodes.py b/impact/subpack_nodes.py
index 13a14e1..4e4394e 100644
--- a/impact/subpack_nodes.py
+++ b/impact/subpack_nodes.py
@@ -24,6 +24,8 @@ def INPUT_TYPES(s):
 
     def doit(self, model_name):
         model_path = folder_paths.get_full_path("ultralytics", model_name)
+        if model_path is None and (model_name.startswith("bbox/") or model_name.startswith("segm/")):    # Allow segm/bbox models to be located in the same folder for shared model storage with other UIs
+            model_path = folder_paths.get_full_path("ultralytics", model_name[5:])
         model = subcore.load_yolo(model_path)
 
         if model_name.startswith("bbox"):