#INIT First version

version 0.0.1 版本

__init__.py

@@ -0,0 +1,18 @@
+from .joy_caption_two_node import Joy_caption_two
+from .joy_caption_two_node import Joy_caption_two_advanced
+from .joy_caption_two_node import Joy_caption_two_load
+from .joy_caption_two_node import Joy_extra_options
+
+NODE_CLASS_MAPPINGS = {
+    "Joy_caption_two": Joy_caption_two,
+    "Joy_caption_two_advanced": Joy_caption_two_advanced,
+    "Joy_caption_two_load": Joy_caption_two_load,
+    "Joy_extra_options": Joy_extra_options,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Joy_caption_two": "Joy Caption Two",
+    "Joy_caption_two_advanced": "Joy Caption Two Advanced",
+    "Joy_caption_two_load": "Joy Caption Two Load",
+    "Joy_extra_options": "Joy Caption Extra Options",
+}

joy_caption_two_node.py

@@ -0,0 +1,602 @@
+import gc
+import time
+
+import numpy as np
+from torch import nn
+from transformers import AutoModel, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, \
+    AutoModelForCausalLM
+from pathlib import Path
+import torch
+import torch.amp.autocast_mode
+from PIL import Image
+import os
+
+import comfy.model_management
+import folder_paths
+import torchvision.transforms.functional as TVF
+
+from comfy.model_management import get_torch_device, unload_all_models, get_free_memory
+
+from comfy.model_management import load_models_gpu, soft_empty_cache
+from comfy.model_patcher import ModelPatcher
+from .uitls import download_hg_model, modify_json_value
+from .joy_config import joy_config
+
+DEVICE = get_torch_device()
+
+BASE_MODEL_PATH = Path(folder_paths.models_dir, "Joy_caption_two")
+
+def tensor2pil(t_image: torch.Tensor)  -> Image:
+    return Image.fromarray(np.clip(255.0 * t_image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
+
+class JoyClipVisionModel:
+    def __init__(self):
+        self.load_device = comfy.model_management.text_encoder_device()
+        self.offload_device = comfy.model_management.text_encoder_offload_device()
+
+        # clip
+        model_id = "google/siglip-so400m-patch14-384"
+        CLIP_PATH = download_hg_model(model_id, "clip")
+
+        clip_model = AutoModel.from_pretrained(
+            CLIP_PATH,
+            trust_remote_code=True
+        )
+
+        clip_model = clip_model.vision_model
+
+        assert (BASE_MODEL_PATH / "clip_model.pt").exists()
+        print("Loading VLM's custom vision model")
+        checkpoint = torch.load(BASE_MODEL_PATH / "clip_model.pt", map_location='cpu', weights_only=True)
+        checkpoint = {k.replace("_orig_mod.module.", ""): v for k, v in checkpoint.items()}
+        clip_model.load_state_dict(checkpoint)
+        del checkpoint
+
+        clip_model.eval()
+        clip_model.requires_grad_(False)
+        self.model = clip_model
+
+        self.patcher = ModelPatcher(self.model, load_device=self.load_device, offload_device=self.offload_device)
+
+    def encode_image(self, pixel_values):
+        #print(f"{id(self)}之前 in JoyClipVisionModel: {next(self.model.parameters()).device}")  # 打印模型参数的设备
+        load_models_gpu([self.patcher], force_full_load=True, force_patch_weights=True)
+        #print(f"之后 in JoyClipVisionModel: {next(self.model.parameters()).device}")  # 打印模型参数的设备
+        vision_outputs = self.model(pixel_values=pixel_values, output_hidden_states=True)
+        return vision_outputs
+
+
+
+
+class ImageAdapter(nn.Module):
+    def __init__(self, input_features: int, output_features: int, ln1: bool, pos_emb: bool, num_image_tokens: int, deep_extract: bool):
+        super().__init__()
+        self.deep_extract = deep_extract
+        if self.deep_extract:
+            input_features = input_features * 5
+
+        self.linear1 = nn.Linear(input_features, output_features)
+        self.activation = nn.GELU()
+        self.linear2 = nn.Linear(output_features, output_features)
+        self.ln1 = nn.Identity() if not ln1 else nn.LayerNorm(input_features)
+        self.pos_emb = None if not pos_emb else nn.Parameter(torch.zeros(num_image_tokens, input_features))
+        # Other tokens (<|image_start|>, <|image_end|>, <|eot_id|>)
+        self.other_tokens = nn.Embedding(3, output_features)
+        self.other_tokens.weight.data.normal_(mean=0.0, std=0.02)   # Matches HF's implementation of llama3
+    def forward(self, vision_outputs: torch.Tensor):
+        if self.deep_extract:
+            x = torch.concat((
+				vision_outputs[-2],
+				vision_outputs[3],
+				vision_outputs[7],
+				vision_outputs[13],
+				vision_outputs[20],
+			), dim=-1)
+            assert len(x.shape) == 3, f"Expected 3, got {len(x.shape)}"  # batch, tokens, features
+            assert x.shape[-1] == vision_outputs[-2].shape[-1] * 5, f"Expected {vision_outputs[-2].shape[-1] * 5}, got {x.shape[-1]}"
+        else:
+            x = vision_outputs[-2]
+
+        x = self.ln1(x)
+
+        if self.pos_emb is not None:
+            assert x.shape[-2:] == self.pos_emb.shape, f"Expected {self.pos_emb.shape}, got {x.shape[-2:]}"
+            x = x + self.pos_emb
+
+        x = self.linear1(x)
+        x = self.activation(x)
+        x = self.linear2(x)
+
+        # <|image_start|>, IMAGE, <|image_end|>
+        other_tokens = self.other_tokens(torch.tensor([0, 1], device=self.other_tokens.weight.device).expand(x.shape[0], -1))
+        assert other_tokens.shape == (x.shape[0], 2, x.shape[2]), f"Expected {(x.shape[0], 2, x.shape[2])}, got {other_tokens.shape}"
+        x = torch.cat((other_tokens[:, 0:1], x, other_tokens[:, 1:2]), dim=1)
+        return x
+
+    def get_eot_embedding(self):
+        return self.other_tokens(torch.tensor([2], device=self.other_tokens.weight.device)).squeeze(0)
+
+
+class JoyImageAdapter:
+    def __init__(self):
+        self.load_device = comfy.model_management.text_encoder_device()
+        self.offload_device = comfy.model_management.text_encoder_offload_device()
+
+        # Image Adapter
+        adapter_path = os.path.join(BASE_MODEL_PATH, "image_adapter.pt")
+
+        image_adapter = ImageAdapter(1152, 4096, False, False, 38,
+                                     False)  # ImageAdapter(clip_model.config.hidden_size, 4096)
+        image_adapter.load_state_dict(torch.load(adapter_path, map_location="cpu", weights_only=True))
+        image_adapter.eval()
+        self.image_adapter = image_adapter
+
+        self.patcher = ModelPatcher(self.image_adapter, load_device=self.load_device,
+                                    offload_device=self.offload_device)
+
+    def embedded_image(self, hidden_states):
+        #print(f"{id(self)}之前device in JoyImageAdapter: {next(self.image_adapter.parameters()).device}")  # 打印模型参数的设备
+        load_models_gpu([self.patcher], force_full_load=True, force_patch_weights=True)
+        #print(f"之后device in JoyImageAdapter: {next(self.image_adapter.parameters()).device}")  # 打印模型参数的设备
+        embedded_images = self.image_adapter(hidden_states)
+        embedded_images.to("cuda")
+        return embedded_images
+
+class JoyLLM:
+    def __init__(self):
+        self.load_device = comfy.model_management.text_encoder_device()
+        self.offload_device = comfy.model_management.text_encoder_offload_device()
+        self.type = comfy.model_management.should_use_fp16()
+
+        print("Loading tokenizer")
+        tokenizer = AutoTokenizer.from_pretrained(os.path.join(BASE_MODEL_PATH, "text_model"), use_fast=True)
+        assert isinstance(tokenizer, PreTrainedTokenizer) or isinstance(tokenizer,
+                                                                        PreTrainedTokenizerFast), f"Tokenizer is of type {type(tokenizer)}"
+
+        self.tokenizer = tokenizer
+        self.text_model = None
+
+    def load_llm_model(self, model_id):
+        if self.text_model is None:
+            print("Loading LLM")
+            LLM_PATH = download_hg_model(model_id, "LLM")
+            text_model_path = os.path.join(BASE_MODEL_PATH, "text_model")
+            modify_json_value(os.path.join(text_model_path, "adapter_config.json"), "base_model_name_or_path",
+                              LLM_PATH)
+            max_retries = 5  # 设置最大重试次数
+            retries = 0
+            while True:
+                free_vram = get_free_memory()/1024/1024
+                print(f"现在的显存{retries}:{free_vram}")
+                if free_vram > 6400:
+                    text_model = AutoModelForCausalLM.from_pretrained(text_model_path,
+                                                              device_map="auto",
+                                                              local_files_only=True,
+                                                              trust_remote_code=True, torch_dtype=torch.bfloat16)
+                    text_model.eval()
+                    self.text_model = text_model
+                    break
+                else:
+                    gc.collect()
+                    unload_all_models()
+                    soft_empty_cache()
+                    retries += 1
+                    if retries > max_retries:
+                        text_model = AutoModelForCausalLM.from_pretrained(text_model_path,
+                                                                          device_map="auto",
+                                                                          local_files_only=True,
+                                                                          trust_remote_code=True,
+                                                                          torch_dtype=torch.bfloat16)
+                        text_model.eval()
+                        self.text_model = text_model
+                        break
+                    time.sleep(1 + retries / 2)
+            print(f"现在呢:{get_free_memory()/1024/1024}")
+        return self.text_model
+
+    def clear_gpu(self, low_vram):
+        del self.text_model
+        self.text_model = None
+        torch.cuda.empty_cache()
+        import gc
+        gc.collect()
+        if low_vram:
+            unload_all_models()
+            soft_empty_cache()
+
+
+class JoyTwoPipeline:
+    def __init__(self):
+        self.clip_model: JoyClipVisionModel | None = None
+        self.image_adapter: JoyImageAdapter | None = None
+        self.llm: JoyLLM | None = None
+        self.parent = None
+        self.model = None
+
+    def clearCache(self):
+        self.clip_model = None
+        self.image_adapter = None
+        self.model = None
+
+    def loadModels(self):
+        # clip
+        self.clip_model = JoyClipVisionModel()
+
+        self.image_adapter = JoyImageAdapter()
+
+    def loadLLM(self):
+        self.llm = JoyLLM()
+
+
+class Joy_caption_two_load:
+
+    def __init__(self):
+        self.model = None
+        self.pipeline = JoyTwoPipeline()
+        self.pipeline.parent = self
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        models = joy_config["model"]
+        return {
+            "required": {
+                "model": (models, ),
+            }
+        }
+
+    CATEGORY = "SLK/LLM"
+    RETURN_TYPES = ("JoyTwoPipeline",)
+    FUNCTION = "generate"
+
+    def loadModels(self):
+        self.pipeline.loadModels()
+
+    def generate(self, model):
+        if self.model is None or self.model != model or self.pipeline is None:
+            self.model = model
+            self.loadModels()
+        self.pipeline.model = model
+        return (self.pipeline,)
+
+
+class Joy_caption_two:
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        caption_lengths = list(joy_config["CAPTION_LENGTH"])
+        caption_types = list(joy_config["CAPTION_TYPE_MAP"].keys())
+        return {
+            "required": {
+                "joy_two_pipeline": ("JoyTwoPipeline",),
+                "image": ("IMAGE",),
+                "caption_type": (caption_types, {}),
+                "caption_length": (caption_lengths, {"default": "long"}),
+                "low_vram": ("BOOLEAN", {"default": False}),
+            }
+        }
+
+    CATEGORY = "SLK/LLM"
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "generate"
+
+    def generate(self, joy_two_pipeline: JoyTwoPipeline, image, caption_type, caption_length, low_vram):
+        torch.cuda.empty_cache()
+
+        if joy_two_pipeline.clip_model is None:
+            joy_two_pipeline.parent.loadModels()
+
+        # 'any' means no length specified
+        length = None if caption_length == "any" else caption_length
+
+        if isinstance(length, str):
+            try:
+                length = int(length)
+            except ValueError:
+                pass
+
+        # Build prompt
+        if length is None:
+            map_idx = 0
+        elif isinstance(length, int):
+            map_idx = 1
+        elif isinstance(length, str):
+            map_idx = 2
+        else:
+            raise ValueError(f"Invalid caption length: {length}")
+
+        caption_type_map = joy_config["CAPTION_TYPE_MAP"]
+        prompt_str = list(caption_type_map[caption_type])[map_idx]
+
+        prompt_str = prompt_str.format(length=caption_length, word_count=caption_length)
+
+        # For debugging
+        # print(f"Prompt: {prompt_str}")
+
+        # Preprocess image
+        # NOTE: I found the default processor for so400M to have worse results than just using PIL directly
+        # image = clip_processor(images=input_image, return_tensors='pt').pixel_values
+        image = tensor2pil(image)
+        image = image.resize((384, 384), Image.LANCZOS)
+        pixel_values = TVF.pil_to_tensor(image).unsqueeze(0) / 255.0
+        pixel_values = TVF.normalize(pixel_values, [0.5], [0.5])
+        pixel_values = pixel_values.to('cuda')
+
+        # Embed image
+        # This results in Batch x Image Tokens x Features
+        with torch.amp.autocast_mode.autocast('cuda', enabled=True):
+            vision_outputs = joy_two_pipeline.clip_model.encode_image(pixel_values)
+            embedded_images = joy_two_pipeline.image_adapter.embedded_image(vision_outputs.hidden_states)
+
+        if low_vram:
+            pixel_values.to("cpu")
+            unload_all_models()
+
+        # Build the conversation
+        convo = [
+            {
+                "role": "system",
+                "content": "You are a helpful image captioner.",
+            },
+            {
+                "role": "user",
+                "content": prompt_str,
+            },
+        ]
+
+        if joy_two_pipeline.llm is None:
+            joy_two_pipeline.loadLLM()
+
+        tokenizer = joy_two_pipeline.llm.tokenizer
+        # Format the conversation
+        convo_string = tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
+        assert isinstance(convo_string, str)
+
+        # Tokenize the conversation
+        # prompt_str is tokenized separately so we can do the calculations below
+        convo_tokens = tokenizer.encode(convo_string, return_tensors="pt", add_special_tokens=False, truncation=False)
+        prompt_tokens = tokenizer.encode(prompt_str, return_tensors="pt", add_special_tokens=False, truncation=False)
+        assert isinstance(convo_tokens, torch.Tensor) and isinstance(prompt_tokens, torch.Tensor)
+        convo_tokens = convo_tokens.squeeze(0)  # Squeeze just to make the following easier
+        prompt_tokens = prompt_tokens.squeeze(0)
+
+        # Calculate where to inject the image
+        eot_id_indices = (convo_tokens == tokenizer.convert_tokens_to_ids("<|eot_id|>")).nonzero(as_tuple=True)[
+            0].tolist()
+        assert len(eot_id_indices) == 2, f"Expected 2 <|eot_id|> tokens, got {len(eot_id_indices)}"
+
+        preamble_len = eot_id_indices[1] - prompt_tokens.shape[0]  # Number of tokens before the prompt
+
+
+        text_model = joy_two_pipeline.llm.load_llm_model(joy_two_pipeline.model)
+        # Embed the tokens
+        convo_embeds = text_model.model.embed_tokens(convo_tokens.unsqueeze(0).to('cuda'))
+
+        # Construct the input
+        input_embeds = torch.cat([
+            convo_embeds[:, :preamble_len],  # Part before the prompt
+            embedded_images.to(dtype=convo_embeds.dtype),  # Image
+            convo_embeds[:, preamble_len:],  # The prompt and anything after it
+        ], dim=1).to('cuda')
+
+        input_ids = torch.cat([
+            convo_tokens[:preamble_len].unsqueeze(0),
+            torch.zeros((1, embedded_images.shape[1]), dtype=torch.long),
+            # Dummy tokens for the image (TODO: Should probably use a special token here so as not to confuse any generation algorithms that might be inspecting the input)
+            convo_tokens[preamble_len:].unsqueeze(0),
+        ], dim=1).to('cuda')
+        attention_mask = torch.ones_like(input_ids)
+
+        # Debugging
+        # print(f"Input to model: {repr(tokenizer.decode(input_ids[0]))}")
+
+        # generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=False, suppress_tokens=None)
+        # generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=True, top_k=10, temperature=0.5, suppress_tokens=None)
+        generate_ids = text_model.generate(input_ids, inputs_embeds=input_embeds, attention_mask=attention_mask,
+                                           max_new_tokens=300, do_sample=True,
+                                           suppress_tokens=None)  # Uses the default which is temp=0.6, top_p=0.9
+
+        # Trim off the prompt
+        generate_ids = generate_ids[:, input_ids.shape[1]:]
+        if generate_ids[0][-1] == tokenizer.eos_token_id or generate_ids[0][-1] == tokenizer.convert_tokens_to_ids(
+                "<|eot_id|>"):
+            generate_ids = generate_ids[:, :-1]
+
+        caption = tokenizer.batch_decode(generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)[0]
+
+        joy_two_pipeline.llm.clear_gpu(low_vram)
+
+        return (caption.strip(), )
+
+class Joy_caption_two_advanced:
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        caption_lengths = list(joy_config["CAPTION_LENGTH"])
+        caption_types = list(joy_config["CAPTION_TYPE_MAP"].keys())
+        return {
+            "required": {
+                "joy_two_pipeline": ("JoyTwoPipeline",),
+                "image": ("IMAGE",),
+                "extra_options": ("Extra_Options", ),
+                "caption_type": (caption_types, {}),
+                "caption_length": (caption_lengths, {"default": "long"}),
+                "name": ("STRING", {"default": ""}),
+                "custom_prompt": ("STRING", {"default": ""}),
+                "low_vram": ("BOOLEAN", {"default": False}),
+            }
+        }
+
+    CATEGORY = "SLK/LLM"
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "generate"
+
+    def generate(self, joy_two_pipeline: JoyTwoPipeline, image, extra_options, caption_type, caption_length, name, custom_prompt, low_vram):
+        torch.cuda.empty_cache()
+
+        if joy_two_pipeline.clip_model == None:
+            joy_two_pipeline.parent.loadModels()
+
+        # 'any' means no length specified
+        length = None if caption_length == "any" else caption_length
+
+        if isinstance(length, str):
+            try:
+                length = int(length)
+            except ValueError:
+                pass
+
+        # Build prompt
+        if length is None:
+            map_idx = 0
+        elif isinstance(length, int):
+            map_idx = 1
+        elif isinstance(length, str):
+            map_idx = 2
+        else:
+            raise ValueError(f"Invalid caption length: {length}")
+
+        caption_type_map = joy_config["CAPTION_TYPE_MAP"]
+        prompt_str = list(caption_type_map[caption_type])[map_idx]
+
+        # Add extra options
+        if len(extra_options) > 0:
+            prompt_str += " " + " ".join(extra_options)
+
+        # Add name, length, word_count
+        prompt_str = prompt_str.format(name=name, length=caption_length, word_count=caption_length)
+
+        if custom_prompt.strip() != "":
+            prompt_str = custom_prompt.strip()
+
+        # For debugging
+        print(f"Prompt: {prompt_str}")
+
+        # Preprocess image
+        # NOTE: I found the default processor for so400M to have worse results than just using PIL directly
+        # image = clip_processor(images=input_image, return_tensors='pt').pixel_values
+        image = tensor2pil(image)
+        image = image.resize((384, 384), Image.LANCZOS)
+        pixel_values = TVF.pil_to_tensor(image).unsqueeze(0) / 255.0
+        pixel_values = TVF.normalize(pixel_values, [0.5], [0.5])
+        pixel_values = pixel_values.to('cuda')
+
+        # Embed image
+        # This results in Batch x Image Tokens x Features
+        with torch.amp.autocast_mode.autocast('cuda', enabled=True):
+            vision_outputs = joy_two_pipeline.clip_model.encode_image(pixel_values)
+            embedded_images = joy_two_pipeline.image_adapter.embedded_image(vision_outputs.hidden_states)
+
+        if low_vram:
+            pixel_values.to("cpu")
+            unload_all_models()
+
+        # Build the conversation
+        convo = [
+            {
+                "role": "system",
+                "content": "You are a helpful image captioner.",
+            },
+            {
+                "role": "user",
+                "content": prompt_str,
+            },
+        ]
+
+        if joy_two_pipeline.llm is None:
+            joy_two_pipeline.loadLLM()
+
+        tokenizer = joy_two_pipeline.llm.tokenizer
+        # Format the conversation
+        convo_string = tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
+        assert isinstance(convo_string, str)
+
+        # Tokenize the conversation
+        # prompt_str is tokenized separately so we can do the calculations below
+        convo_tokens = tokenizer.encode(convo_string, return_tensors="pt", add_special_tokens=False, truncation=False)
+        prompt_tokens = tokenizer.encode(prompt_str, return_tensors="pt", add_special_tokens=False, truncation=False)
+        assert isinstance(convo_tokens, torch.Tensor) and isinstance(prompt_tokens, torch.Tensor)
+        convo_tokens = convo_tokens.squeeze(0)  # Squeeze just to make the following easier
+        prompt_tokens = prompt_tokens.squeeze(0)
+
+        # Calculate where to inject the image
+        eot_id_indices = (convo_tokens == tokenizer.convert_tokens_to_ids("<|eot_id|>")).nonzero(as_tuple=True)[
+            0].tolist()
+        assert len(eot_id_indices) == 2, f"Expected 2 <|eot_id|> tokens, got {len(eot_id_indices)}"
+
+        preamble_len = eot_id_indices[1] - prompt_tokens.shape[0]  # Number of tokens before the prompt
+
+
+        text_model = joy_two_pipeline.llm.load_llm_model(joy_two_pipeline.model)
+        # Embed the tokens
+        convo_embeds = text_model.model.embed_tokens(convo_tokens.unsqueeze(0).to('cuda'))
+        print(f"convo_embeds device: {convo_embeds.device}")  # 打印 convo_embeds 的设备
+        # Construct the input
+        input_embeds = torch.cat([
+            convo_embeds[:, :preamble_len],  # Part before the prompt
+            embedded_images.to(dtype=convo_embeds.dtype),  # Image
+            convo_embeds[:, preamble_len:],  # The prompt and anything after it
+        ], dim=1).to('cuda')
+
+        input_ids = torch.cat([
+            convo_tokens[:preamble_len].unsqueeze(0),
+            torch.zeros((1, embedded_images.shape[1]), dtype=torch.long),
+            # Dummy tokens for the image (TODO: Should probably use a special token here so as not to confuse any generation algorithms that might be inspecting the input)
+            convo_tokens[preamble_len:].unsqueeze(0),
+        ], dim=1).to('cuda')
+        attention_mask = torch.ones_like(input_ids)
+
+        # Debugging
+        # print(f"Input to model: {repr(tokenizer.decode(input_ids[0]))}")
+
+        # generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=False, suppress_tokens=None)
+        # generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=True, top_k=10, temperature=0.5, suppress_tokens=None)
+        generate_ids = text_model.generate(input_ids, inputs_embeds=input_embeds, attention_mask=attention_mask,
+                                           max_new_tokens=300, do_sample=True,
+                                           suppress_tokens=None)  # Uses the default which is temp=0.6, top_p=0.9
+
+        # Trim off the prompt
+        generate_ids = generate_ids[:, input_ids.shape[1]:]
+        if generate_ids[0][-1] == tokenizer.eos_token_id or generate_ids[0][-1] == tokenizer.convert_tokens_to_ids(
+                "<|eot_id|>"):
+            generate_ids = generate_ids[:, :-1]
+
+        caption = tokenizer.batch_decode(generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)[0]
+
+        joy_two_pipeline.llm.clear_gpu(low_vram)
+
+        return (caption.strip(), )
+
+class Joy_extra_options:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        options = list(joy_config["EXTRA_OPTIONS"])
+        required = {}
+        for option in options:
+            required[option] = ("BOOLEAN", {"default": False})
+        return {
+            "required": required
+        }
+
+    CATEGORY = "SLK/LLM"
+    RETURN_TYPES = ("Extra_Options",)
+    FUNCTION = "run"
+
+    def run(self, **kwargs):
+        # 转为列表
+        options_selected = list(kwargs.values())
+        options = list(joy_config["EXTRA_OPTIONS"])
+        values = []
+        for selected, option in zip(options_selected, options):
+            if selected:
+                values.append(option)
+        return (values, )

joy_config.json

@@ -0,0 +1,105 @@
+{
+  "CAPTION_TYPE_MAP": {
+    "Descriptive": [
+      "Write a descriptive caption for this image in a formal tone.",
+      "Write a descriptive caption for this image in a formal tone within {word_count} words.",
+      "Write a {length} descriptive caption for this image in a formal tone."
+    ],
+    "Descriptive (Informal)": [
+      "Write a descriptive caption for this image in a casual tone.",
+      "Write a descriptive caption for this image in a casual tone within {word_count} words.",
+      "Write a {length} descriptive caption for this image in a casual tone."
+    ],
+    "Training Prompt": [
+      "Write a stable diffusion prompt for this image.",
+      "Write a stable diffusion prompt for this image within {word_count} words.",
+      "Write a {length} stable diffusion prompt for this image."
+    ],
+    "MidJourney": [
+      "Write a MidJourney prompt for this image.",
+      "Write a MidJourney prompt for this image within {word_count} words.",
+      "Write a {length} MidJourney prompt for this image."
+    ],
+    "Booru tag list": [
+      "Write a list of Booru tags for this image.",
+      "Write a list of Booru tags for this image within {word_count} words.",
+      "Write a {length} list of Booru tags for this image."
+    ],
+    "Booru-like tag list": [
+      "Write a list of Booru-like tags for this image.",
+      "Write a list of Booru-like tags for this image within {word_count} words.",
+      "Write a {length} list of Booru-like tags for this image."
+    ],
+    "Art Critic": [
+      "Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc.",
+      "Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc. Keep it within {word_count} words.",
+      "Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc. Keep it {length}."
+    ],
+    "Product Listing": [
+      "Write a caption for this image as though it were a product listing.",
+      "Write a caption for this image as though it were a product listing. Keep it under {word_count} words.",
+      "Write a {length} caption for this image as though it were a product listing."
+    ],
+    "Social Media Post": [
+      "Write a caption for this image as if it were being used for a social media post.",
+      "Write a caption for this image as if it were being used for a social media post. Limit the caption to {word_count} words.",
+      "Write a {length} caption for this image as if it were being used for a social media post."
+    ]
+  },
+  "CAPTION_LENGTH": [
+    "any",
+    "very short",
+    "short",
+    "medium-length",
+    "long",
+    "very long",
+    "20",
+    "30",
+    "40",
+    "50",
+    "60",
+    "70",
+    "80",
+    "90",
+    "100",
+    "110",
+    "120",
+    "130",
+    "140",
+    "150",
+    "160",
+    "170",
+    "180",
+    "190",
+    "200",
+    "210",
+    "220",
+    "230",
+    "240",
+    "250",
+    "260"
+  ],
+  "EXTRA_OPTIONS": [
+    "If there is a person/character in the image you must refer to them as {name}.",
+	"Do NOT include information about people/characters that cannot be changed (like ethnicity, gender, etc), but do still include changeable attributes (like hair style).",
+	"Include information about lighting.",
+	"Include information about camera angle.",
+	"Include information about whether there is a watermark or not.",
+	"Include information about whether there are JPEG artifacts or not.",
+	"If it is a photo you MUST include information about what camera was likely used and details such as aperture, shutter speed, ISO, etc.",
+	"Do NOT include anything sexual; keep it PG.",
+	"Do NOT mention the image's resolution.",
+	"You MUST include information about the subjective aesthetic quality of the image from low to very high.",
+	"Include information on the image's composition style, such as leading lines, rule of thirds, or symmetry.",
+	"Do NOT mention any text that is in the image.",
+	"Specify the depth of field and whether the background is in focus or blurred.",
+	"If applicable, mention the likely use of artificial or natural lighting sources.",
+	"Do NOT use any ambiguous language.",
+	"Include whether the image is sfw, suggestive, or nsfw.",
+	"ONLY describe the most important elements of the image."
+  ],
+  "model": [
+    "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
+    "meta-llama/Meta-Llama-3.1-8B"
+  ]
+}

joy_config.py

@@ -0,0 +1,5 @@
+import os
+from .uitls import read_json_file
+
+joy_base_path = os.path.dirname(os.path.realpath(__file__))
+joy_config = read_json_file(os.path.join(joy_base_path, "joy_config.json"))

readme.md

@@ -0,0 +1,68 @@
+
+## Recent changes 
+* [2024-10-12] v0.0.1: 基本完成[JoyCaptionAlpha Two](https://huggingface.co/spaces/fancyfeast/joy-caption-alpha-two)到ComfyUI的实现
+
+
+## ComfyUI上JoyCaptionAlpha Two的实现
+
+参考自 [Comfyui_CXH_joy_caption](https://github.com/StartHua/Comfyui_CXH_joy_caption), 以及 [JoyCaptionAlpha Two](https://huggingface.co/spaces/fancyfeast/joy-caption-alpha-two)
+
+参考工作流在examples/workflow.png中获取:
+![image](./examples/workflow.png)
+
+### 依赖安装
+
+1. 把仓库下载克隆到 custom_nodes 子文件夹下。
+```
+cd custom_nodes
+git clone https://github.com/EvilBT/ComfyUI_SLK_joy_caption_two.git
+```
+2. 安装相关依赖：
+```angular2html
+pip install -r ComfyUI_SLK_joy_caption_two\requirements.txt
+```
+3. 下载相关模型。
+4. 重启ComfyUI。
+
+### 相关模型下载
+以下的models目录是指ComfyUI根目录下的models文件夹
+#### 1. google/siglip-so400m-patch14-384:
+
+国外：[google/siglip-so400m-patch14-384](https://huggingface.co/google/siglip-so400m-patch14-384)
+
+国内：[hf/google/siglip-so400m-patch14-384](https://hf-mirror.com/google/siglip-so400m-patch14-384)
+
+会自动下载，也可以手动下载整个仓库，并把siglip-so400m-patch14-384内的文件全部复制到`models/clip/siglip-so400m-patch14-384`
+
+#### 2. Llama3.1-8B 模型下载
+
+支持两个版本：bnb-4bit是小显存的福音，我是使用这个版本的，原版的我没有测试过，可自行测试。程序会自动下载，可自行下载。
+
+2.1 [unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit](https://huggingface.co/unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit)
+
+国内可以从镜像网站下载[hf/unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit](https://hf-mirror.com/unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit)
+
+把整个文件夹内的内容复制到 `models\LLM\Meta-Llama-3.1-8B-Instruct-bnb-4bit` 下
+
+2.2 meta-llama/Meta-Llama-3.1-8B
+
+国外:[meta-llama/Meta-Llama-3.1-8B](https://huggingface.co/meta-llama/Llama-3.1-8B)
+
+国内:[meta-llama/Meta-Llama-3.1-8B](https://hf-mirror.com/meta-llama/Llama-3.1-8B)
+
+把下载后的整个文件夹的内容复制到`models\LLM\Meta-Llama-3.1-8B`下
+
+
+#### 3. Joy-Caption-alpha-two 模型下载（必须手动下载）
+
+把 [Joy-Caption-alpha-two](https://huggingface.co/spaces/fancyfeast/joy-caption-alpha-two/tree/main) 下的`cgrkzexw-599808`
+文件夹的所有内容下载复制到`models/Joy_caption_two` 下
+
+### 重装ComfyUI之后就可以添加使用了，具体可以参考上面的图片
+
+### 其他
+
+如果你安装了 [AIGODLIKE-ComfyUI-Translation](https://github.com/AIGODLIKE/AIGODLIKE-ComfyUI-Translation) 语言包插件，你可以复制 `translation` 文件夹下的中文翻译到对应的语言包路径下，重启就可以使用中文版的了。
+把 `translation/zh-CN/Nodes/Comfyui_SLK_joy_caption_two.json` 复制到目录：`AIGODLIKE-ComfyUI-Translation\zh-CN\Nodes` 即可
+
+有问题可以开issue问我，未完全测试，我是8G显存的环境

requirements.txt

@@ -0,0 +1,7 @@
+huggingface_hub==0.23.4
+transformers>=4.44.0
+numpy
+sentencepiece
+pillow>=10.1.0
+bitsandbytes>=0.44.1
+peft>=0.12.0

translation/zh-CN/Nodes/Comfyui_SLK_joy_caption_two.json

@@ -0,0 +1,92 @@
+{
+    "Joy_caption_two_load": {
+        "title": "加载JoyCaptionTwo",
+        "widgets": {
+            "model": "模型"
+        },
+        "outputs": {
+            "JoyTwoPipeline": "JoyCaptionTwo"
+        }
+    },
+    "Joy_caption_two": {
+        "title": "JoyCaptionTwo",
+        "inputs": {
+            "joy_two_pipeline": "JoyCaptionTwo",
+            "image": "图像"
+        },
+        "widgets": {
+            "low_vram": "低显存",
+            "caption_type": "提示词类型",
+                "Descriptive": "描述性",
+                "Descriptive (Informal)": "描述性（非正式）",
+                "Training Prompt": "训练提示",
+                "MidJourney": "MidJourney",
+                "Booru tag list": "Booru 标签列表",
+                "Booru-like tag list": "类似 Booru 的标签列表",
+                "Art Critic": "艺术评论家",
+                "Product Listing": "产品清单",
+                "Social Media Post": "社交媒体帖子",
+            "caption_length": "提示词长度",
+                "any": "任意",
+                "very short": "非常短",
+                "short": "短",
+                "medium-length": "中等长度",
+                "long": "长",
+                "very long": "非常长"
+        },
+        "outputs": {
+            "STRING": "提示词"
+        }
+    },
+    "Joy_caption_two_advanced": {
+        "title": "高级JoyCaptionTwo",
+        "inputs": {
+            "joy_two_pipeline": "JoyCaptionTwo",
+            "image": "图像",
+            "extra_options": "附加选项"
+        },
+        "widgets": {
+            "name": "人名",
+            "custom_prompt": "自定义引导词",
+            "caption_type": "提示词类型",
+                "Descriptive": "描述性",
+                "Descriptive (Informal)": "描述性（非正式）",
+            "caption_length": "提示词长度",
+                "any": "任意",
+                "very short": "非常短",
+                "short": "短",
+                "medium-length": "中等长度",
+                "long": "长",
+                "very long": "非常长",
+            "low_vram": "低显存"
+        },
+        "outputs": {
+            "STRING": "提示词"
+        }
+    },
+    "Joy_extra_options": {
+        "title": "JoyCaption附加选项",
+        "widgets": {
+            "If there is a person/character in the image you must refer to them as {name}.": "如果图像中有人物/角色，你必须将其称为 {name}。",
+            "Do NOT include information about people/characters that cannot be changed (like ethnicity, gender, etc), but do still include changeable attributes (like hair style).": "不要包含关于人物/角色不可改变的信息（如种族、性别等），但仍然要包含可改变的属性（如发型）。",
+            "Include information about lighting.": "包含有关光照的信息。",
+            "Include information about camera angle.": "包含有关相机角度的信息。",
+            "Include information about whether there is a watermark or not.": "包含是否有水印的信息。",
+            "Include information about whether there are JPEG artifacts or not.": "包含是否有 JPEG 伪影的信息。",
+            "If it is a photo you MUST include information about what camera was likely used and details such as aperture, shutter speed, ISO, etc.": "如果它是照片，你必须包含有关可能使用的相机的信息，以及光圈、快门速度、ISO 等细节。",
+            "Do NOT include anything sexual; keep it PG.": "不要包含任何色情内容；保持 PG 级。",
+            "Do NOT mention the image's resolution.": "不要提及图像的分辨率。",
+            "You MUST include information about the subjective aesthetic quality of the image from low to very high.": "你必须包含关于图像主观审美质量的信息，从低到非常高。",
+            "Include information on the image's composition style, such as leading lines, rule of thirds, or symmetry.": "包含有关图像构图风格的信息，例如引导线、三分法则或对称性。",
+            "Do NOT mention any text that is in the image.": "不要提及图像中的任何文字。",
+            "Specify the depth of field and whether the background is in focus or blurred.": "指定景深以及背景是聚焦还是模糊。",
+            "If applicable, mention the likely use of artificial or natural lighting sources.": "如果适用，请提及可能使用的人工或自然光源。",
+            "Do NOT use any ambiguous language.": "不要使用任何模棱两可的语言。",
+            "Include whether the image is sfw, suggestive, or nsfw.": "包括图像是否是 sfw、暗示性或 nsfw。",
+            "ONLY describe the most important elements of the image.": "只描述图像中最重要的元素。"
+    },
+        "outputs": {
+            "Extra_Options": "附加选项"
+        }
+    }
+}

uitls.py

@@ -0,0 +1,56 @@
+import json
+import os
+import folder_paths
+
+# 下载hg 模型到本地
+def download_hg_model(model_id:str, exDir:str=''):
+    # 下载本地
+    model_checkpoint = os.path.join(folder_paths.models_dir, exDir, os.path.basename(model_id))
+    print(model_checkpoint)
+    if not os.path.exists(model_checkpoint):
+        from huggingface_hub import snapshot_download
+        snapshot_download(repo_id=model_id, local_dir=model_checkpoint, local_dir_use_symlinks=False)
+    return model_checkpoint
+
+def modify_json_value(file_path, key_to_modify, new_value):
+  """
+  读取 JSON 文件，修改指定 key 的 value 值，并保存修改后的文件。
+
+  Args:
+    file_path: JSON 文件路径。
+    key_to_modify: 需要修改的 key。
+    new_value:  新的 value 值。
+  """
+  try:
+    with open(file_path, 'r', encoding='utf-8') as f:
+      data = json.load(f)
+
+    # 查找并修改 key 的 value
+    if key_to_modify in data:
+      data[key_to_modify] = new_value
+    else:
+      print(f"Warning: Key '{key_to_modify}' not found in JSON file.")
+
+    # 保存修改后的 JSON 文件
+    with open(file_path, 'w', encoding='utf-8') as f:
+      json.dump(data, f, indent=4)  # 使用 indent 参数格式化输出
+
+    print(f"Successfully modified '{key_to_modify}' value in '{file_path}'.")
+
+  except FileNotFoundError:
+    print(f"Error: File '{file_path}' not found.")
+  except json.JSONDecodeError:
+    print(f"Error: Invalid JSON format in '{file_path}'.")
+
+def read_json_file(file_path):
+  """读取 JSON 文件并转换为 Python 字典。"""
+  try:
+    with open(file_path, 'r', encoding='utf-8') as f:
+      data = json.load(f)
+    return data
+  except FileNotFoundError:
+    print(f"Error: File '{file_path}' not found.")
+    return None
+  except json.JSONDecodeError:
+    print(f"Error: Invalid JSON format in '{file_path}'.")
+    return None