From 537541bdb39d3e7178b663c1df1ff6d1e3388728 Mon Sep 17 00:00:00 2001
From: AIExplorer <abuzar332@gmail.com>
Date: Sat, 17 May 2025 02:06:36 +0530
Subject: [PATCH] Add files via upload

---
 __init__.py        |   3 ++
 image_captioner.py | 117 +++++++++++++++++++++++++++++++++++++++++++++
 requirements.txt   |   0
 3 files changed, 120 insertions(+)
 create mode 100644 __init__.py
 create mode 100644 image_captioner.py
 create mode 100644 requirements.txt

diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..d0f235d
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,3 @@
+from .chatgpt_helper 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/image_captioner.py b/image_captioner.py
new file mode 100644
index 0000000..b5673c0
--- /dev/null
+++ b/image_captioner.py
@@ -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.",
+}
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..e69de29