diff --git a/README.md b/README.md index 04dd351..b8c14e3 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,8 @@ Includes nodes for loading and running VLMs: Along with some utility nodes for working with text: - Parse Bounding Boxes +- Parse Points +- Plot Points - Regex Split String - Regex Search - Regex Find All @@ -53,18 +55,26 @@ Unfortunately, the Pixtral nf4 model has considerably degraded performance on so ## Examples Example Pixtral image captioning (not saving the output to a text file in this example): -![Example Pixtral image captioning workflow](pixtral_caption_example.jpg) +![Example Pixtral image captioning workflow](examples/pixtral_caption_example.jpg) Both models should work very well for image captioning, even in 4-bit quantization. You can also customize your captioning instructions. Example Pixtral image comparison: -![Example Pixtral image comparison workflow](pixtral_comparison_example.jpg) +![Example Pixtral image comparison workflow](examples/pixtral_comparison_example.jpg) I haven't been able to get image comparison to work well at all with Llama Vision. It doesn't give any errors, but the multi-image understanding just isn't there. The image tokens have to be **before** the question/instruction and consecutive for the model to even be able to see both images at once (I found this out by looking at the image preprocessor cross-attention implementation), and even then, it seems to randomly mix up which is the first/second, left/right, the colors between them and other details. It doesn't seem usable for purposes involving two images in the same message, in my opinion. Not sure whether the non-quantized model is better at this. Since Pixtral directly tokenizes the input images, it's able to handle them inline in the context, with any number of images of any aspect ratio, but it's limited by token lengths, since each image can be around 1000 tokens. Example Llama Vision object detection with bounding box: -![Example Llama Vision object detection with bounding box workflow](llama_vision_bounding_box_example.jpg) +![Example Llama Vision object detection with bounding box workflow](examples/llama_vision_bounding_box_example.jpg) Both models kind of work for this, but not that well. They definitely have some understanding of the positions of objects in the image, though. Maybe it needs a better prompt. Or a non-quantized model. Or a finetune. But it does sometimes work. + +Example Molmo counting: +![Example Molmo counting example](examples/molmo_counting_example.jpg) + +Example Molmo pointing, with labels: +![Example Molmo pointing example](examples/molmo_multi_pointing_example.jpg) + +I wasn't able to get it to point at both objects with a single prompt for some reason (it would just assign both labels to both points), but splitting it into two simple prompts like this isn't too bad. \ No newline at end of file diff --git a/llama_vision_bounding_box_example.jpg b/examples/llama_vision_bounding_box_example.jpg similarity index 100% rename from llama_vision_bounding_box_example.jpg rename to examples/llama_vision_bounding_box_example.jpg diff --git a/llama_vision_bounding_box_workflow.json b/examples/llama_vision_bounding_box_workflow.json similarity index 100% rename from llama_vision_bounding_box_workflow.json rename to examples/llama_vision_bounding_box_workflow.json diff --git a/examples/molmo_count_example.jpg b/examples/molmo_count_example.jpg new file mode 100644 index 0000000..031ce2d Binary files /dev/null and b/examples/molmo_count_example.jpg differ diff --git a/examples/molmo_counting_workflow.json b/examples/molmo_counting_workflow.json new file mode 100644 index 0000000..c0a1f0b --- /dev/null +++ b/examples/molmo_counting_workflow.json @@ -0,0 +1,363 @@ +{ + "last_node_id": 58, + "last_link_id": 114, + "nodes": [ + { + "id": 56, + "type": "ParsePoints", + "pos": { + "0": 1086, + "1": 335 + }, + "size": [ + 315, + 122 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "string", + "type": "STRING", + "link": 111, + "widget": { + "name": "string" + } + } + ], + "outputs": [ + { + "name": "POINT", + "type": "POINT", + "links": [ + 112 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "STRING", + "type": "STRING", + "links": null, + "shape": 3 + }, + { + "name": "STRING", + "type": "STRING", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "ParsePoints" + }, + "widgets_values": [ + "", + "" + ] + }, + { + "id": 54, + "type": "MolmoModelLoader", + "pos": { + "0": 313, + "1": 248 + }, + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VISION_MODEL", + "type": "VISION_MODEL", + "links": [ + 108 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "MolmoModelLoader" + }, + "widgets_values": [ + "molmo-7B-D-bnb-4bit" + ] + }, + { + "id": 58, + "type": "PreviewImage", + "pos": { + "0": 1771, + "1": 334 + }, + "size": [ + 382.1099039489918, + 403.9147862091281 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 113 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 4, + "type": "LoadImage", + "pos": { + "0": 324, + "1": 358 + }, + "size": { + "0": 294.7367248535156, + "1": 375.8291931152344 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 109, + 114 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "slot_index": 1, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "Flux_00207_ (1).png", + "image" + ] + }, + { + "id": 57, + "type": "PlotPoints", + "pos": { + "0": 1429, + "1": 340 + }, + "size": { + "0": 315, + "1": 150 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "points", + "type": "POINT", + "link": 112 + }, + { + "name": "image", + "type": "IMAGE", + "link": 114 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 113 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "PlotPoints" + }, + "widgets_values": [ + 10, + 40, + "#ff00ff", + "" + ] + }, + { + "id": 55, + "type": "MolmoGenerateText", + "pos": { + "0": 656, + "1": 301 + }, + "size": { + "0": 400, + "1": 312 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "molmo_model", + "type": "VISION_MODEL", + "link": 108 + }, + { + "name": "images", + "type": "IMAGE", + "link": 109 + } + ], + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": [ + 110, + 111 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "MolmoGenerateText" + }, + "widgets_values": [ + "Count the people in this image.", + 512, + true, + 0.3, + 0.9, + 40, + "<|endoftext|>", + 2130269981, + "randomize", + false + ] + }, + { + "id": 3, + "type": "Display Any (rgthree)", + "pos": { + "0": 1109, + "1": 6 + }, + "size": [ + 608.1341091642305, + 278.6934670890276 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "source", + "type": "*", + "link": 110, + "dir": 3 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "Display Any (rgthree)" + }, + "widgets_values": [ + "" + ] + } + ], + "links": [ + [ + 108, + 54, + 0, + 55, + 0, + "VISION_MODEL" + ], + [ + 109, + 4, + 0, + 55, + 1, + "IMAGE" + ], + [ + 110, + 55, + 0, + 3, + 0, + "*" + ], + [ + 111, + 55, + 0, + 56, + 0, + "STRING" + ], + [ + 112, + 56, + 0, + 57, + 0, + "POINT" + ], + [ + 113, + 57, + 0, + 58, + 0, + "IMAGE" + ], + [ + 114, + 4, + 0, + 57, + 1, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1.2284597357367266, + "offset": [ + -254.51223899431685, + 74.38924093465424 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/examples/molmo_multi_pointing_example.jpg b/examples/molmo_multi_pointing_example.jpg new file mode 100644 index 0000000..2c8b7fa Binary files /dev/null and b/examples/molmo_multi_pointing_example.jpg differ diff --git a/pixtral_caption_example.jpg b/examples/pixtral_caption_example.jpg similarity index 100% rename from pixtral_caption_example.jpg rename to examples/pixtral_caption_example.jpg diff --git a/pixtral_caption_workflow.json b/examples/pixtral_caption_workflow.json similarity index 100% rename from pixtral_caption_workflow.json rename to examples/pixtral_caption_workflow.json diff --git a/pixtral_comparison_example.jpg b/examples/pixtral_comparison_example.jpg similarity index 100% rename from pixtral_comparison_example.jpg rename to examples/pixtral_comparison_example.jpg diff --git a/nodes.py b/nodes.py index 63a5a12..608a9c2 100644 --- a/nodes.py +++ b/nodes.py @@ -9,14 +9,17 @@ from transformers import ( AutoProcessor, BitsAndBytesConfig, GenerationConfig, + StopStringCriteria, set_seed ) from torchvision.transforms.functional import to_pil_image +import numpy as np +import torch import json import os from pathlib import Path -from PIL import Image +from PIL import Image, ImageDraw import re import time @@ -110,6 +113,7 @@ class PixtralGenerateText: "top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.1}), "top_k": ("INT", {"default": 40, "min": 1}), "repetition_penalty": ("FLOAT", {"default": 1.1}), + "stop_strings": ("STRING", {"default": ""}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}), "include_prompt_in_output": ("BOOLEAN", {"default": False}), } @@ -120,7 +124,7 @@ class PixtralGenerateText: CATEGORY = "PixtralLlamaVision/Pixtral" TITLE = "Generate Text with Pixtral" - def generate_text(self, pixtral_model, images, prompt, max_new_tokens, do_sample, temperature, top_p, top_k, repetition_penalty, seed, include_prompt_in_output): + def generate_text(self, pixtral_model, images, prompt, max_new_tokens, do_sample, temperature, top_p, top_k, repetition_penalty, stop_strings, seed, include_prompt_in_output): device = pixtral_model['model'].device # I'm sure there is a way to do this without converting back to numpy and then PIL... # Pixtral requires PIL input for some reason, and the to_pil_image function requires channels to be the first dimension for a Tensor but the last dimension for a numpy array... Yeah idk @@ -129,6 +133,7 @@ class PixtralGenerateText: inputs = pixtral_model['processor'](images=image_list, text=prompt, return_tensors="pt").to(device) prompt_tokens = len(inputs['input_ids'][0]) print(f"Prompt tokens: {prompt_tokens}") + stop_strings_list = stop_strings.split(",") set_seed(seed) t0 = time.time() generate_ids = pixtral_model['model'].generate( @@ -140,7 +145,8 @@ class PixtralGenerateText: top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty, - ) + ), + stopping_criteria=[StopStringCriteria(tokenizer=pixtral_model['processor'].tokenizer, stop_strings=stop_strings_list)], ) t1 = time.time() total_time = t1 - t0 @@ -203,6 +209,7 @@ class LlamaVisionGenerateText: "top_k": ("INT", {"default": 40, "min": 1}), # For some reason, including this causes the CUDA kernel to fail catastrophically? Didn't have this issue with Pixtral #"repetition_penalty": ("FLOAT", {"default": 1.1}), + "stop_strings": ("STRING", {"default": "<|eot_id|>"}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}), "include_prompt_in_output": ("BOOLEAN", {"default": False}), } @@ -213,7 +220,7 @@ class LlamaVisionGenerateText: CATEGORY = "PixtralLlamaVision/LlamaVision" TITLE = "Generate Text with Llama Vision" - def generate_text(self, llama_vision_model, images, prompt, max_new_tokens, do_sample, temperature, top_p, top_k, seed, include_prompt_in_output): + def generate_text(self, llama_vision_model, images, prompt, max_new_tokens, do_sample, temperature, top_p, top_k, stop_strings, seed, include_prompt_in_output): device = llama_vision_model['model'].device # I'm sure there is a way to do this without converting back to numpy and then PIL... # Llama Vision also requires PIL input for some reason, and the to_pil_image function requires channels to be the first dimension for a Tensor but the last dimension for a numpy array... Yeah idk @@ -222,6 +229,7 @@ class LlamaVisionGenerateText: inputs = llama_vision_model['processor'](images=image_list, text=prompt, return_tensors="pt").to(device) prompt_tokens = len(inputs['input_ids'][0]) print(f"Prompt tokens: {prompt_tokens}") + stop_strings_list = stop_strings.split(",") set_seed(seed) t0 = time.time() generate_ids = llama_vision_model['model'].generate( @@ -234,6 +242,7 @@ class LlamaVisionGenerateText: top_k=top_k, #repetition_penalty=repetition_penalty, ), + stopping_criteria=[StopStringCriteria(tokenizer=llama_vision_model['processor'].tokenizer, stop_strings=stop_strings_list)], ) t1 = time.time() total_time = t1 - t0 @@ -302,6 +311,7 @@ class MolmoGenerateText: "top_k": ("INT", {"default": 40, "min": 1}), # This doesn't seem to work for this model #"repetition_penalty": ("FLOAT", {"default": 1.1}), + "stop_strings": ("STRING", {"default": "<|endoftext|>"}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}), "include_prompt_in_output": ("BOOLEAN", {"default": False}), } @@ -312,7 +322,7 @@ class MolmoGenerateText: CATEGORY = "PixtralLlamaVision/Molmo" TITLE = "Generate Text with Molmo" - def generate_text(self, molmo_model, images, prompt, max_new_tokens, do_sample, temperature, top_p, top_k, seed, include_prompt_in_output): + def generate_text(self, molmo_model, images, prompt, max_new_tokens, do_sample, temperature, top_p, top_k, stop_strings, seed, include_prompt_in_output): device = molmo_model['model'].device print(f"Batch of {images.shape} images") image_list = [to_pil_image(image.numpy()) for image in images] @@ -323,6 +333,8 @@ class MolmoGenerateText: prompt_tokens = inputs["input_ids"].size(1) print(f"Prompt tokens: {prompt_tokens}") + stop_strings_list = stop_strings.split(",") + set_seed(seed) t0 = time.time() output = molmo_model['model'].generate_from_batch( @@ -334,8 +346,8 @@ class MolmoGenerateText: top_p=top_p, top_k=top_k, #repetition_penalty=repetition_penalty, - stop_string="<|endoftext|>", ), + stopping_criteria=[StopStringCriteria(tokenizer=molmo_model['processor'].tokenizer, stop_strings=stop_strings_list)], tokenizer=molmo_model['processor'].tokenizer, ) t1 = time.time() @@ -458,6 +470,132 @@ class ParseBoundingBoxes: return (bboxes,) +class ParsePoints: + """eyes""" + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "string": ("STRING",), + "filter": ("STRING",), + } + } + + RETURN_TYPES = ("POINT", "STRING", "STRING") + FUNCTION = "generate_points" + CATEGORY = "PixtralLlamaVision/Utility" + TITLE = "Parse Points" + + def generate_points(self, string, filter): + point_batches = [] + label_batches = [] + alt_label_batches = [] + if type(string) != list: + string = [string] # batch 1 + for s in string: + points = [] + labels = [] + alt_labels = [] + # Tried to design this regex in a way where even if the message gets cut off by the token limit, it finds the points + # Another absolutely ridiculous looking regex + for match in re.findall(r"""|$)>?([^<]*)""", s, flags=re.M): + try: + data = match[0] + if len(match) > 1: + alt = match[1] + if len(match) > 2: + inner = match[2] + else: + inner = "" + else: + alt = "" + inner = "" + + # Roughly matching + if alt == "" or filter.lower() in alt.lower() or filter.lower() in inner.lower(): + data_parts = data.split(" ") + for i in range(len(data_parts)//2): + # Points from Molmo are expressed as percentages + x = float(data_parts[2*i].split('"')[1])/100.0 + y = float(data_parts[2*i+1].split('"')[1])/100.0 + + # Check for duplicates + for point, label, alt_label in zip(points, labels, alt_labels): + if point[0] == x and point[1] == y and label == inner and alt_label == alt: + print(f"Duplicate point ({x}, {y}, {alt}, {inner})") + continue + + points.append([x, y]) + labels.append(inner) + alt_labels.append(alt) # I'm not really convinced alt even matters + else: + print(f"Non-matching filter for {match}") + except Exception as e: + print(f"Failed to parse points: {match}: {e}") + point_batches.append(points) + label_batches.append(labels) + alt_label_batches.append(alt_labels) + return (np.array(point_batches), np.array(label_batches), np.array(alt_label_batches)) + + +class PlotPoints: + @classmethod + def INPUT_TYPES(s): + return { + "optional": { + "labels": ("STRING",), + }, + "required": { + "points": ("POINT",), + "image": ("IMAGE",), + "size": ("INT", {"default": 5, "min": 1, "step": 1}), + "font_size": ("INT", {"default": 40}), + "color": ("STRING", {"default": "#0000ff"}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "plot_points" + CATEGORY = "PixtralLlamaVision/Utility" + TITLE = "Plot Points" + + def plot_points(self, points, labels, image, size, font_size, color): + image_width = image.shape[2] + image_height = image.shape[1] + + if labels is None or len(labels) == 0 or font_size == 0: + labels = np.array([["" for point in point_list] for point_list in points]) + + batch_size = image.shape[0] + if len(points) != len(labels) or len(points) != image.shape[0]: + print(f"Warning: Batch size mismatch: Image {image.shape}, points {points.shape}, labels {labels.shape}") + batch_size = min(image.shape[0], len(points), len(labels)) + + # font = ImageFont.truetype("Pillow/Tests/fonts/FreeMono.ttf", font_size) + # I might have overengineered this, it doesn't seem like the model can label separate objects in a single call. But you can concatenate the strings anyway. + colors = [color] + if "," in color: + colors = color.split(",") + color_map = {"": colors[0]} + for i, label in enumerate(np.unique(labels)): + color_map[label] = colors[i%len(colors)] + + # Add points to image (which is a tensor of floats of shape (batch, height, width, channels) + changed_images = [] + for img, point_list, label_list in zip(image, points, labels): + temp_image = to_pil_image(img.numpy()) + draw = ImageDraw.Draw(temp_image) + for point, label in zip(point_list, label_list): + x = int(image_width*point[0]) + y = int(image_height*point[1]) + draw.circle((x, y), fill=color_map[label], outline=color_map[label], radius=size) + if label != "": + draw.text((x, y-size), label, fill=color_map[label], font_size=font_size, anchor='md') + output_image = np.array(temp_image)/0xff + changed_images.append(output_image) + return (torch.Tensor(np.array(changed_images)),) + + def process_regex_flags(flags): flag_value = re.NOFLAG if 'a' in flags.lower(): @@ -655,6 +793,8 @@ NODE_CLASS_MAPPINGS = { "RegexSubstitution": RegexSubstitution, "JoinString": JoinString, "ParseBoundingBoxes": ParseBoundingBoxes, + "ParsePoints": ParsePoints, + "PlotPoints": PlotPoints, "SelectIndex": SelectIndex, "SliceList": SliceList, } diff --git a/pyproject.toml b/pyproject.toml index a79e21e..070cc44 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-pixtralllamavision" description = "For loading and running Pixtral, Llama 3.2 Vision, and Molmo models. Put models in the models/LLM folder." -version = "2.2.0" +version = "2.2.1" license = {file = "LICENSE"} [project.urls]