From 65cae62c7149048951fc70d601fc9c0d5434f69c Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 28 Aug 2023 16:49:06 -0400 Subject: [PATCH] No need to check filename extensions to detect shuffle controlnet. --- comfy/controlnet.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/comfy/controlnet.py b/comfy/controlnet.py index b18ab967..7098186f 100644 --- a/comfy/controlnet.py +++ b/comfy/controlnet.py @@ -1,5 +1,6 @@ import torch import math +import os import comfy.utils import comfy.model_management import comfy.model_detection @@ -386,7 +387,8 @@ def load_controlnet(ckpt_path, model=None): control_model = control_model.half() global_average_pooling = False - if ckpt_path.endswith("_shuffle.pth") or ckpt_path.endswith("_shuffle.safetensors") or ckpt_path.endswith("_shuffle_fp16.safetensors"): #TODO: smarter way of enabling global_average_pooling + filename = os.path.splitext(ckpt_path)[0] + if filename.endswith("_shuffle") or filename.endswith("_shuffle_fp16"): #TODO: smarter way of enabling global_average_pooling global_average_pooling = True control = ControlNet(control_model, global_average_pooling=global_average_pooling)