Add files

2024-09-25 03:58:44 -07:00
parent a6030efba7
commit 3dd272bf9e
7 changed files with 325 additions and 0 deletions
@@ -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)
@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
@@ -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": "<s>[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",
+}
@@ -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": [
+        "<s>[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
+}
@@ -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 = ""
@@ -0,0 +1 @@
+bitsandbytes