From 9b59d919642255aea88d247351d8d41bb9a76198 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Thu, 22 Aug 2024 22:42:15 +0300
Subject: [PATCH] Add max_image_size to the captioning node

---
 nodes.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/nodes.py b/nodes.py
index 88ad401..c1e18b0 100644
--- a/nodes.py
+++ b/nodes.py
@@ -216,11 +216,12 @@ class OneVisionCaptionFolder:
             "required": {
                 "llava_model": ("LLAVAMODEL", ),
                 "folder_path": ("STRING", ),
-                "prompt": ("STRING", {"default": "", "multiline": True} ),
-                "max_tokens": ("INT", {"default": 4096, "min": 1, "max": 8192}),
+                "prompt": ("STRING", {"default": "You are AI captioning tool, you caption images in very elaborate detail without referring to the image as 'the image', the results should be useful for image model training purposes. You focus on the composition, style and action any possible subject is performing. You don't make assumptions or try to tell a story. You also describe the background of the image separately. Caption this image:", "multiline": True} ),
+                "max_tokens": ("INT", {"default": 512, "min": 1, "max": 8192}),
                 "keep_model_loaded": ("BOOLEAN", {"default": True}),
-                "temperature": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "temperature": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01}),
                 "seed": ("INT", {"default": 1, "min": 1, "max": 0xffffffffffffffff}),
+                "max_image_size": ("INT", {"default": 1024, "min": 256, "max": 8192}),
             },
         }
 
@@ -229,7 +230,7 @@ class OneVisionCaptionFolder:
     OUTPUT_IS_LIST = (True,)
     CATEGORY = "LLaVA-OneVision"
 
-    def caption(self, folder_path, llava_model, prompt, max_tokens, keep_model_loaded, temperature, seed):
+    def caption(self, folder_path, llava_model, prompt, max_tokens, keep_model_loaded, temperature, seed, max_image_size):
         from PIL import Image
         image_files = []
         for filename in os.listdir(folder_path):
@@ -244,6 +245,12 @@ class OneVisionCaptionFolder:
             img_path = os.path.join(folder_path, filename)
             try:
                 img = Image.open(img_path).convert('RGB')
+                # Resize image if it exceeds max_image_size
+                if max(img.size) > max_image_size:
+                    aspect_ratio = min(max_image_size / img.size[0], max_image_size / img.size[1])
+                    new_size = (int(img.size[0] * aspect_ratio), int(img.size[1] * aspect_ratio))
+                    img = img.resize(new_size, Image.LANCZOS)
+
                 img_tensor = transform(img)
                 img_tensor = img_tensor.unsqueeze(0).permute(0, 2, 3, 1)
             except IOError: