Add files via upload

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"
+}

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.
 

__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']
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

requirements.txt

@@ -1 +1,5 @@
+transformers>=4.41.1 
+torch
+Pillow
+huggingface_hub
 groq