# 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:
committed by
Gabi Dobocan
parent
44b7a47b14
commit
71d21095ff
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user