diff --git a/Custom_ComfyUI_YoloWorld_EfficientSAM/YOLO_WORLD_EfficientSAM.py b/Custom_ComfyUI_YoloWorld_EfficientSAM/YOLO_WORLD_EfficientSAM.py index 1f58ff3..46ffe50 100644 --- a/Custom_ComfyUI_YoloWorld_EfficientSAM/YOLO_WORLD_EfficientSAM.py +++ b/Custom_ComfyUI_YoloWorld_EfficientSAM/YOLO_WORLD_EfficientSAM.py @@ -210,7 +210,7 @@ NODE_CLASS_MAPPINGS = { } NODE_DISPLAY_NAME_MAPPINGS = { - "Yoloworld_ModelLoader_Zho": "🔎Yoloworld Model Loader", - "ESAM_ModelLoader_Zho": "🔎ESAM Model Loader", + "Yoloworld_ModelLoader_Zho": "🔎Custom Yoloworld Model Loader", + "ESAM_ModelLoader_Zho": "🔎Custom ESAM Model Loader", "Yoloworld_ESAM_Zho": "🔎Custom Yoloworld ESAM", } diff --git a/Custom_ComfyUI_YoloWorld_EfficientSAM/YOLO_WORLD_SEGS.py b/Custom_ComfyUI_YoloWorld_EfficientSAM/YOLO_WORLD_SEGS.py deleted file mode 100644 index ea86126..0000000 --- a/Custom_ComfyUI_YoloWorld_EfficientSAM/YOLO_WORLD_SEGS.py +++ /dev/null @@ -1,322 +0,0 @@ -from .YOLO_WORLD_EfficientSAM import * -from collections import namedtuple -from PIL import Image - -SEG = namedtuple("SEG", - ['cropped_image', 'cropped_mask', 'confidence', 'crop_region', 'bbox', 'label', 'control_net_wrapper'], - defaults=[None]) - - -def crop_ndarray4(npimg, crop_region): - x1 = crop_region[0] - y1 = crop_region[1] - x2 = crop_region[2] - y2 = crop_region[3] - - cropped = npimg[:, y1:y2, x1:x2, :] - - return cropped - - -def crop_ndarray2(npimg, crop_region): - x1 = crop_region[0] - y1 = crop_region[1] - x2 = crop_region[2] - y2 = crop_region[3] - - cropped = npimg[y1:y2, x1:x2] - - return cropped - - -crop_tensor4 = crop_ndarray4 - - -def crop_image(image, crop_region): - return crop_tensor4(image, crop_region) - - -def create_segmasks(results): - bboxs = results[1] - segms = results[2] - confidence = results[3] - - results = [] - for i in range(len(segms)): - item = (bboxs[i], segms[i].astype(np.float32), confidence[i]) - results.append(item) - return results - - -def dilate_masks(segmasks, dilation_factor, iter=1): - if dilation_factor == 0: - return segmasks - - dilated_masks = [] - kernel = np.ones((abs(dilation_factor), abs(dilation_factor)), np.uint8) - - for i in range(len(segmasks)): - cv2_mask = segmasks[i][1] - - if dilation_factor > 0: - dilated_mask = cv2.dilate(cv2_mask, kernel, iter) - else: - dilated_mask = cv2.erode(cv2_mask, kernel, iter) - - item = (segmasks[i][0], dilated_mask, segmasks[i][2]) - dilated_masks.append(item) - - return dilated_masks - - -def make_crop_region(w, h, bbox, crop_factor, crop_min_size=None): - x1 = bbox[0] - y1 = bbox[1] - x2 = bbox[2] - y2 = bbox[3] - - bbox_w = x2 - x1 - bbox_h = y2 - y1 - - crop_w = bbox_w * crop_factor - crop_h = bbox_h * crop_factor - - if crop_min_size is not None: - crop_w = max(crop_min_size, crop_w) - crop_h = max(crop_min_size, crop_h) - - kernel_x = x1 + bbox_w / 2 - kernel_y = y1 + bbox_h / 2 - - new_x1 = int(kernel_x - crop_w / 2) - new_y1 = int(kernel_y - crop_h / 2) - - # make sure position in (w,h) - new_x1, new_x2 = normalize_region(w, new_x1, crop_w) - new_y1, new_y2 = normalize_region(h, new_y1, crop_h) - - return [new_x1, new_y1, new_x2, new_y2] - - -def normalize_region(limit, startp, size): - if startp < 0: - new_endp = min(limit, size) - new_startp = 0 - elif startp + size > limit: - new_startp = max(0, limit - size) - new_endp = limit - else: - new_startp = startp - new_endp = min(limit, startp+size) - - return int(new_startp), int(new_endp) - - -def combine_masks(masks): - if len(masks) == 0: - return None - else: - initial_cv2_mask = np.array(masks[0][1]) - combined_cv2_mask = initial_cv2_mask - - for i in range(1, len(masks)): - cv2_mask = np.array(masks[i][1]) - - if combined_cv2_mask.shape == cv2_mask.shape: - combined_cv2_mask = cv2.bitwise_or(combined_cv2_mask, cv2_mask) - else: - # do nothing - incompatible mask - pass - - mask = torch.from_numpy(combined_cv2_mask) - return mask - - -def inference_bbox(yolo_world_model, categories, iou_threshold, with_class_agnostic_nms, image, confidence): - img = np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8) - yolo_world_model.set_classes(categories) - - results = yolo_world_model.infer(img, confidence=confidence) - detections = sv.Detections.from_inference(results) - detections = detections.with_nms(class_agnostic=with_class_agnostic_nms, threshold=iou_threshold) - - bboxes = detections.xyxy - cv2_image = np.array(img) - if len(cv2_image.shape) == 3: - cv2_image = cv2_image[:, :, ::-1].copy() # Convert RGB to BGR for cv2 processing - else: - # Handle the grayscale image here - # For example, you might want to convert it to a 3-channel grayscale image for consistency: - cv2_image = cv2.cvtColor(cv2_image, cv2.COLOR_GRAY2BGR) - cv2_gray = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2GRAY) - - segms = [] - for x0, y0, x1, y1 in bboxes: - cv2_mask = np.zeros(cv2_gray.shape, np.uint8) - cv2.rectangle(cv2_mask, (int(x0), int(y0)), (int(x1), int(y1)), 255, -1) - cv2_mask_bool = cv2_mask.astype(bool) - segms.append(cv2_mask_bool) - - n, m = bboxes.shape - if n == 0: - return [[], [], [], []] - - results = [[], [], [], []] - for i in range(len(bboxes)): - results[0].append(detections.data['class_name'][i]) - results[1].append(bboxes[i]) - results[2].append(segms[i]) - results[3].append(detections.confidence[i]) - - return results - - -def inference_segm(yolo_world_model, esam_model, categories, iou_threshold, with_class_agnostic_nms, image, confidence): - img = np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8) - yolo_world_model.set_classes(categories) - results = yolo_world_model.infer(img, confidence=confidence) - detections = sv.Detections.from_inference(results) - detections = detections.with_nms(class_agnostic=with_class_agnostic_nms, threshold=iou_threshold) - segms = inference_with_boxes( - image=img, - xyxy=detections.xyxy, - model=esam_model, - device=DEVICE - ) - - bboxes = detections.xyxy - n, m = bboxes.shape - if n == 0: - return [[], [], [], []] - - results = [[], [], [], []] - for i in range(len(bboxes)): - results[0].append(detections.data['class_name'][i]) - results[1].append(bboxes[i]) - - mask = torch.from_numpy(segms[i]) - scaled_mask = torch.nn.functional.interpolate(mask.float().unsqueeze(0).unsqueeze(0), size=(img.shape[0], img.shape[1]), mode='bilinear', align_corners=False) - scaled_mask = scaled_mask.squeeze().squeeze() - - results[2].append(scaled_mask.numpy()) - results[3].append(detections.confidence[i]) - - return results - - -class YoloworldBboxDetector: - def __init__(self, yolo_world_model, categories, iou_threshold, with_class_agnostic_nms): - self.yolo_world_model = yolo_world_model - self.categories = process_categories(categories) - self.iou_threshold = iou_threshold - self.with_class_agnostic_nms = with_class_agnostic_nms - - def detect(self, image, threshold, dilation, crop_factor, drop_size=1, detailer_hook=None, esam_model=None): - drop_size = max(drop_size, 1) - if esam_model is None: - detected_results = inference_bbox(self.yolo_world_model, self.categories, self.iou_threshold, self.with_class_agnostic_nms, image, threshold) - else: - detected_results = inference_segm(self.yolo_world_model, esam_model, self.categories, self.iou_threshold, self.with_class_agnostic_nms, image, threshold) - - segmasks = create_segmasks(detected_results) - - if dilation > 0: - segmasks = dilate_masks(segmasks, dilation) - - items = [] - h = image.shape[1] - w = image.shape[2] - - for x, label in zip(segmasks, detected_results[0]): - item_bbox = x[0] - item_mask = x[1] - - y1, x1, y2, x2 = item_bbox - - if x2 - x1 > drop_size and y2 - y1 > drop_size: # minimum dimension must be (2,2) to avoid squeeze issue - crop_region = make_crop_region(w, h, item_bbox, crop_factor) - - if detailer_hook is not None: - crop_region = detailer_hook.post_crop_region(w, h, item_bbox, crop_region) - - cropped_image = crop_image(image, crop_region) - cropped_mask = crop_ndarray2(item_mask, crop_region) - confidence = x[2] - - item = SEG(cropped_image, cropped_mask, confidence, crop_region, item_bbox, label, None) - - items.append(item) - - shape = image.shape[1], image.shape[2] - segs = shape, items - - if detailer_hook is not None and hasattr(detailer_hook, "post_detection"): - segs = detailer_hook.post_detection(segs) - - return segs - - def detect_combined(self, image, threshold, dilation): - detected_results = inference_bbox(self.yolo_world_model, self.categories, self.iou_threshold, self.with_class_agnostic_nms, image, threshold) - segmasks = create_segmasks(detected_results) - if dilation > 0: - segmasks = dilate_masks(segmasks, dilation) - - return combine_masks(segmasks) - - -class YoloworldSegmDetector: - def __init__(self, bbox_detector, esam_model): - self.bbox_detector = bbox_detector - self.esam_model = esam_model - - def detect(self, image, threshold, dilation, crop_factor, drop_size=1, detailer_hook=None): - return self.bbox_detector.detect(image, threshold, dilation, crop_factor, drop_size, detailer_hook=detailer_hook, esam_model=self.esam_model) - - def detect_combined(self, image, threshold, dilation): - bb = self.bbox_detector - detected_results = inference_segm(bb.yolo_world_model, self.esam_model, bb.categories, bb.iou_threshold, bb.with_class_agnostic_nms, image, threshold) - segmasks = create_segmasks(detected_results) - if dilation > 0: - segmasks = dilate_masks(segmasks, dilation) - - return combine_masks(segmasks) - - -class Yoloworld_ESAM_DetectorProvider_Zho: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "yolo_world_model": ("YOLOWORLDMODEL",), - "categories": ("STRING", {"default": "", "placeholder": "Please enter the objects to be detected separated by commas.", "multiline": True}), - "iou_threshold": ("FLOAT", {"default": 0.1, "min": 0, "max": 1, "step": 0.01}), - "with_class_agnostic_nms": ("BOOLEAN", {"default": False}), - }, - - "optional": { - "esam_model_opt": ("ESAMMODEL",), - } - } - - RETURN_TYPES = ("BBOX_DETECTOR", "SEGM_DETECTOR") - FUNCTION = "doit" - - CATEGORY = "🔎YOLOWORLD_ESAM" - - def doit(self, yolo_world_model, categories, iou_threshold, with_class_agnostic_nms, esam_model_opt=None): - bbox_detector = YoloworldBboxDetector(yolo_world_model, categories, iou_threshold, with_class_agnostic_nms) - if esam_model_opt is not None: - segm_detector = YoloworldSegmDetector(bbox_detector, esam_model_opt) - else: - segm_detector = None - - return bbox_detector, segm_detector - - -NODE_CLASS_MAPPINGS = { - "Yoloworld_ESAM_DetectorProvider_Zho": Yoloworld_ESAM_DetectorProvider_Zho, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "Yoloworld_ESAM_DetectorProvider_Zho": "🔎Yoloworld ESAM Detector Provider", -} diff --git a/Custom_comfyui_llm_assistant/AssistantNodes.py b/Custom_comfyui_llm_assistant/AssistantNodes.py index 70882bb..5b71fba 100644 --- a/Custom_comfyui_llm_assistant/AssistantNodes.py +++ b/Custom_comfyui_llm_assistant/AssistantNodes.py @@ -117,224 +117,11 @@ class GenerateStableDiffsutionPromptLLM: return (response,) -class TranslateTextLLM: - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls): - data = load_config() - template_system = "" - template_user = "" - models = [] - - if data: - _prompt = data["default_translate_prompt"] - if _prompt: - _system = _prompt["system"] - _user = _prompt["user"] - if _system: - template_system = _system - if _user: - template_user = _user - _models = data["models"] - if isinstance(_models, list) and len(_models) > 0: - models = _models - default_model = data["default_model"] - if default_model and len(_models) > 0: - default_model_index = models.index(default_model) - if default_model_index > 0: - models.pop(default_model_index) - models.insert(0, default_model) - return { - "required": { - "prompt": ("STRING", { - "dynamicPrompts": False, - "multiline": True, - "default": "" - }), - - "template_system": ("STRING", { - "dynamicPrompts": False, - "multiline": True, - "default": template_system, - "display": "textarea", - - }), - "template_user": ("STRING", { - "dynamicPrompts": False, - "multiline": True, - "default": template_user, - "display": "textarea" - }), - "stop": ("STRING", { - "dynamicPrompts": False, - "multiline": False, - "default": "" - }), - "response_pattern": ("STRING", { - "dynamicPrompts": False, - "multiline": False, - "default": "" - }), - "temperature": ("FLOAT", { - "default": 1.0, - "min": 0.0, - "max": 1.0, - "step": 0.01, - "round": 0.01, - "display": "number" - }), - - "max_tokens": ("INT", { - "default": 300, - "min": -1, - "max": 2048, - "display": "number" - }), - "model_name": (models, { - "default": default_model, - "display": "select" - }) - } - } - - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("text",) - - FUNCTION = "translateText" - - CATEGORY = "LLM" - - def translateText(self, prompt, template_system, template_user, stop, response_pattern, temperature, max_tokens, model_name): - config = load_config() - openai_config = config["openai"] - if openai_config is None: - return (prompt,) - api_base = openai_config["api_base"] - api_key = openai_config["api_key"] - if not (is_valid_string(api_base) and is_valid_string(api_key)): - return (prompt,) - _prompt = prompt - if is_valid_string(template_user): - _prompt = template_user.format(prompt) - response = get_completion(_prompt, response_pattern, api_base, api_key, temperature, template_system, - max_tokens, stop, model_name) - return (response,) - - -class ChatWithLLM: - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls): - data = load_config() - template_system = "" - template_user = "" - models = [] - - if data: - _prompt = data["default_chat_prompt"] - if _prompt: - _system = _prompt["system"] - _user = _prompt["user"] - if _system: - template_system = _system - if _user: - template_user = _user - _models = data["models"] - if isinstance(_models, list) and len(_models) > 0: - models = _models - default_model = data["default_model"] - if default_model and len(_models) > 0: - default_model_index = models.index(default_model) - if default_model_index > 0: - models.pop(default_model_index) - models.insert(0, default_model) - return { - "required": { - "prompt": ("STRING", { - "dynamicPrompts": False, - "multiline": True, - "default": "" - }), - "template_system": ("STRING", { - "dynamicPrompts": False, - "multiline": True, - "default": template_system, - "display": "textarea" - }), - "template_user": ("STRING", { - "dynamicPrompts": False, - "multiline": True, - "default": template_user, - "display": "textarea" - }), - "stop": ("STRING", { - "dynamicPrompts": False, - "multiline": False, - "default": "" - }), - "response_pattern": ("STRING", { - "dynamicPrompts": False, - "multiline": False, - "default": "" - }), - "temperature": ("FLOAT", { - "default": 1.0, - "min": 0.0, - "max": 1.0, - "step": 0.01, - "round": 0.01, - "display": "number" - }), - - "max_tokens": ("INT", { - "default": 300, - "min": -1, - "max": 2048, - "display": "number" - }), - "model_name": (models, { - "default": default_model, - "display": "select" - }) - } - } - - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("text",) - - FUNCTION = "chatLLM" - - CATEGORY = "LLM" - - def chatLLM(self, prompt, template_system, template_user, stop, response_pattern, temperature, max_tokens, model_name): - config = load_config() - openai_config = config["openai"] - if openai_config is None: - return (prompt,) - api_base = openai_config["api_base"] - api_key = openai_config["api_key"] - if not (is_valid_string(api_base) and is_valid_string(api_key)): - return (prompt,) - _prompt = prompt - if is_valid_string(template_user): - _prompt = template_user.format(prompt) - response = get_completion(_prompt, response_pattern, api_base, api_key, temperature, template_system, - max_tokens, stop, model_name) - return (response,) - NODE_CLASS_MAPPINGS = { "Generate Stable Diffsution Prompt With LLM duci": GenerateStableDiffsutionPromptLLM, - "Translate Text With LLM": TranslateTextLLM, - "Chat With LLM": ChatWithLLM, } NODE_DISPLAY_NAME_MAPPINGS = { "Generate Stable Diffsution Prompt With LLM duci": "Custom Generate Stable Diffsution Prompt With LLM", - "Translate Text With LLM": "Translate Text With LLM", - "Chat With LLM": "Chat With LLM", }