File size: 7,954 Bytes
c6d47bc 4626098 c6d47bc 4626098 c6d47bc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | 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})
# Mask queries: the last four coordinates before each <OBJ> inside a box.
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)
# top_k=1 is greedy decoding that still applies the reference repetition penalty (ignored by do_sample=False).
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()
|