diff --git a/examples/workflows.png b/examples/workflows.png new file mode 100644 index 0000000..ed4b9b6 Binary files /dev/null and b/examples/workflows.png differ diff --git a/joy_caption_two_node.py b/joy_caption_two_node.py index 9467774..d1e61fe 100644 --- a/joy_caption_two_node.py +++ b/joy_caption_two_node.py @@ -431,6 +431,8 @@ class Joy_caption_two_advanced: "name": ("STRING", {"default": ""}), "custom_prompt": ("STRING", {"default": ""}), "low_vram": ("BOOLEAN", {"default": False}), + "top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.01}), + "temperature": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 1.0, "step": 0.01}), } } @@ -438,7 +440,7 @@ class Joy_caption_two_advanced: RETURN_TYPES = ("STRING",) FUNCTION = "generate" - def generate(self, joy_two_pipeline: JoyTwoPipeline, image, extra_options, caption_type, caption_length, name, custom_prompt, low_vram): + def generate(self, joy_two_pipeline: JoyTwoPipeline, image, extra_options, caption_type, caption_length, name, custom_prompt, low_vram, top_p, temperature): torch.cuda.empty_cache() if joy_two_pipeline.clip_model == None: @@ -559,6 +561,7 @@ class Joy_caption_two_advanced: # 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, + top_p=top_p, temperature=temperature, suppress_tokens=None) # Uses the default which is temp=0.6, top_p=0.9 # Trim off the prompt @@ -782,6 +785,8 @@ class Batch_joy_caption_two_advanced: "name": ("STRING", {"default": ""}), "custom_prompt": ("STRING", {"default": ""}), "low_vram": ("BOOLEAN", {"default": False}), + "top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.01}), + "temperature": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 1.0, "step": 0.01}), } } @@ -789,7 +794,7 @@ class Batch_joy_caption_two_advanced: RETURN_TYPES = ("STRING",) FUNCTION = "generate" - def generate_caption(self, joy_two_pipeline: JoyTwoPipeline, image, prompt, low_vram=True): + def generate_caption(self, joy_two_pipeline: JoyTwoPipeline, image, prompt, top_p=0.9, temperature=0.6, low_vram=True): torch.cuda.empty_cache() pixel_values = TVF.pil_to_tensor(image).unsqueeze(0) / 255.0 pixel_values = TVF.normalize(pixel_values, [0.5], [0.5]) @@ -865,6 +870,7 @@ class Batch_joy_caption_two_advanced: # 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, + top_p=top_p, temperature=temperature, suppress_tokens=None) # Uses the default which is temp=0.6, top_p=0.9 # Trim off the prompt @@ -877,7 +883,7 @@ class Batch_joy_caption_two_advanced: return caption.strip() - def generate(self, joy_two_pipeline: JoyTwoPipeline, input_dir, output_dir, extra_options, caption_type, caption_length, name, custom_prompt, low_vram): + def generate(self, joy_two_pipeline: JoyTwoPipeline, input_dir, output_dir, extra_options, caption_type, caption_length, name, custom_prompt, low_vram, top_p, temperature): torch.cuda.empty_cache() if joy_two_pipeline.clip_model == None: @@ -943,7 +949,7 @@ class Batch_joy_caption_two_advanced: img = img.convert('RGB') pbar.update_absolute(step, image_count) image = img.resize((384, 384), Image.LANCZOS) - caption = self.generate_caption(joy_two_pipeline, image, prompt_str) + caption = self.generate_caption(joy_two_pipeline, image, prompt_str, top_p, temperature) with open(text_path, 'w', encoding='utf-8') as f: f.write(caption) finished_image_count += 1 diff --git a/joy_config.json b/joy_config.json index f71ba9c..8ceb480 100644 --- a/joy_config.json +++ b/joy_config.json @@ -100,6 +100,8 @@ ], "model": [ "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit", - "unsloth/Meta-Llama-3.1-8B-Instruct" + "unsloth/Meta-Llama-3.1-8B-Instruct", + "John6666/Llama-3.1-8B-Lexi-Uncensored-V2-nf4", + "Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2" ] } \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index ff9c1bd..d39b737 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui_slk_joy_caption_two" description = "NODES:Joy Caption Two, Joy Caption Two Advanced, Joy Caption Two Load, Joy Caption Extra Options" -version = "0.0.5" +version = "0.0.6" license = {file = "LICENSE"} dependencies = ["huggingface_hub==0.23.4", "transformers>=4.44.0", "numpy", "sentencepiece", "pillow>=10.1.0", "bitsandbytes>=0.44.1", "peft==0.12.0"] diff --git a/readme.md b/readme.md index 4c7e500..5f9991a 100644 --- a/readme.md +++ b/readme.md @@ -2,6 +2,8 @@ [English](./readme_us.md) | 中文 ## Recent changes +* [2024-10-16] v0.0.6: 高级模式增加top_p与temperature，给予更多的选择，添加更多的大模型选择，我试了一下 [John6666/Llama-3.1-8B-Lexi-Uncensored-V2-nf4](https://huggingface.co/John6666/Llama-3.1-8B-Lexi-Uncensored-V2-nf4) +效果不错，你们也可以尝试使用，另外也添加了原版的模型 [Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2](https://huggingface.co/Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2)，可以自行选择 * [2024-10-15] v0.0.5: 修复批处理时图片有透明通道 RGBA 时的BUG * [2024-10-15] v0.0.4: 添加指量处理节点：字幕保存目录为空时则保存在图片文件夹下，参考工作流可以在examples目录下查看。 * [2024-10-15] v0.0.3: 修复'cuda:0'部分出错的问题，直接设置为 'cuda' @@ -13,9 +15,8 @@ 参考自 [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) -![image](./examples/batch_workflow.png) +参考工作流在examples/workflows.png中获取: +![image](./examples/workflows.png) ### 安装 @@ -72,8 +73,7 @@ pip install -r ComfyUI_SLK_joy_caption_two\requirements.txt 文件夹的所有内容下载复制到`models/Joy_caption_two` 下 ![image](./examples/joy_caption.png) ### 重启ComfyUI之后就可以添加使用了，具体可以参考下面的图片 -![image](./examples/workflow.png) -![image](./examples/workflow_flux.png) +![image](./examples/workflows.png) ### 其他 diff --git a/readme_us.md b/readme_us.md index 180edf1..9732935 100644 --- a/readme_us.md +++ b/readme_us.md @@ -1,6 +1,7 @@ # JoyCaptionAlpha Two for ComfyUI English | [中文](./readme.md) ## Recent changes +* [2024-10-16] v0.0.6: Added `top_p` and `temperature` parameters to the advanced mode for greater control. Expanded the selection of large language models. I tested [John6666/Llama-3.1-8B-Lexi-Uncensored-V2-nf4](https://huggingface.co/John6666/Llama-3.1-8B-Lexi-Uncensored-V2-nf4) and found the results quite good; you can also try it out. Additionally, the original [Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2](https://huggingface.co/Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2) model has been added as an option. * [2024-10-15] v0.0.5: Fix the bug when processing images with an alpha channel (RGBA) in batch. * [2024-10-15] v0.0.4: Added batch processing nodes: When the output directory is empty, it will be saved in the image folder. You can find the example workflow in the examples directory. * [2024-10-15] v0.0.3: Fixed an issue where specifying 'cuda:0' would partially fail, now defaults to 'cuda' @@ -12,9 +13,8 @@ English | [中文](./readme.md) Referred to [Comfyui_CXH_joy_caption](https://github.com/StartHua/Comfyui_CXH_joy_caption) and [JoyCaptionAlpha Two](https://huggingface.co/spaces/fancyfeast/joy-caption-alpha-two). -Refer to the example workflow in examples/workflow.png: -![image](./examples/workflow.png) -![image](./examples/batch_workflow.png) +Refer to the example workflow in examples/workflows.png: +![image](./examples/workflows.png) ### Installation @@ -80,8 +80,7 @@ Download and copy all the contents of the `cgrkzexw-599808` folder under [Joy-Ca ![image](./examples/joy_caption.png) ### After restarting ComfyUI, you can add and use it. Refer to the following images for details: -![image](./examples/workflow.png) -![image](./examples/workflow_flux.png) +![image](./examples/workflows.png) ### Others diff --git a/translation/zh-CN/Nodes/Comfyui_SLK_joy_caption_two.json b/translation/zh-CN/Nodes/Comfyui_SLK_joy_caption_two.json index 4c80cc3..20597b8 100644 --- a/translation/zh-CN/Nodes/Comfyui_SLK_joy_caption_two.json +++ b/translation/zh-CN/Nodes/Comfyui_SLK_joy_caption_two.json @@ -58,7 +58,9 @@ "medium-length": "中等长度", "long": "长", "very long": "非常长", - "low_vram": "低显存" + "low_vram": "低显存", + "top_p": "概率筛选(top_p)", + "temperature": "创造力温度(temperature)" }, "outputs": { "STRING": "提示词" @@ -118,7 +120,9 @@ "medium-length": "中等长度", "long": "长", "very long": "非常长", - "low_vram": "低显存" + "low_vram": "低显存", + "top_p": "概率筛选(top_p)", + "temperature": "创造力温度(temperature)" }, "outputs": { "STRING": "处理结果"