From 9eaa40b45235cc358efc683a1bbf6ea779eafe57 Mon Sep 17 00:00:00 2001
From: Cyber-BlackCat <rayncaine@gmail.com>
Date: Wed, 12 Mar 2025 01:25:25 +0800
Subject: [PATCH] joy caption2 added

---
 auto_caption.py  | 8 +++++---
 auto_caption2.py | 2 +-
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/auto_caption.py b/auto_caption.py
index 6e52f1c..0fe643c 100644
--- a/auto_caption.py
+++ b/auto_caption.py
@@ -72,7 +72,7 @@ class Joy_Model_load:
        
          # clip
         model_id = "google/siglip-so400m-patch14-384"
-        CLIP_PATH = download_hg_model(model_id,"clip")
+        CLIP_PATH = download_hg_model(model_id,"clip_vision")
 
         clip_processor = AutoProcessor.from_pretrained(CLIP_PATH) 
         clip_model = AutoModel.from_pretrained(
@@ -88,14 +88,16 @@ class Joy_Model_load:
        
         # LLM
         MODEL_PATH = download_hg_model(self.model,"LLM")
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH,use_fast=False)
+        LORA_PATH = os.path.join(folder_paths.models_dir, "loras-LLM", "cgrkzexw-599808")
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_fast=True)
+        # tokenizer = AutoTokenizer.from_pretrained(os.path.join(CAPTION_PATH, "text_model"), use_fast=True)
         assert isinstance(tokenizer, PreTrainedTokenizer) or isinstance(tokenizer, PreTrainedTokenizerFast), f"Tokenizer is of type {type(tokenizer)}"
 
         text_model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map="auto",trust_remote_code=True)
         text_model.eval()
 
         # Image Adapter
-        adapter_path =  os.path.join(folder_paths.models_dir,"Auto_Caption","image_adapter.pt")
+        adapter_path =  os.path.join(folder_paths.models_dir,"loras-LLM","image_adapter.pt")
 
         image_adapter = ImageAdapter(clip_model.config.hidden_size, text_model.config.hidden_size) # ImageAdapter(clip_model.config.hidden_size, 4096) 
         image_adapter.load_state_dict(torch.load(adapter_path, map_location="cpu"))
diff --git a/auto_caption2.py b/auto_caption2.py
index c8998c0..f3aa313 100644
--- a/auto_caption2.py
+++ b/auto_caption2.py
@@ -395,7 +395,7 @@ class Auto_Caption2:
     CATEGORY = "Auto Caption"
     RETURN_TYPES = ("STRING",)
     RETURN_NAMES = ("prompt",)
-    OUTPUT_NODE = True
+    # OUTPUT_NODE = True
     OUTPUT_IS_LIST = (True,)
     FUNCTION = "gen"