Upload 2 files

ComfyUI/custom_nodes/img2txt-comfyui-nodes/src/blip_img2txt.py

@@ -1,3 +1,4 @@
+import os
 from PIL import Image
 from transformers import (
     BlipProcessor,
@@ -9,7 +10,7 @@ from transformers import (
 
 import torch
 import model_management
-
+import folder_paths
 
 class BLIPImg2Txt:
     def __init__(
@@ -21,21 +22,24 @@ class BLIPImg2Txt:
         repetition_penalty: float,
         search_beams: int,
         model_id: str = "Salesforce/blip-image-captioning-large",
+        custom_model_path: str = None,
     ):
         self.conditional_caption = conditional_caption
         self.model_id = model_id
+        self.custom_model_path = custom_model_path
+
+        if self.custom_model_path and os.path.exists(self.custom_model_path):
+            self.model_path = self.custom_model_path
+        else:
+            self.model_path = folder_paths.get_full_path("blip", model_id)
 
-        # Determine do_sample and num_beams
         if temperature > 1.1 or temperature < 0.90:
             do_sample = True
-            num_beams = 1  # Sampling does not use beam search
+            num_beams = 1
         else:
             do_sample = False
-            num_beams = (
-                search_beams if search_beams > 1 else 1
-            )  # Use beam search if num_beams > 1
+            num_beams = search_beams if search_beams > 1 else 1
 
-        # Initialize text config kwargs
         self.text_config_kwargs = {
             "do_sample": do_sample,
             "max_length": max_words,
@@ -51,18 +55,25 @@ class BLIPImg2Txt:
         if image.mode != "RGB":
             image = image.convert("RGB")
 
-        processor = BlipProcessor.from_pretrained(self.model_id)
+        if self.model_path and os.path.exists(self.model_path):
+            model_path = self.model_path
+            local_files_only = True
+        else:
+            model_path = self.model_id
+            local_files_only = False
+
+        processor = BlipProcessor.from_pretrained(model_path, local_files_only=local_files_only)
 
-        # Update and apply configurations
-        config_text = BlipTextConfig.from_pretrained(self.model_id)
+        config_text = BlipTextConfig.from_pretrained(model_path, local_files_only=local_files_only)
         config_text.update(self.text_config_kwargs)
-        config_vision = BlipVisionConfig.from_pretrained(self.model_id)
+        config_vision = BlipVisionConfig.from_pretrained(model_path, local_files_only=local_files_only)
         config = BlipConfig.from_text_vision_configs(config_text, config_vision)
 
         model = BlipForConditionalGeneration.from_pretrained(
-            self.model_id,
+            model_path,
             config=config,
             torch_dtype=torch.float16,
+            local_files_only=local_files_only
         ).to(model_management.get_torch_device())
 
         inputs = processor(
@@ -78,4 +89,4 @@ class BLIPImg2Txt:
         del model
         torch.cuda.empty_cache()
 
-        return ret
+        return ret

ComfyUI/custom_nodes/img2txt-comfyui-nodes/src/img2txt_node.py

@@ -14,6 +14,8 @@ from .mini_cpm_img2txt import MiniPCMImg2Txt
 
 from typing import Tuple
 
+import os
+import folder_paths
 
 class Img2TxtNode:
     CATEGORY = "img2txt"
@@ -145,6 +147,11 @@ class Img2TxtNode:
 
         captions = []
         if use_all_models or use_blip_model:
+            blip_model_path = folder_paths.get_folder_paths("blip")[0]
+            print(f"blip_model_path: {blip_model_path}")
+            if not blip_model_path or not os.path.exists(blip_model_path):
+                raise ValueError("BLIP model 'blip-image-captioning-large' not found in ComfyUI models directory. Please ensure it's in the 'models/blip' folder.")
+            
             blip = BLIPImg2Txt(
                 conditional_caption=blip_caption_prefix,
                 min_words=min_words,
@@ -152,6 +159,7 @@ class Img2TxtNode:
                 temperature=temperature,
                 repetition_penalty=repetition_penalty,
                 search_beams=search_beams,
+                custom_model_path=blip_model_path
             )
             captions.append(blip.generate_caption(raw_image))