# Checks to see if the model exists,

# If not, cwd is likely the folder just up from ComfyUI
    # There is likely a better way to do this, but this is the best to make sure those without issues now will still work after this fix.
    # I assume this is an issue because of the start.bat file for portable Windows users.
This commit is contained in:
Wesley Pyburn
2025-03-12 22:33:47 +02:00
committed by Gabi Dobocan
parent 44b7a47b14
commit 71d21095ff
+11
View File
@@ -67,6 +67,14 @@ def _xywh2cs(x, y, w, h, aspect_ratio):
scale = np.array([w, h], dtype=np.float32)
return center, scale
def check_model_path(model_path):
# Checks to see if the model exists, if not try adding ComfyUI/ to the start to fix possible errors on Windows (maybe others too)
if not os.path.exists(model_path):
new_model_path = os.path.join("ComfyUI", model_path)
if os.path.exists(new_model_path):
return new_model_path
return model_path
def generate(image, type, device):
num_classes = dataset_settings[type]['num_classes']
input_size = dataset_settings[type]['input_size']
@@ -78,6 +86,9 @@ def generate(image, type, device):
elif type == 'pascal':
model_path = 'models/schp/exp-schp-201908270938-pascal-person-part.pth'
# Check and adjust the model path if necessary
model_path = check_model_path(model_path)
model = networks.init_model('resnet101', num_classes=num_classes, pretrained=None)
state_dict = torch.load(model_path)['state_dict']
new_state_dict = OrderedDict()