diff --git a/README.md b/README.md index 7bdeda5..616792e 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,13 @@ # ComfyUI-Pixtral For loading and running Pixtral models + +Includes two nodes, PixtralModelLoader and PixtralGenerateText. These should be self-explanatory. + +Install the latest version of transformers, which has support for Pixtral models: +`python_embeded\python.exe -m pip install git+https://github.com/huggingface/transformers` +Requires transformers 4.45.0 + +Also install bitsandbytes if you don't have it already: +`python_embeded\python.exe -m pip install bitsandbytes` + +![Example workflow](pixtral_caption_example.jpg) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..2e96bd6 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..b70b314 --- /dev/null +++ b/nodes.py @@ -0,0 +1,100 @@ +import comfy.utils +import comfy.model_management as mm +import folder_paths + +from transformers import LlavaForConditionalGeneration, AutoProcessor, BitsAndBytesConfig, set_seed +from torchvision.transforms.functional import to_pil_image +from PIL import Image +import time +import os +from pathlib import Path + +class PixtralModelLoader: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model_name": ([item.name for item in Path(folder_paths.models_dir, "pixtral").iterdir() if item.is_dir()],), + } + } + + RETURN_TYPES = ("PIXTRAL_MODEL",) + FUNCTION = "load_model" + CATEGORY = "Pixtral" + TITLE = "PixtralModelLoader" + + def load_model(self, model_name): + model_path = os.path.join(folder_paths.models_dir, "pixtral", model_name) + device = mm.get_torch_device() + model = LlavaForConditionalGeneration.from_pretrained( + model_path, + use_safetensors=True, + device_map=device, + ) + processor = AutoProcessor.from_pretrained(model_path) + pixtral_model = { + 'model': model, + 'processor': processor, + } + return (pixtral_model,) + +class PixtralGenerateText: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "pixtral_model": ("PIXTRAL_MODEL",), + "images": ("IMAGE",), + "prompt": ("STRING", {"default": "[INST]Caption this image:\n[IMG][/INST]", "multiline": True}), + "max_new_tokens": ("INT", {"default": 256, "min": 1, "max": 4096}), + "do_sample": ("BOOLEAN", {"default": True}), + "temperature": ("FLOAT", {"default": 0.5}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}), + } + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "generate_text" + CATEGORY = "Pixtral" + TITLE = "PixtralGenerateText" + + def generate_text(self, pixtral_model, images, prompt, max_new_tokens, do_sample, temperature, seed): + device = mm.get_torch_device() + print(type(images)) + # How does batched input work? I really don't know + # Also 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 + print(f"Batch of {images.shape} images") + image_list = [to_pil_image(image.numpy()) for image in images] + 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}") + set_seed(seed) + t0 = time.time() + generate_ids = pixtral_model['model'].generate( + **inputs, + max_new_tokens=max_new_tokens, + do_sample=do_sample, + temperature=temperature, + ) + 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)") + print(len(generate_ids[0][prompt_tokens:])) + output = pixtral_model['processor'].decode(generate_ids[0][prompt_tokens:], skip_special_tokens=True, clean_up_tokenization_spaces=False) + print(output) + return (output,) + +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, +} +NODE_DISPLAY_NAME_MAPPINGS = { + "PixtralModelLoader": "PixtralModelLoader", + "PixtralGenerateText": "PixtralGenerateText", +} diff --git a/pixtral_caption_example.jpg b/pixtral_caption_example.jpg new file mode 100644 index 0000000..ad6bfee Binary files /dev/null and b/pixtral_caption_example.jpg differ diff --git a/pixtral_caption_workflow.json b/pixtral_caption_workflow.json new file mode 100644 index 0000000..4e8d063 --- /dev/null +++ b/pixtral_caption_workflow.json @@ -0,0 +1,196 @@ +{ + "last_node_id": 19, + "last_link_id": 38, + "nodes": [ + { + "id": 10, + "type": "PixtralModelLoader", + "pos": { + "0": 364, + "1": 304 + }, + "size": { + "0": 315, + "1": 58 + }, + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "PIXTRAL_MODEL", + "type": "PIXTRAL_MODEL", + "links": [ + 36 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "PixtralModelLoader" + }, + "widgets_values": [ + "pixtral-12b-nf4" + ] + }, + { + "id": 4, + "type": "LoadImage", + "pos": { + "0": 362, + "1": 463 + }, + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 37 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "test.png", + "image" + ] + }, + { + "id": 19, + "type": "PixtralGenerateText", + "pos": { + "0": 835, + "1": 380 + }, + "size": [ + 405.81818181818176, + 258.090909090909 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "pixtral_model", + "type": "PIXTRAL_MODEL", + "link": 36 + }, + { + "name": "images", + "type": "IMAGE", + "link": 37 + } + ], + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": [ + 38 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "PixtralGenerateText" + }, + "widgets_values": [ + "[INST]Caption this image:\n[IMG][/INST]", + 256, + true, + 0.5, + 3722012111, + "randomize" + ] + }, + { + "id": 3, + "type": "Display Any (rgthree)", + "pos": { + "0": 1333, + "1": 368 + }, + "size": [ + 397.81818181818176, + 337.36363636363626 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "source", + "type": "*", + "link": 38, + "dir": 3 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "Display Any (rgthree)" + }, + "widgets_values": [ + "" + ] + } + ], + "links": [ + [ + 36, + 10, + 0, + 19, + 0, + "PIXTRAL_MODEL" + ], + [ + 37, + 4, + 0, + 19, + 1, + "IMAGE" + ], + [ + 38, + 19, + 0, + 3, + 0, + "*" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1.1, + "offset": [ + 60.090909090908944, + 74.63636363636358 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..e2839cc --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "comfyui-pixtral" +description = "For loading and running Pixtral models" +version = "1.0.0" +license = {file = "LICENSE"} + +[project.urls] +Repository = "https://github.com/SeanScripts/ComfyUI-Pixtral" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "seanscripts" +DisplayName = "ComfyUI-Pixtral" +Icon = "" \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..38cb110 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +bitsandbytes