Add files via upload

2025-05-17 02:06:36 +05:30
parent 1a8131705d
commit 537541bdb3
3 changed files with 120 additions and 0 deletions
@@ -0,0 +1,3 @@
+from .chatgpt_helper import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
@@ -0,0 +1,117 @@
+import torch
+import os
+from PIL import Image
+from transformers import (
+    AutoModelForImageTextToText,
+    AutoProcessor
+)
+
+class ImageCaptioner:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "rename_images_sequencially": ("BOOLEAN", {
+                    "default": True, "label_on": "Yes", "label_off": "No"
+                }),
+                 "image_folder_path": ("STRING", {"multiline": True, "default": ""}),
+                "model_id": ("STRING", {"multiline": True, "default": ""}),
+                "system_message": ("STRING", {"multiline": True, "default": ""}),
+            }
+        }
+
+    
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("updated_status",)
+    FUNCTION = "generate_captions"
+    OUTPUT_NODE = True
+    CATEGORY = "utils"
+
+    # Class variable to store the last inputs and outputs
+    _cache = {}
+    
+    def generate_captions(self,rename_images_sequencially,image_folder_path, model_id, system_message,):
+        # Handle case where inputs are not lists
+        print("inside generate_captions function...")
+        if rename_images_sequencially:
+            # Rename images in the folder sequentially
+            for i, image_name in enumerate(os.listdir(image_folder_path)):
+                if image_name.endswith(".jpg") or image_name.endswith(".png"):
+                    new_image_name = f"{i+1}.jpg"
+                    os.rename(os.path.join(image_folder_path, image_name), os.path.join(image_folder_path, new_image_name))
+                    print(f"Renamed {image_name} to {new_image_name}")
+        else:
+            # Rename images in the folder sequentially
+            for i, image_name in enumerate(os.listdir(image_folder_path)):
+                if image_name.endswith(".jpg") or image_name.endswith(".png"):
+                    new_image_name = f"{i+1}.jpg"
+                    os.rename(os.path.join(image_folder_path, image_name), os.path.join(image_folder_path, new_image_name))
+                    print(f"Renamed {image_name} to {new_image_name}")
+        # Load the model and processor
+        model = AutoModelForImageTextToText.from_pretrained(
+        model_id,
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2", # Use "flash_attention_2" when running on Ampere or newer GPU or use "eager" for older GPUs
+        )
+        min_pixels = 512*28*28
+        max_pixels = 1600*28*28
+
+        processor = AutoProcessor.from_pretrained(model_id, max_pixels=max_pixels, min_pixels=min_pixels)
+        status=[]
+        system_message = "You are an expert image describer."
+        for image in os.listdir(image_folder_path):
+            if image.endswith(".jpg") or image.endswith(".png"):
+                image_path = os.path.join(image_folder_path, image)
+                image = Image.open(image_path).convert("RGB")
+                messages = [
+                {
+                    "role": "system",
+                    "content": [{"type": "text", "text": system_message}],
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": "Describe this image."},
+                        {"type": "image", "image": image},
+                    ],
+                },
+                ]
+                text = processor.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True
+                )
+                inputs = processor(
+                text=[text],
+                images=image,
+                padding=True,
+                return_tensors="pt",
+                )
+                inputs = inputs.to(model.device)
+                generated_ids = model.generate(**inputs, max_new_tokens=512, min_p=0.1, do_sample=True, temperature=1.5)
+                generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+                output_text = processor.batch_decode(
+                    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+                )
+                caption = output_text[0]
+                print(f"Caption for {image}: {caption}")
+                # Save the caption to a file
+                caption_file_path = os.path.join(image_folder_path, f"{os.path.splitext(image)[0]}.txt")
+                with open(caption_file_path, "w") as caption_file:
+                    caption_file.write(caption)
+                status.append(f"Caption for {image}: {caption}")
+
+
+
+
+        
+        return ", ".join(status),
+
+   
+
+NODE_CLASS_MAPPINGS = {
+    "ImageCaptioner": ImageCaptioner
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ImageCaptioner": "ImageCaptioner helps to generate captions for images in a folder and rename them sequentially.",
+}