diff --git a/PaliGemmaPixelProse.py b/PaliGemmaPixelProse.py new file mode 100644 index 0000000..cf25d1c --- /dev/null +++ b/PaliGemmaPixelProse.py @@ -0,0 +1,60 @@ +import os +import io +from pathlib import Path +import torch +from PIL import Image +from torchvision.transforms import ToPILImage +from transformers import AutoProcessor, PaliGemmaForConditionalGeneration +import folder_paths + +class PaliGemmaPixelProse: + def __init__(self): + self.model_id = "gokaygokay/PaliGemma-PixelProse" + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.model = PaliGemmaForConditionalGeneration.from_pretrained(self.model_id).to(self.device).eval() + self.processor = AutoProcessor.from_pretrained(self.model_id) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "prompt": ("STRING", {"multiline": False, "default": "caption en"}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("caption",) + FUNCTION = "generate_caption" + CATEGORY = "image/text" + + def generate_caption(self, image, prompt): + # Convert the image tensor to PIL Image + pil_image = ToPILImage()(image[0].permute(2, 0, 1)) + + # Process the image and prompt + model_inputs = self.processor(text=prompt, images=pil_image, return_tensors="pt").to(self.device) + + # Generate caption + with torch.inference_mode(): + input_len = model_inputs["input_ids"].shape[-1] + generation = self.model.generate( + **model_inputs, + repetition_penalty=1.05, + max_new_tokens=512, + do_sample=False + ) + + # Decode the generated text + generation = generation[0][input_len:] + decoded = self.processor.decode(generation, skip_special_tokens=True) + + return (decoded,) + +NODE_CLASS_MAPPINGS = { + "PaliGemmaPixelProse": PaliGemmaPixelProse +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PaliGemmaPixelProse": "PaliGemma PixelProse Caption" +} \ No newline at end of file diff --git a/README.md b/README.md index 984a085..84cb93d 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,9 @@ # comfy-groqchat -Content generation with open source models in comfyui via graq api implementation. +Updated image backpropagation nodes based on Google's open-source PaliGemma vision model. The model used in the node is a fine-tuned version of gokayfem's big brother, special thanks for that! Content generation with open source models in comfyui via graq api implementation. # Usage: +PaliGemmaPixelProse image backpropagation node. +![](demo01.png) The following four models are currently supported. * gemma-7b-it * llama3-70b-8192 @@ -19,5 +21,5 @@ The use of nodes enables the generation of positive and negative cues in 2 secon Picture with workflow ____ ##Statement: -The GroqChat node follows the MIT license agreement, and some of the functional code comes from other open source projects. Thanks to the original authors. For commercial use, please refer to the original project license agreement for authorization. +The GroqChat node follows the MIT license agreement, and some of the functional code comes from other open source projects. Thanks to the original authors. For commercial use, please refer to the original project license agreement for authorization.Thanks also to Google's open source PaliGemma model. diff --git a/__init__.py b/__init__.py index 3927f1a..1075dd7 100644 --- a/__init__.py +++ b/__init__.py @@ -1,7 +1,9 @@ import os import json +import sys from .groqchat import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from .PaliGemmaPixelProse import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS -__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/demo01.png b/demo01.png new file mode 100644 index 0000000..2c80256 Binary files /dev/null and b/demo01.png differ diff --git a/requirements.txt b/requirements.txt index 079cc8a..aeb7cb3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,5 @@ +transformers>=4.41.1 +torch +Pillow +huggingface_hub groq \ No newline at end of file