896 lines
36 KiB
Python
896 lines
36 KiB
Python
import comfy.utils
|
|
import comfy.model_management as mm
|
|
#from comfy.model_patcher import ModelPatcher
|
|
import folder_paths
|
|
|
|
from transformers import (
|
|
LlavaForConditionalGeneration,
|
|
MllamaForConditionalGeneration,
|
|
AutoModelForCausalLM,
|
|
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, ImageDraw
|
|
import re
|
|
import time
|
|
|
|
# Using folder ComfyUI/models/LLM -- Place each model inside its own folder here, e.g. ComfyUI/models/LLM/pixtral-12b-nf4/model.safetensors
|
|
# Also include other config files and tokenizer files in that same folder
|
|
llm_model_dir = os.path.join(folder_paths.models_dir, "LLM")
|
|
# Add LLM folder if not present
|
|
if not os.path.exists(llm_model_dir):
|
|
os.makedirs(llm_model_dir)
|
|
|
|
model_type_map = {
|
|
"LlavaForConditionalGeneration": LlavaForConditionalGeneration,
|
|
"MllamaForConditionalGeneration": MllamaForConditionalGeneration,
|
|
"MolmoForCausalLM": AutoModelForCausalLM,
|
|
# Other vision models can be added here as needed but will require importing
|
|
"AutoModelForCausalLM": AutoModelForCausalLM,
|
|
}
|
|
|
|
def get_models_with_config():
|
|
models = []
|
|
for model_path in Path(llm_model_dir).iterdir():
|
|
if model_path.is_dir():
|
|
if os.path.exists(os.path.join(model_path, "config.json")):
|
|
models.append(model_path.parts[-1])
|
|
return models
|
|
|
|
def get_model_type(model_path):
|
|
config_path = os.path.join(model_path, "config.json")
|
|
if os.path.exists(config_path):
|
|
with open(config_path, 'r') as config_file:
|
|
config = json.load(config_file)
|
|
return config["architectures"][0]
|
|
print(f"Invalid model config for model {model_path}")
|
|
return "Invalid model config"
|
|
|
|
def get_models_of_type(model_type):
|
|
models = []
|
|
for model_path in Path(llm_model_dir).iterdir():
|
|
if model_path.is_dir():
|
|
current_model_type = get_model_type(model_path)
|
|
if current_model_type == model_type:
|
|
models.append(model_path.parts[-1])
|
|
return models
|
|
|
|
class PixtralModelLoader:
|
|
"""Loads a Pixtral model. Add models as folders inside the `ComfyUI/models/LLM` folder. Each model folder should contain a standard transformers loadable safetensors model along with a tokenizer and any config files needed."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_name": (get_models_of_type("LlavaForConditionalGeneration"),),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VISION_MODEL",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "PixtralLlamaVision/Pixtral"
|
|
TITLE = "Load Pixtral Model"
|
|
|
|
def load_model(self, model_name):
|
|
model_path = os.path.join(llm_model_dir, model_name)
|
|
print(f"Setting Pixtral model: {model_name}")
|
|
# Don't load the full model until needed for generation
|
|
processor = AutoProcessor.from_pretrained(model_path)
|
|
pixtral_model = {
|
|
'path': model_path,
|
|
'processor': processor,
|
|
}
|
|
return (pixtral_model,)
|
|
|
|
|
|
class PixtralGenerateText:
|
|
"""Generates text using a Pixtral model. Takes a list of images and a string prompt as input. The prompt must contain an equal number of [IMG] tokens to the number of images passed in."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"optional": {
|
|
"images": ("IMAGE",),
|
|
},
|
|
"required": {
|
|
"pixtral_model": ("VISION_MODEL",),
|
|
#"system_prompt": ("STRING", {"default": "", "multiline": True}),
|
|
"prompt": ("STRING", {"default": "Caption this image:\n[IMG]", "multiline": True}),
|
|
"max_new_tokens": ("INT", {"default": 256, "min": 1, "max": 4096}),
|
|
"do_sample": ("BOOLEAN", {"default": True}),
|
|
"temperature": ("FLOAT", {"default": 0.3, "min": 0, "step": 0.1}),
|
|
"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}),
|
|
"unload_after_generate": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "generate_text"
|
|
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, stop_strings, seed, include_prompt_in_output, unload_after_generate):
|
|
# Load model now if needed
|
|
device = mm.get_torch_device()
|
|
if pixtral_model['path'] and 'model' not in pixtral_model:
|
|
pixtral_model['model'] = LlavaForConditionalGeneration.from_pretrained(
|
|
pixtral_model['path'],
|
|
use_safetensors=True,
|
|
device_map=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
|
|
if images != None and len(images) > 0:
|
|
print(f"Batch of {images.shape} images")
|
|
image_list = [to_pil_image(image.numpy()) for image in images]
|
|
|
|
# Process prompt
|
|
# Example: <s>[INST]Caption this image:\n[IMG][/INST]
|
|
# Images can be placed anywhere, unlike the other models
|
|
image_tag_count = prompt.count("[IMG]")
|
|
added_image_tags = ""
|
|
if image_tag_count > 0 and (images is None or len(images) == 0):
|
|
print("Warning: Prompt contains image tags but no image")
|
|
elif image_tag_count < len(images):
|
|
added_image_tags = "[IMG]"*(len(images) - image_tag_count)
|
|
print("Warning: Adding extra images to the beginning of the prompt")
|
|
elif image_tag_count > len(images):
|
|
print("Warning: Too many image tags")
|
|
|
|
# Not sure how system vs user input differs for this model yet
|
|
final_prompt = "<s>"
|
|
#if system_prompt != "":
|
|
# final_prompt += f"[INST]{system_prompt}[/INST]"
|
|
final_prompt += f"[INST]{added_image_tags}{prompt}[/INST]"
|
|
|
|
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(
|
|
**inputs,
|
|
generation_config=GenerationConfig(
|
|
max_new_tokens=max_new_tokens,
|
|
do_sample=do_sample,
|
|
temperature=temperature,
|
|
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
|
|
generated_tokens = len(generate_ids[0]) - prompt_tokens
|
|
time_per_token = generated_tokens/total_time
|
|
print(f"Generated {generated_tokens} tokens in {total_time:.3f} s ({time_per_token:.3f} tok/s)")
|
|
output_tokens = generate_ids[0] if include_prompt_in_output else generate_ids[0][prompt_tokens:]
|
|
output = pixtral_model['processor'].decode(output_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=False)
|
|
print(output)
|
|
|
|
# Unload model
|
|
if unload_after_generate:
|
|
del pixtral_model['model']
|
|
torch.cuda.empty_cache()
|
|
print("Pixtral model unloaded")
|
|
|
|
return (output,)
|
|
|
|
|
|
class LlamaVisionModelLoader:
|
|
"""Loads a Llama 3.2 Vision model. Add models as folders inside the `ComfyUI/models/LLM` folder. Each model folder should contain a standard transformers loadable safetensors model along with a tokenizer and any config files needed."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_name": (get_models_of_type("MllamaForConditionalGeneration"),),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VISION_MODEL",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "PixtralLlamaVision/LlamaVision"
|
|
TITLE = "Load Llama Vision Model"
|
|
|
|
def load_model(self, model_name):
|
|
model_path = os.path.join(llm_model_dir, model_name)
|
|
print(f"Setting Llama Vision model: {model_name}")
|
|
# Don't load the full model until needed for generation
|
|
processor = AutoProcessor.from_pretrained(model_path)
|
|
llama_vision_model = {
|
|
'path': model_path,
|
|
'processor': processor,
|
|
}
|
|
return (llama_vision_model,)
|
|
|
|
|
|
class LlamaVisionGenerateText:
|
|
"""Generates text using a Llama 3.2 Vision model. The prompt must contain an equal number of <|image|> tokens to the number of images passed in. Image tokens must also be sequential and before the text you want them to apply to for the image attention to work as intended."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"optional": {
|
|
"images": ("IMAGE",),
|
|
},
|
|
"required": {
|
|
"llama_vision_model": ("VISION_MODEL",),
|
|
"system_prompt": ("STRING", {"default": "", "multiline": True}),
|
|
"prompt": ("STRING", {"default": "Caption this image.", "multiline": True}),
|
|
"max_new_tokens": ("INT", {"default": 256, "min": 1, "max": 4096}),
|
|
"do_sample": ("BOOLEAN", {"default": True}),
|
|
"temperature": ("FLOAT", {"default": 0.3, "min": 0.0, "step": 0.1}),
|
|
"top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.1}),
|
|
"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}),
|
|
"unload_after_generate": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "generate_text"
|
|
CATEGORY = "PixtralLlamaVision/LlamaVision"
|
|
TITLE = "Generate Text with Llama Vision"
|
|
|
|
# TODO: Support batching
|
|
|
|
def generate_text(self, llama_vision_model, images, system_prompt, prompt, max_new_tokens, do_sample, temperature, top_p, top_k, stop_strings, seed, include_prompt_in_output, unload_after_generate):
|
|
# Load model now if needed
|
|
device = mm.get_torch_device()
|
|
if llama_vision_model['path'] and 'model' not in llama_vision_model:
|
|
llama_vision_model['model'] = MllamaForConditionalGeneration.from_pretrained(
|
|
llama_vision_model['path'],
|
|
use_safetensors=True,
|
|
device_map=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
|
|
|
|
if images != None and len(images) > 0:
|
|
print(f"Batch of {images.shape} images")
|
|
image_list = [to_pil_image(image.numpy()) for image in images]
|
|
|
|
# Process prompt
|
|
image_tags = "<|image|>"*len(images)
|
|
final_prompt = "<|begin_of_text|>"
|
|
if system_prompt != "":
|
|
final_prompt += f"<|start_header_id|>system<|end_header_id|>\n\n{system_prompt}<|eot_id|>\n\n"
|
|
final_prompt += f"<|start_header_id|>user<|end_header_id|>\n\n{image_tags}{prompt}<|eot_id|>\n\n"
|
|
final_prompt += "<|start_header_id|>assistant<|end_header_id|>\n\n"
|
|
|
|
inputs = llama_vision_model['processor'](images=image_list, text=final_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(
|
|
**inputs,
|
|
generation_config=GenerationConfig(
|
|
max_new_tokens=max_new_tokens,
|
|
do_sample=do_sample,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
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
|
|
generated_tokens = len(generate_ids[0]) - prompt_tokens
|
|
time_per_token = generated_tokens/total_time
|
|
print(f"Generated {generated_tokens} tokens in {total_time:.3f} s ({time_per_token:.3f} tok/s)")
|
|
output_tokens = generate_ids[0] if include_prompt_in_output else generate_ids[0][prompt_tokens:]
|
|
output = llama_vision_model['processor'].decode(output_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=False)
|
|
print(output)
|
|
|
|
# Unload model
|
|
if unload_after_generate:
|
|
del llama_vision_model['model']
|
|
torch.cuda.empty_cache()
|
|
print("Llama vision model unloaded")
|
|
|
|
return (output,)
|
|
|
|
|
|
class MolmoModelLoader:
|
|
"""Loads a Molmo model. Add models as folders inside the `ComfyUI/models/LLM` folder. Each model folder should contain a standard transformers loadable safetensors model along with a tokenizer and any config files needed."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_name": (get_models_of_type("MolmoForCausalLM"),),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VISION_MODEL",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "PixtralLlamaVision/Molmo"
|
|
TITLE = "Load Molmo Model"
|
|
|
|
def load_model(self, model_name):
|
|
model_path = os.path.join(llm_model_dir, model_name)
|
|
print(f"Setting Molmo model: {model_name}")
|
|
# Don't load the full model until needed for generation
|
|
processor = AutoProcessor.from_pretrained(
|
|
model_path,
|
|
torch_dtype="auto",
|
|
trust_remote_code=True,
|
|
)
|
|
molmo_model = {
|
|
'path': model_path,
|
|
'processor': processor,
|
|
}
|
|
return (molmo_model,)
|
|
|
|
|
|
class MolmoGenerateText:
|
|
"""Generates text using a Molmo model. Takes a list of images and a string prompt as input. The prompt must contain an equal number of [IMG] tokens to the number of images passed in."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"optional": {
|
|
"images": ("IMAGE",),
|
|
},
|
|
"required": {
|
|
"molmo_model": ("VISION_MODEL",),
|
|
"system_prompt": ("STRING", {"default": "", "multiline": True}),
|
|
"prompt": ("STRING", {"default": "Describe this image. ", "multiline": True}),
|
|
"max_new_tokens": ("INT", {"default": 256, "min": 1, "max": 4096}),
|
|
"do_sample": ("BOOLEAN", {"default": True}),
|
|
"temperature": ("FLOAT", {"default": 0.3, "min": 0, "step": 0.1}),
|
|
"top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.1}),
|
|
"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}),
|
|
"unload_after_generate": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "generate_text"
|
|
CATEGORY = "PixtralLlamaVision/Molmo"
|
|
TITLE = "Generate Text with Molmo"
|
|
|
|
# TODO: Support batching
|
|
|
|
def generate_text(self, molmo_model, images, system_prompt, prompt, max_new_tokens, do_sample, temperature, top_p, top_k, stop_strings, seed, include_prompt_in_output, unload_after_generate):
|
|
# Load model now if needed
|
|
device = mm.get_torch_device()
|
|
if molmo_model['path'] and 'model' not in molmo_model:
|
|
molmo_model['model'] = AutoModelForCausalLM.from_pretrained(
|
|
molmo_model['path'],
|
|
use_safetensors=True,
|
|
device_map=device,
|
|
torch_dtype="auto",
|
|
trust_remote_code=True,
|
|
)
|
|
|
|
if images != None and len(images) > 0:
|
|
print(f"Batch of {images.shape} images")
|
|
image_list = [to_pil_image(image.numpy()) for image in images]
|
|
|
|
# Process prompt
|
|
final_prompt = ""
|
|
if system_prompt != "":
|
|
final_prompt += f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
|
|
final_prompt += f"<|im_start|>user\n{prompt}<|im_end|>\n"
|
|
final_prompt += "<|im_start|>assistant\n"
|
|
|
|
inputs = molmo_model['processor'].process(
|
|
images=image_list,
|
|
#system_prompt=system_prompt, # Doesn't do anything
|
|
text=final_prompt,
|
|
message_format="none",
|
|
always_start_with_space=False,
|
|
)
|
|
inputs = {k: v.to(device).unsqueeze(0) for k, v in inputs.items()}
|
|
|
|
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(
|
|
inputs,
|
|
generation_config=GenerationConfig(
|
|
max_new_tokens=max_new_tokens,
|
|
do_sample=do_sample,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
#repetition_penalty=repetition_penalty,
|
|
),
|
|
stopping_criteria=[StopStringCriteria(tokenizer=molmo_model['processor'].tokenizer, stop_strings=stop_strings_list)],
|
|
tokenizer=molmo_model['processor'].tokenizer,
|
|
)
|
|
t1 = time.time()
|
|
|
|
total_time = t1 - t0
|
|
generated_tokens = output.size(1) - prompt_tokens
|
|
time_per_token = generated_tokens/total_time
|
|
print(f"Generated {generated_tokens} tokens in {total_time:.3f} s ({time_per_token:.3f} tok/s)")
|
|
|
|
output_tokens = output[0] if include_prompt_in_output else output[0, prompt_tokens:]
|
|
generated_text = molmo_model['processor'].tokenizer.decode(output_tokens, skip_special_tokens=True)
|
|
print(generated_text)
|
|
|
|
# Unload model
|
|
if unload_after_generate:
|
|
del molmo_model['model']
|
|
torch.cuda.empty_cache()
|
|
print("Molmo model unloaded")
|
|
|
|
return (generated_text,)
|
|
|
|
|
|
class AutoVisionModelLoader:
|
|
"""Loads a vision model. Add models as folders inside the `ComfyUI/models/LLM` folder. Each model folder should contain a standard transformers loadable safetensors model along with a tokenizer and any config files needed. Use `trust_remote_code` at your own risk."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_name": (get_models_with_config(),),
|
|
"trust_remote_code": ("BOOLEAN", {"default": False}), # No longer very useful. I can add a bit of code checking this when loading Pixtral/Llama Vision if there are custom finetunes of them in the future.
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VISION_MODEL",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "PixtralLlamaVision/VLM"
|
|
TITLE = "Load Vision Model"
|
|
|
|
def load_model(self, model_name, trust_remote_code):
|
|
model_path = os.path.join(llm_model_dir, model_name)
|
|
device = mm.get_torch_device()
|
|
# Don't load the full model until needed for generation
|
|
try:
|
|
model_type_name = get_model_type(model_path)
|
|
print(f"Setting vision model: {model_name} of type {model_type_name}")
|
|
'''
|
|
model_type = model_type_map.get(model_type_name, AutoModelForCausalLM)
|
|
model = model_type.from_pretrained(
|
|
model_path,
|
|
use_safetensors=True,
|
|
device_map=device,
|
|
torch_dtype="auto",
|
|
trust_remote_code=trust_remote_code,
|
|
)
|
|
'''
|
|
processor = AutoProcessor.from_pretrained(
|
|
model_path,
|
|
torch_dtype="auto",
|
|
trust_remote_code=trust_remote_code,
|
|
)
|
|
vision_model = {
|
|
'path': model_path,
|
|
'model_type_name': model_type_name, # Not used yet
|
|
'trust_remote_code': trust_remote_code, # Not used yet
|
|
'processor': processor,
|
|
}
|
|
return (vision_model,)
|
|
except Exception as e:
|
|
print(f"Error loading vision model: {e}")
|
|
raise
|
|
|
|
|
|
# Utility for bounding boxes (I'm sure this has been done before but I just wanted to try it out to see how well Pixtral can do it)
|
|
class ParseBoundingBoxes:
|
|
"""Uses a regular expression to find bounding boxes in a string, returning a list of bbox objects (compatible with mtb). `relative` means the bounding box uses float values between 0 and 1 if true and absolute image coordinates if false. `corners_only` means the bounding box is [(x1, y1), (x2, y2)] if true and [(x1, y1), (width, height)] if false. Parentheses are treated as optional."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"string": ("STRING",),
|
|
"relative": ("BOOLEAN", {"default": True}),
|
|
"corners_only": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BBOX",)
|
|
FUNCTION = "generate_bboxes"
|
|
CATEGORY = "PixtralLlamaVision/Utility"
|
|
TITLE = "Parse Bounding Boxes"
|
|
|
|
def generate_bboxes(self, image, string, relative, corners_only):
|
|
image_width = image.shape[2]
|
|
image_height = image.shape[1]
|
|
|
|
bboxes = []
|
|
# Ridiculous-looking regex
|
|
for match in re.findall(r"""\[?\(?([0-9\.]+),\s*([0-9\.]+)\)?,\s*\(?([0-9\.]+),\s*([0-9\.]+)\)?\]?""", string, flags=re.M):
|
|
try:
|
|
x1_raw = float(match[0])
|
|
y1_raw = float(match[1])
|
|
x2_raw = float(match[2])
|
|
y2_raw = float(match[3])
|
|
|
|
if relative:
|
|
x1 = int(image_width*x1_raw)
|
|
y1 = int(image_height*y1_raw)
|
|
x2 = int(image_width*x2_raw)
|
|
y2 = int(image_height*y2_raw)
|
|
else:
|
|
x1 = int(x1_raw)
|
|
y1 = int(y1_raw)
|
|
x2 = int(x2_raw)
|
|
y2 = int(y2_raw)
|
|
|
|
if corners_only:
|
|
width = x2 - x1
|
|
height = y2 - y1
|
|
else:
|
|
width = x2
|
|
height = y2
|
|
|
|
if width <= 0 or width > image_width or height <= 0 or height > image_height:
|
|
print(f"Invalid bbox: ({x1}, {y1}, {width}, {height})")
|
|
continue
|
|
bbox = (x1, y1, width, height)
|
|
bboxes.append(bbox)
|
|
except Exception as e:
|
|
print(f"Failed to parse bbox: {match}")
|
|
|
|
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
|
|
valid = True
|
|
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})")
|
|
valid = False
|
|
break
|
|
if valid:
|
|
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):
|
|
# Workaround for Python 3.10 not having re.NOFLAG
|
|
flag_value = 0 # re.NOFLAG
|
|
if 'a' in flags.lower():
|
|
flag_value |= re.A
|
|
if 'i' in flags.lower():
|
|
flag_value |= re.I
|
|
if 'l' in flags.lower():
|
|
flag_value |= re.L
|
|
if 'm' in flags.lower():
|
|
flag_value |= re.M
|
|
if 's' in flags.lower():
|
|
flag_value |= re.S
|
|
if 'u' in flags.lower(): # u for useless
|
|
flag_value |= re.U
|
|
if 'x' in flags.lower():
|
|
flag_value |= re.X
|
|
return flag_value
|
|
|
|
# Utility nodes that I couldn't find elsewhere, not sure why?
|
|
class RegexSplitString:
|
|
"""Uses a regular expression to split in a string by a pattern into a list of strings"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pattern": ("STRING",),
|
|
"string": ("STRING",),
|
|
"flags": ("STRING", {"default": "M"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "split_string"
|
|
CATEGORY = "PixtralLlamaVision/Utility"
|
|
TITLE = "Regex Split String"
|
|
|
|
def split_string(self, pattern, string, flags):
|
|
return (re.split(pattern, string, flags=process_regex_flags(flags)),)
|
|
|
|
|
|
class RegexSearch:
|
|
"""Uses a regular expression to search for the first occurrence of a pattern in a string, returning whether the pattern was found, the start and end positions if found, and the list of match groups if found"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pattern": ("STRING",),
|
|
"string": ("STRING",),
|
|
"flags": ("STRING", {"default": "M"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BOOLEAN", "INT", "INT", "STRING")
|
|
FUNCTION = "search"
|
|
CATEGORY = "PixtralLlamaVision/Utility"
|
|
TITLE = "Regex Search"
|
|
|
|
def search(self, pattern, string, flags):
|
|
match = re.search(pattern, string, flags=process_regex_flags(flags))
|
|
if match:
|
|
span = match.span()
|
|
groups = list(match.groups())
|
|
return (True, span[0], span[1], groups)
|
|
return (False, 0, 0, [])
|
|
|
|
|
|
class RegexFindAll:
|
|
"""Uses a regular expression to find all matches of a pattern in a string, returning a list of match groups (which could be strings or tuples of strings if you have more than one match group)"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pattern": ("STRING",),
|
|
"string": ("STRING",),
|
|
"flags": ("STRING", {"default": "M"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "find_all"
|
|
CATEGORY = "PixtralLlamaVision/Utility"
|
|
TITLE = "Regex Find All"
|
|
|
|
def find_all(self, pattern, string, flags):
|
|
return (re.findall(pattern, string, flags=process_regex_flags(flags)),)
|
|
|
|
|
|
# This one is also available in Derfuu_ComfyUI_ModdedNodes
|
|
class RegexSubstitution:
|
|
"""Uses a regular expression to find and replace text in a string"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pattern": ("STRING",),
|
|
"string": ("STRING",),
|
|
"replace": ("STRING",),
|
|
"flags": ("STRING", {"default": "M"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "sub"
|
|
CATEGORY = "PixtralLlamaVision/Utility"
|
|
TITLE = "Regex Substitution"
|
|
|
|
def sub(self, pattern, string, replace, flags):
|
|
return (re.sub(pattern, replace, string, flags=process_regex_flags(flags)),)
|
|
|
|
|
|
class JoinString:
|
|
"""Joins a list of strings with a delimiter between them"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"string_list": ("STRING",),
|
|
"delimiter": ("STRING", {"default": ", "}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "join_string"
|
|
CATEGORY = "PixtralLlamaVision/Utility"
|
|
TITLE = "Join String"
|
|
|
|
def join_string(self, string_list, delimiter):
|
|
# Convert to strings just in case? Or is this a bad idea? Well, it'll error if they're not strings, so I guess this will have to do
|
|
return (delimiter.join([str(string) for string in string_list]),)
|
|
|
|
|
|
# Arbitrary data type for list/tuple indexing
|
|
class AnyType(str):
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
ANY = AnyType("*")
|
|
|
|
# These ones are especially weird to not be doable in ComfyUI base
|
|
class SelectIndex:
|
|
"""Returns list[index]"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"list": (ANY,),
|
|
"index": ("INT", {"default": 0}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (ANY,)
|
|
FUNCTION = "select_index"
|
|
CATEGORY = "PixtralLlamaVision/Utility"
|
|
TITLE = "Select Index"
|
|
|
|
def select_index(self, list, index):
|
|
return (list[index],)
|
|
|
|
class SliceList:
|
|
"""Returns list[start_index:end_index]"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"list": (ANY,),
|
|
"start_index": ("INT", {"default": 0}),
|
|
"end_index": ("INT", {"default": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (ANY,)
|
|
FUNCTION = "select_index"
|
|
CATEGORY = "PixtralLlamaVision/Utility"
|
|
TITLE = "Slice List"
|
|
|
|
def select_index(self, list, start_index, end_index):
|
|
return (list[start_index:end_index],)
|
|
|
|
# Batch Count works for getting list length
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PixtralModelLoader": PixtralModelLoader,
|
|
"PixtralGenerateText": PixtralGenerateText,
|
|
# Not really much need to work with the image tokenization directly for something like image captioning, but might be interesting later...
|
|
#"PixtralImageEncode": PixtralImageEncode,
|
|
#"PixtralTextEncode": PixtralTextEncode,
|
|
"LlamaVisionModelLoader": LlamaVisionModelLoader,
|
|
"LlamaVisionGenerateText": LlamaVisionGenerateText,
|
|
"MolmoModelLoader": MolmoModelLoader,
|
|
"MolmoGenerateText": MolmoGenerateText,
|
|
"AutoVisionModelLoader": AutoVisionModelLoader,
|
|
"RegexSplitString": RegexSplitString,
|
|
"RegexSearch": RegexSearch,
|
|
"RegexFindAll": RegexFindAll,
|
|
"RegexSubstitution": RegexSubstitution,
|
|
"JoinString": JoinString,
|
|
"ParseBoundingBoxes": ParseBoundingBoxes,
|
|
"ParsePoints": ParsePoints,
|
|
"PlotPoints": PlotPoints,
|
|
"SelectIndex": SelectIndex,
|
|
"SliceList": SliceList,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {k:v.TITLE for k,v in NODE_CLASS_MAPPINGS.items()} |