| import json |
| import re |
|
|
| import torch |
| from typing_extensions import override |
|
|
| import comfy.model_management |
| import comfy.sd |
| import comfy.supported_models_base |
| import comfy.utils |
| import folder_paths |
| from comfy.text_encoders.hunyuan_video import llama_detect |
| from comfy_api.latest import ComfyExtension, io |
|
|
| from .wevisone_model import DECODER_MEMORY, SYSTEM_PROMPT, WeVisOneTEModel, WeVisOneTokenizer |
|
|
| TASKS = ["detection", "referring", "point", "count_box", "count_point", "matting", "ocr", "ocr_grounding", "freeform"] |
| REF_RE = re.compile(r"<ref>(.*?)</ref>\s*<(box|point|polygon)>(.*?)(?:</(?:box|point|polygon)>|<ref>|$)", re.DOTALL) |
| COORD_RE = re.compile(r"<(\d{1,4})>") |
|
|
|
|
| def build_prompt(task, query): |
| labels = [line.strip() for line in query.splitlines() if line.strip()] |
| text = " ".join(labels) |
| phrase = text if not text or text[-1] in ".?!" else text + "." |
| if task == "detection": |
| return "Locate all the instances that match the following description: " + ". ".join(labels) + "." |
| if task == "referring": |
| return "Locate the instances that match the following description: " + phrase |
| if task == "point": |
| return "point to: " + ". ".join(labels) |
| if task == "count_box": |
| return "How many {} in the image? Please count them with boxes.".format(text or "objects") |
| if task == "count_point": |
| return "How many {} in the image?".format(text or "objects") |
| if task == "matting": |
| return "Locate the instances that match the following description and provide the alpha mattes: " + phrase |
| if task == "ocr": |
| return "detect all the word in the box format" |
| if task == "ocr_grounding": |
| return "Locate the following text in the image: " + text |
| return query |
|
|
|
|
| def parse_response(text, width, height): |
| boxes, points, seg_boxes = [], [], [] |
| for label, shape, body in REF_RE.findall(text): |
| label = label.strip() or "object" |
| for segment in body.split(","): |
| values = [int(v) for v in COORD_RE.findall(segment)] |
| if shape == "point": |
| for i in range(0, len(values) - 1, 2): |
| points.append({"x": round(values[i] / 1000 * width), "y": round(values[i + 1] / 1000 * height), "label": label}) |
| elif shape == "box": |
| for i in range(0, len(values) - 3, 4): |
| x1, y1, x2, y2 = values[i:i + 4] |
| if x2 > x1 and y2 > y1: |
| boxes.append({"x": x1 / 1000 * width, "y": y1 / 1000 * height, "width": (x2 - x1) / 1000 * width, "height": (y2 - y1) / 1000 * height, "label": label}) |
| |
| for body in re.findall(r"<box>(.*?)(?:</box>|$)", text, re.DOTALL): |
| for prefix in body.split("<OBJ>")[:-1]: |
| values = [int(v) for v in COORD_RE.findall(prefix)] |
| if len(values) >= 4: |
| x1, y1, x2, y2 = values[-4:] |
| if x2 > x1 and y2 > y1 and max(x1, y1, x2, y2) <= 1000: |
| seg_boxes.append([x1 / 1000, y1 / 1000, x2 / 1000, y2 / 1000]) |
| return boxes, points, seg_boxes |
|
|
|
|
| class WeVisOneVLLoader(io.ComfyNode): |
| @classmethod |
| def define_schema(cls): |
| return io.Schema( |
| node_id="WeVisOneVLLoader", |
| display_name="WeVisOne-VL Loader", |
| category="WeVisOne-VL", |
| inputs=[io.Combo.Input("clip_name", options=folder_paths.get_filename_list("text_encoders"))], |
| outputs=[io.Clip.Output()], |
| ) |
|
|
| @classmethod |
| def execute(cls, clip_name) -> io.NodeOutput: |
| sd = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("text_encoders", clip_name), safe_load=True) |
| target = comfy.supported_models_base.ClipTarget(WeVisOneTokenizer, WeVisOneTEModel) |
| detect = llama_detect(sd) |
| model_options = {"dtype": detect["dtype_llama"]} |
| if "llama_quantization_metadata" in detect: |
| model_options["quantization_metadata"] = detect["llama_quantization_metadata"] |
| clip = comfy.sd.CLIP(target, embedding_directory=folder_paths.get_folder_paths("embeddings"), parameters=comfy.utils.calculate_parameters(sd), |
| state_dict=[sd], model_options=model_options) |
| return io.NodeOutput(clip) |
|
|
|
|
| class WeVisOneVLPerceive(io.ComfyNode): |
| @classmethod |
| def define_schema(cls): |
| return io.Schema( |
| node_id="WeVisOneVLPerceive", |
| display_name="WeVisOne-VL Perceive", |
| category="WeVisOne-VL", |
| search_aliases=["detect", "grounding", "segment", "matting", "ocr", "count", "point"], |
| inputs=[ |
| io.Clip.Input("clip"), |
| io.Image.Input("image"), |
| io.Combo.Input("task", options=TASKS, default="matting"), |
| io.String.Input("query", multiline=True, default="", tooltip="Object labels (one per line for detection/point), a referring phrase, text to find, or the full prompt for freeform."), |
| io.Int.Input("max_length", default=1024, min=1, max=8192), |
| io.Boolean.Input("individual_masks", default=False, tooltip="Output one alpha matte per instance instead of their union."), |
| ], |
| outputs=[ |
| io.String.Output(display_name="text"), |
| io.BoundingBox.Output(display_name="bboxes"), |
| io.String.Output(display_name="points", tooltip="JSON [{\"x\": int, \"y\": int}, ...] in pixels, as SAM3 point prompts."), |
| io.Mask.Output(display_name="mask", tooltip="Alpha mattes for the matting task, or filled boxes for the other box tasks: their union per image, or every instance with individual_masks."), |
| ], |
| ) |
|
|
| @classmethod |
| def execute(cls, clip, image, task, query, max_length, individual_masks=False) -> io.NodeOutput: |
| height, width = image.shape[1:3] |
| prompt = build_prompt(task, query) |
| texts, all_boxes, all_points, masks = [], [], [], [] |
| for i in range(image.shape[0]): |
| frame = image[i:i + 1] |
| tokens = clip.tokenize(prompt, image=frame, system_prompt=SYSTEM_PROMPT) |
| |
| ids = clip.generate(tokens, do_sample=True, max_length=max_length, temperature=1.0, top_k=1, top_p=1.0, min_p=0.0, repetition_penalty=1.05, seed=0) |
| text = clip.decode(ids, skip_special_tokens=False).split("<|im_end|>", 1)[0].split("<|endoftext|>", 1)[0].strip() |
| boxes, points, seg_boxes = parse_response(text, width, height) |
| mattes = torch.zeros((0, height, width)) |
| if seg_boxes: |
| comfy.model_management.load_models_gpu([clip.patcher], memory_required=DECODER_MEMORY) |
| clip.cond_stage_model.set_clip_options({"execution_device": clip.patcher.load_device}) |
| mattes = clip.cond_stage_model.decode_masks(tokens, ids, frame, seg_boxes).cpu() |
| elif boxes: |
| mattes = torch.zeros((len(boxes), height, width)) |
| for m, b in zip(mattes, boxes): |
| m[round(b["y"]):round(b["y"] + b["height"]), round(b["x"]):round(b["x"] + b["width"])] = 1.0 |
| texts.append(text) |
| all_boxes.append(boxes) |
| all_points.extend({"x": p["x"], "y": p["y"]} for p in points) |
| masks.append(mattes if individual_masks else mattes.amax(0, keepdim=True) if len(mattes) else torch.zeros((1, height, width))) |
| return io.NodeOutput("\n".join(texts), all_boxes, json.dumps(all_points), torch.cat(masks)) |
|
|
|
|
| class WeVisOneVLExtension(ComfyExtension): |
| @override |
| async def get_node_list(self) -> list[type[io.ComfyNode]]: |
| return [WeVisOneVLLoader, WeVisOneVLPerceive] |
|
|
|
|
| async def comfy_entrypoint() -> WeVisOneVLExtension: |
| return WeVisOneVLExtension() |
|
|
