nodes/SeanScripts-ComfyUI-PixtralLlamaMolmoVision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add stop strings and Molmo pointing utility nodes

Browse Source
This commit is contained in:
SeanScripts
2024-09-28 02:26:49 -07:00
parent beda1d4708
commit 5f30ee91ff
11 changed files with 523 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+13 -3
View File
@@ -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.
llama_vision_bounding_box_example.jpg → examples/llama_vision_bounding_box_example.jpg
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 324 KiB

After

Width:  |  Height:  |  Size: 324 KiB

llama_vision_bounding_box_workflow.json → examples/llama_vision_bounding_box_workflow.json
View File
examples/molmo_count_example.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 564 KiB

examples/molmo_counting_workflow.json
+363
View File
@@ -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
}
examples/molmo_multi_pointing_example.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 452 KiB

pixtral_caption_example.jpg → examples/pixtral_caption_example.jpg
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 226 KiB

After

Width:  |  Height:  |  Size: 226 KiB

pixtral_caption_workflow.json → examples/pixtral_caption_workflow.json
View File
pixtral_comparison_example.jpg → examples/pixtral_comparison_example.jpg
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 547 KiB

After

Width:  |  Height:  |  Size: 547 KiB

nodes.py
+146 -6
View File
@@ -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": "</s>"}),
"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:
"""<points x1=\"43.6\" y1=\"37.7\" x2=\"69.6\" y2=\"41.1\" alt=\"eyes\">eyes</points>"""
@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"""<points?\s*([xy\d\.="\s]*?)\s*(?:alt="([^"]*)")?(?=>|$)>?([^<]*)""", 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,
}
pyproject.toml
+1 -1
View File
@@ -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]