From 0ca81c19cf717c953913cc75ae6c25d1aae074f6 Mon Sep 17 00:00:00 2001 From: yemiancheng Date: Wed, 11 Jun 2025 14:17:28 +0800 Subject: [PATCH] build(core): init --- .editorconfig | 12 + .gitattributes | 69 +++++ .github/workflows/publish_to_comfy.yml | 21 ++ .gitignore | 2 + LICENSE | 21 ++ README.md | 122 ++++++++ __init__.py | 110 +++++++ conf.py | 2 + docs.template.md | 72 +++++ docs/models.md | 37 +++ docs/nodes.md | 12 + docs/why.md | 3 + joy.py | 382 +++++++++++++++++++++++++ lib/__init__.py | 0 lib/xfile.py | 32 +++ lib/ximg.py | 167 +++++++++++ lib/xmodel.py | 34 +++ pyproject.toml | 15 + requirements.txt | 23 ++ 19 files changed, 1136 insertions(+) create mode 100644 .editorconfig create mode 100644 .gitattributes create mode 100644 .github/workflows/publish_to_comfy.yml create mode 100644 .gitignore create mode 100644 LICENSE create mode 100644 README.md create mode 100644 __init__.py create mode 100644 conf.py create mode 100644 docs.template.md create mode 100644 docs/models.md create mode 100644 docs/nodes.md create mode 100644 docs/why.md create mode 100644 joy.py create mode 100644 lib/__init__.py create mode 100644 lib/xfile.py create mode 100644 lib/ximg.py create mode 100644 lib/xmodel.py create mode 100644 pyproject.toml create mode 100644 requirements.txt diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000..c1322dc --- /dev/null +++ b/.editorconfig @@ -0,0 +1,12 @@ +# EditorConfig is awesome: https://EditorConfig.org + +# top-most EditorConfig file +root = true + +[*] +indent_style = space +indent_size = 4 +end_of_line = lf +charset = utf-8 +trim_trailing_whitespace = false +insert_final_newline = false \ No newline at end of file diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..f187747 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,69 @@ +# s: set *.js and other as js file for js project +# disable other with linguist-detectable for ps1 +*.ps1 linguist-language=false +*.md linguist-detectable=false +*.py linguist-detectable=true +*.html linguist-detectable=false +*.js linguist-detectable=false +*.css linguist-detectable=false +# s: set *.js and other as js file for js project + +# refer: +# https://huggingface.co/guoyww/animatediff/blob/main/.gitattributes +# https://huggingface.co/datasets/hdparmar/itt_specdata/blob/main/.gitattributes +# https://github.com/Richienb/git-files/blob/master/files/.gitattributes + +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text +# Audio files - uncompressed +*.pcm filter=lfs diff=lfs merge=lfs -text +*.sam filter=lfs diff=lfs merge=lfs -text +*.raw filter=lfs diff=lfs merge=lfs -text +# Audio files - compressed +*.aac filter=lfs diff=lfs merge=lfs -text +*.flac filter=lfs diff=lfs merge=lfs -text +*.mp3 filter=lfs diff=lfs merge=lfs -text +*.ogg filter=lfs diff=lfs merge=lfs -text +*.wav filter=lfs diff=lfs merge=lfs -text +# Image files - uncompressed +*.bmp filter=lfs diff=lfs merge=lfs -text +*.gif filter=lfs diff=lfs merge=lfs -text +*.png filter=lfs diff=lfs merge=lfs -text +*.tiff filter=lfs diff=lfs merge=lfs -text +# Image files - compressed +*.jpg filter=lfs diff=lfs merge=lfs -text +*.jpeg filter=lfs diff=lfs merge=lfs -text +*.webp filter=lfs diff=lfs merge=lfs -text \ No newline at end of file diff --git a/.github/workflows/publish_to_comfy.yml b/.github/workflows/publish_to_comfy.yml new file mode 100644 index 0000000..b278ccd --- /dev/null +++ b/.github/workflows/publish_to_comfy.yml @@ -0,0 +1,21 @@ +name: Publish to Comfy registry +on: + workflow_dispatch: + push: + branches: + - main + paths: + - "pyproject.toml" + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + steps: + - name: Check out code + uses: actions/checkout@v4 + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@main + with: + ## Add your own personal access token to your Github Repository secrets and reference it here. + personal_access_token: ${{ secrets.COMFY_REGISTRY_TOKEN }} diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..10a149b --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +models +__pycache__ \ No newline at end of file diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..c91f0a4 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 ymc-github + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..08f93ff --- /dev/null +++ b/README.md @@ -0,0 +1,122 @@ +
+

ymc_node_joy

+

+ 🤖 comfyui custom nodes to caption image with joy +

+ +
+ + + + +## Why + + + +- use joy to caption image for aigc. +- use joy to caption image files in diretory to do sth. (eg. lora training) + + +## Features + + + + +## Nodes + + +- nodes show in console: +
+ console + +
+ +- nodes show in right mouse menu: + +
+ + right mouse menu +
+ +## Install + +```bash +# cd to comfyui/custom_nodes +git clone https://github.com/ymc-github/ymc_node_joy +``` +- **deps will be installed automatically** if deps in requirements.txt were not installed when comfyui up + + + +## Required Models + +Essential components: +- `google/siglip-so400m-patch14-384` (Vision model) +- `unsloth/Meta-Llama-3.1-8B-bnb-4bit` or `meta-llama/Meta-Llama-3.1-8B` (LLM) +- `Joy_caption/image_adapter.pt` (Custom adapter) + +```plaintext +/ +├── models/ +│ ├── clip/ # SigLIP Vision Model +│ │ └── siglip-so400m-patch14-384/ +│ ├── llm/ # Llama Language Model +│ │ ├── Meta-Llama-3.1-8B-bnb-4bit/ +│ │ └── Meta-Llama-3.1-8B/ +│ └── Joy_caption/ # Custom Components +│ └── image_adapter.pt # Dimension Adapter +``` + +### 1. SigLIP Vision Model (google/siglip-so400m-patch14-384) +**International**: https://huggingface.co/google/siglip-so400m-patch14-384 +**China Mirror**: https://hf-mirror.com/google/siglip-so400m-patch14-384 + +### 2. Llama Language Models +#### 4bit Quantized (unsloth/Meta-Llama-3.1-8B-bnb-4bit) +**International**: https://huggingface.co/unsloth/Meta-Llama-3.1-8B-bnb-4bit +**China Mirror**: https://hf-mirror.com/unsloth/Meta-Llama-3.1-8B-bnb-4bit + +#### Original (meta-llama/Meta-Llama-3.1-8B) +**International**: https://huggingface.co/meta-llama/Meta-Llama-3.1-8B (Access approval required) +**China Mirror**: https://hf-mirror.com/meta-llama/Meta-Llama-3.1-8B + +### 3. Image Adapter (Joy_caption/image_adapter.pt) +**International**: https://huggingface.co/spaces/fancyfeast/joy-caption-pre-alpha/tree/main/wpkklhc6 +**China Mirror**: https://www.modelscope.cn/models/fireicewolf/joy-caption-pre-alpha/files + + + + +## Usage + +- you can find it in search box : double click + typing keywords (eg: `joy - `, `caption`) +- you can find it in right mouse menu : `ymc/caption` +- ~~you can find it in right mouse menu : `utils/ymc/caption` (as alias)~~ + +## Based-on + +- pypi package [yors_comfyui_node_setup](https://pypi.org/project/yors_comfyui_node_setup/) - setup comfyui custom nodes easily +- pypi package [yors_pano_ansi_color](https://pypi.org/project/yors_pano_ansi_color/) - info msg in console with color in your comfyui custom nodes easily +- ~~pypi package [yors_pano_zero_field](https://pypi.org/project/yors_pano_zero_field/) - set nodes input field to be HQ in your comfyui custom nodes easily~~ + + +## Published to Comfy registry + +- get more details in [publish_to_comfy.yml](.github/workflows/publish_to_comfy.yml) + +- [docs for publishing to comfy registey](https://docs.comfy.org/registry/overview) + +- installed with comfy-cli ? `comfy node registry-install ymc_node_joy` + +## Author + +ymc-github + +name|email|desciption +:--|:--|:-- +yemiancheng||Main developer and code maintainer| +chenxinghua|<455758525@qq.com>|Code reference from [StartHua/Comfyui_CXH_joy_caption](https://github.com/StartHua/Comfyui_CXH_joy_caption)| + +## License + +MIT \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..b22fbe1 --- /dev/null +++ b/__init__.py @@ -0,0 +1,110 @@ +# Standard library imports +from pathlib import Path +# import os +# import toml + +# Third-party imports +from yors_comfyui_node_setup import entry +from yors_pano_ansi_color import info_status, info_step, msg_padd, log_msg + +# def get_version_from_pyproject(file: str, fallback: str = '1.0.0'): +# """ +# get version from pyproject.toml 's project.version +# """ +# try: +# file = Path(file) +# if not file.exists(): +# raise FileNotFoundError +# if not os.access(file, os.R_OK): +# raise PermissionError(f"Permission denied: {file}") + +# pyproject_content = file.read_text() +# pyproject = toml.loads(pyproject_content) +# version = pyproject.get('project', {}).get('version', fallback) +# except FileNotFoundError: +# log_msg(f"pyproject.toml not found at {file}, using default version {fallback}") +# version = fallback +# except PermissionError as e: +# log_msg(f"Permission error: {e}, using default version {fallback}") +# version = fallback +# except Exception as e: +# log_msg(f"Error reading pyproject.toml: {e}, using default version {fallback}") +# version = fallback +# return version + +# def get_version_from_txt(file: str, fallback: str = '1.0.0'): +# """ +# get version from version.txt file +# """ +# try: +# file = Path(file) +# if not file.exists(): +# raise FileNotFoundError +# if not os.access(file, os.R_OK): +# raise PermissionError(f"Permission denied: {file}") + +# version = file.read_text().strip() +# if not version: +# raise ValueError("Empty version string") +# return version +# except Exception as e: +# log_msg(f"Error reading version.txt: {e}, using default version {fallback}") +# return fallback + +def path_resolve(path: str,root:str): + """ + resolve path to absolute path with root path + + INIT_PY_REl='../' + root = path_resolve(INIT_PY_REl,str(Path(__file__))) + path_resolve(root,'pyproject.toml') + """ + return str(Path(root).joinpath(path).resolve().as_posix()) + +def path_dirname(path: str): + """ + get dirname of path + """ + return str(Path(path).parent) + +def path_parse(path:str): + flag = Path(path) + name = flag.name + stem = flag.stem + suffix = flag.suffix + parent = str(flag.parent.as_posix()) + return (name, stem, suffix, parent) + + +__all__,NODE_CLASS_MAPPINGS,NODE_DISPLAY_NAME_MAPPINGS,NODE_MENU_NAMES = entry(__name__,__file__,False) + +info_step(f"__all__ + WEB_DIRECTORY") +WEB_DIRECTORY = "./web" +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] + +# build(core): get root path with __file and INIT_PY_REl +INIT_PY_REl="../" +root = path_resolve(INIT_PY_REl,str(Path(__file__))) + +# build(core): get name from root path +name, stem, suffix, parent=path_parse(root) + + +# - build(core): get version from pyproject.toml +# pyproject_file=path_resolve(root,'pyproject.toml') +# version= get_version_from_pyproject(Path(pyproject_file)) +# version_file=path_resolve(root,'version.txt') +# version=get_version_from_txt(version_file) + +# name='ymc_node_joy' +# version='1.0.0' +log_msg(msg_padd("=",60,"=")) +log_msg(msg_padd(f'welocme to {name}',60,"=")) +# log_msg(f'version: {version}') +log_msg(f'node counts:{len(NODE_MENU_NAMES)}') +log_msg(f'node menu names:') +NODE_MENU_NAMES.sort() +for node_name in NODE_MENU_NAMES: + # log_msg(f'node name:{node_name}') + info_status(f'{node_name}',0) +log_msg(msg_padd("=",60,"=")) \ No newline at end of file diff --git a/conf.py b/conf.py new file mode 100644 index 0000000..5df810b --- /dev/null +++ b/conf.py @@ -0,0 +1,2 @@ +CURRENT_CATEGORY="ymc/caption" +CURRENT_FUNCTION="exec" \ No newline at end of file diff --git a/docs.template.md b/docs.template.md new file mode 100644 index 0000000..98a1e1c --- /dev/null +++ b/docs.template.md @@ -0,0 +1,72 @@ +
+

{node_name}

+

+ 🤖 {node_desc} +

+ +
+ + + + +## Why + + + + + +## Features + + + + +## Nodes + + + + +## Install + +```bash +# cd to comfyui/custom_nodes +git clone https://github.com/ymc-github/{node_name} +``` +- **deps will be installed automatically** if deps in requirements.txt were not installed when comfyui up + + + + + +## Usage + +- you can find it in search box : double click + typing keywords (eg: `joy - `, `caption`) +- you can find it in right mouse menu : `{node_right_menu}` +- ~~you can find it in right mouse menu : `utils/{node_right_menu}` (as alias)~~ + +## Based-on + +- pypi package [yors_comfyui_node_setup](https://pypi.org/project/yors_comfyui_node_setup/) - setup comfyui custom nodes easily +- pypi package [yors_pano_ansi_color](https://pypi.org/project/yors_pano_ansi_color/) - info msg in console with color in your comfyui custom nodes easily +- ~~pypi package [yors_pano_zero_field](https://pypi.org/project/yors_pano_zero_field/) - set nodes input field to be HQ in your comfyui custom nodes easily~~ + + +## Published to Comfy registry + +- get more details in [publish_to_comfy.yml](.github/workflows/publish_to_comfy.yml) + +- [docs for publishing to comfy registey](https://docs.comfy.org/registry/overview) + +- installed with comfy-cli ? `comfy node registry-install {node_name}` + +## Author + +ymc-github + +name|email|desciption +:--|:--|:-- +yemiancheng||Main developer and code maintainer| +chenxinghua|<455758525@qq.com>|Code reference from [StartHua/Comfyui_CXH_joy_caption](https://github.com/StartHua/Comfyui_CXH_joy_caption)| + +## License + +MIT \ No newline at end of file diff --git a/docs/models.md b/docs/models.md new file mode 100644 index 0000000..45ebb1d --- /dev/null +++ b/docs/models.md @@ -0,0 +1,37 @@ + +## Required Models + +Essential components: +- `google/siglip-so400m-patch14-384` (Vision model) +- `unsloth/Meta-Llama-3.1-8B-bnb-4bit` or `meta-llama/Meta-Llama-3.1-8B` (LLM) +- `Joy_caption/image_adapter.pt` (Custom adapter) + +```plaintext +/ +├── models/ +│ ├── clip/ # SigLIP Vision Model +│ │ └── siglip-so400m-patch14-384/ +│ ├── llm/ # Llama Language Model +│ │ ├── Meta-Llama-3.1-8B-bnb-4bit/ +│ │ └── Meta-Llama-3.1-8B/ +│ └── Joy_caption/ # Custom Components +│ └── image_adapter.pt # Dimension Adapter +``` + +### 1. SigLIP Vision Model (google/siglip-so400m-patch14-384) +**International**: https://huggingface.co/google/siglip-so400m-patch14-384 +**China Mirror**: https://hf-mirror.com/google/siglip-so400m-patch14-384 + +### 2. Llama Language Models +#### 4bit Quantized (unsloth/Meta-Llama-3.1-8B-bnb-4bit) +**International**: https://huggingface.co/unsloth/Meta-Llama-3.1-8B-bnb-4bit +**China Mirror**: https://hf-mirror.com/unsloth/Meta-Llama-3.1-8B-bnb-4bit + +#### Original (meta-llama/Meta-Llama-3.1-8B) +**International**: https://huggingface.co/meta-llama/Meta-Llama-3.1-8B (Access approval required) +**China Mirror**: https://hf-mirror.com/meta-llama/Meta-Llama-3.1-8B + +### 3. Image Adapter (Joy_caption/image_adapter.pt) +**International**: https://huggingface.co/spaces/fancyfeast/joy-caption-pre-alpha/tree/main/wpkklhc6 +**China Mirror**: https://www.modelscope.cn/models/fireicewolf/joy-caption-pre-alpha/files + diff --git a/docs/nodes.md b/docs/nodes.md new file mode 100644 index 0000000..2797c94 --- /dev/null +++ b/docs/nodes.md @@ -0,0 +1,12 @@ +- nodes show in console: +
+ console + +
+ +- nodes show in right mouse menu: + +
+ + right mouse menu +
\ No newline at end of file diff --git a/docs/why.md b/docs/why.md new file mode 100644 index 0000000..1638e81 --- /dev/null +++ b/docs/why.md @@ -0,0 +1,3 @@ + +- use joy to caption image for aigc. +- use joy to caption image files in diretory to do sth. (eg. lora training) \ No newline at end of file diff --git a/joy.py b/joy.py new file mode 100644 index 0000000..44d60a5 --- /dev/null +++ b/joy.py @@ -0,0 +1,382 @@ +import os +from pathlib import Path +from typing import Optional, Tuple, List + +import torch +import torch.nn as nn +from PIL import Image +from transformers import ( + AutoModel, + AutoProcessor, + AutoTokenizer, + PreTrainedTokenizer, + PreTrainedTokenizerFast, + BitsAndBytesConfig, + AutoModelForCausalLM + +) +import folder_paths +from model_management import get_torch_device + +from .lib.ximg import tensor2pil, pil2tensor +from .lib.xmodel import download_hg_model +from .conf import CURRENT_CATEGORY, CURRENT_FUNCTION + +DEVICE = get_torch_device() + +class ImageAdapter(nn.Module): + """Adapter to transform image features to match text model hidden size.""" + def __init__(self, input_features: int, output_features: int): + super().__init__() + self.linear1 = nn.Linear(input_features, output_features) + self.activation = nn.GELU() + self.linear2 = nn.Linear(output_features, output_features) + + def forward(self, vision_outputs: torch.Tensor) -> torch.Tensor: + x = self.linear1(vision_outputs) + x = self.activation(x) + x = self.linear2(x) + return x + +class JoyPipeline: + """Pipeline for handling image captioning models.""" + def __init__(self): + self.clip_model: Optional[nn.Module] = None + self.clip_processor: Optional[AutoProcessor] = None + self.tokenizer: Optional[PreTrainedTokenizer] = None + self.text_model: Optional[nn.Module] = None + self.image_adapter: Optional[nn.Module] = None + self.parent: Optional[object] = None + + def clear_cache(self) -> None: + """Clear all model components from memory.""" + for attr in ['clip_model', 'clip_processor', 'tokenizer', 'text_model', 'image_adapter']: + setattr(self, attr, None) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + +class JoyCaptionBase: + """Base class for Joy captioning functionality.""" + def __init__(self): + self.model: Optional[str] = None + self.pipeline = JoyPipeline() + self.pipeline.parent = self + + def load_checkpoint(self, model_id: str) -> None: + """Load all required models and components.""" + + print(f"Loading model: {model_id}") + if self.pipeline.clip_model is not None and self.model == model_id: + print("Model already loaded, skipping") + return + + print("Loading CLIP model...") + self.pipeline.clear_cache() + self.model = model_id + + # Load CLIP model + clip_model_id = "google/siglip-so400m-patch14-384" + clip_path = download_hg_model(clip_model_id, "clip") + + self.pipeline.clip_processor = AutoProcessor.from_pretrained(clip_path) + clip_model = AutoModel.from_pretrained( + clip_path, + trust_remote_code=True + ).vision_model + + clip_model.eval() + clip_model.requires_grad_(False) + clip_model.to(DEVICE) + self.pipeline.clip_model = clip_model + + # Load LLM + model_path = download_hg_model(model_id, "LLM") + self.pipeline.tokenizer = AutoTokenizer.from_pretrained( + model_path, + use_fast=False + ) + + # todo: fix Unused kwargs: ['_load_in_4bit', '_load_in_8bit', 'quant_method'] + bnb_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_use_double_quant=True, + bnb_4bit_compute_dtype=torch.float16 + ) + self.pipeline.text_model = AutoModelForCausalLM.from_pretrained( + model_path, + quantization_config=bnb_config, + device_map="auto", + trust_remote_code=True, + torch_dtype=torch.float16, + max_memory={0: "10GiB"} + ) + + self.pipeline.text_model.eval() + + # Load image adapter + adapter_path = os.path.join( + folder_paths.models_dir, + "Joy_caption", + "image_adapter.pt" + ) + if not os.path.exists(adapter_path): + raise FileNotFoundError(f"Image adapter not found at {adapter_path}") + + image_adapter = ImageAdapter( + clip_model.config.hidden_size, + self.pipeline.text_model.config.hidden_size + ) + image_adapter.load_state_dict(torch.load(adapter_path, map_location="cpu")) + image_adapter.eval() + image_adapter.to(DEVICE) + self.pipeline.image_adapter = image_adapter + + def generate_caption( + self, + image: torch.Tensor, + prompt: str, + max_new_tokens: int = 1024, + temperature: float = 0.7 + ) -> str: + """Generate caption for a single image.""" + if self.pipeline.clip_processor is None: + raise RuntimeError("Pipeline not initialized. Call load_checkpoint first.") + + # Convert and preprocess image + input_image = tensor2pil(image) + p_image = self.pipeline.clip_processor( + images=input_image, + return_tensors='pt' + ).pixel_values.to(DEVICE) + + # Tokenize prompt + prompt_tokens = self.pipeline.tokenizer.encode( + prompt, + return_tensors='pt', + add_special_tokens=False + ).to(DEVICE) + + with torch.autocast(device_type=str(DEVICE), enabled=True): + # Get image features + vision_outputs = self.pipeline.clip_model( + pixel_values=p_image, + output_hidden_states=True + ) + image_features = vision_outputs.hidden_states[-2] + embedded_images = self.pipeline.image_adapter(image_features) + + # Prepare embeddings + prompt_embeds = self.pipeline.text_model.model.embed_tokens(prompt_tokens) + embedded_bos = self.pipeline.text_model.model.embed_tokens( + torch.tensor( + [[self.pipeline.tokenizer.bos_token_id]], + device=DEVICE + ) + ) + + # Construct input embeddings + inputs_embeds = torch.cat([ + embedded_bos.expand(embedded_images.shape[0], -1, -1), + embedded_images.to(dtype=embedded_bos.dtype), + prompt_embeds.expand(embedded_images.shape[0], -1, -1), + ], dim=1) + + # Generate caption + generate_ids = self.pipeline.text_model.generate( + inputs_embeds=inputs_embeds, + max_new_tokens=max_new_tokens, + do_sample=True, + top_k=10, + temperature=temperature + ) + + # Decode and clean caption + caption = self.pipeline.tokenizer.decode( + generate_ids[0], + skip_special_tokens=True + ).strip() + + return caption + +class JoyCaptionLoad(JoyCaptionBase): + """Node to load Joy captioning models.""" + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ([ + "unsloth/Meta-Llama-3.1-8B-bnb-4bit", + "meta-llama/Meta-Llama-3.1-8B" + ],), + } + } + + CATEGORY = CURRENT_CATEGORY + # FUNCTION = CURRENT_FUNCTION + NODE_DESC = "joy model loader" + RETURN_TYPES = ("JoyPipeline",) + FUNCTION = "load" + OUTPUT_NODE = True + def IS_CHANGED(cls, **kwargs): + return float("NaN") + def load(self, model: str) -> Tuple[JoyPipeline]: + self.load_checkpoint(model) + return (self.pipeline,) + +class JoyCaption(JoyCaptionBase): + """Node to generate captions for single images.""" + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "joy_pipeline": ("JoyPipeline",), + "image": ("IMAGE",), + "prompt": ("STRING", { + "multiline": True, + "default": "A descriptive caption for this image" + }), + "max_new_tokens": ("INT", { + "default": 1024, + "min": 10, + "max": 4096, + "step": 1 + }), + "temperature": ("FLOAT", { + "default": 0.7, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + "cache": ("BOOLEAN", {"default": False}), + } + } + + CATEGORY = CURRENT_CATEGORY + # FUNCTION = CURRENT_FUNCTION + + NODE_DESC = "joy image caption" + OUTPUT_NODE = True + RETURN_TYPES = ("STRING",) + FUNCTION = "generate" + def generate( + self, + joy_pipeline: JoyPipeline, + image: torch.Tensor, + prompt: str, + max_new_tokens: int, + temperature: float, + cache: bool + ) -> Tuple[str]: + try: + self.pipeline = joy_pipeline + caption = self.generate_caption(image, prompt, max_new_tokens, temperature) + + if not cache: + self.pipeline.clear_cache() + + # ensure caption to be in oneline + caption = ' '.join(caption.split()) + return (caption,) + except Exception as e: + raise RuntimeError(f"Caption generation failed: {str(e)}") + def IS_CHANGED(cls, **kwargs): + return float("NaN") + +class JoyCaptionFromDir(JoyCaptionBase): + """Node to generate captions for all images in a directory.""" + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "joy_pipeline": ("JoyPipeline",), + "image_dir": ("STRING", { + "default": "", + "multiline": False + }), + "prompt": ("STRING", { + "multiline": True, + "default": "A descriptive caption for these images" + }), + "max_new_tokens": ("INT", { + "default": 1024, + "min": 10, + "max": 4096, + "step": 1 + }), + "temperature": ("FLOAT", { + "default": 0.7, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + "cache": ("BOOLEAN", {"default": False}), + } + } + + CATEGORY = CURRENT_CATEGORY + # FUNCTION = CURRENT_FUNCTION + NODE_DESC = "joy diretory caption" + OUTPUT_NODE = True + + RETURN_TYPES = ("STRING", "STRING") + RETURN_NAMES = ("image_files", "captions") + FUNCTION = "generate_for_dir" + + def IS_CHANGED(cls, **kwargs): + return float("NaN") + def generate_for_dir( + self, + joy_pipeline: JoyPipeline, + image_dir: str, + prompt: str, + max_new_tokens: int, + temperature: float, + cache: bool + ) -> Tuple[str, str]: + + if not os.path.isdir(image_dir): + raise ValueError(f"Directory not found: {image_dir}") + + # Get all image files from directory + image_extensions = ('.png', '.jpg', '.jpeg', '.webp') + image_files = [f for f in os.listdir(image_dir) + if f.lower().endswith(image_extensions)] + + if not image_files: + print(f"No images found in directory: {image_dir}") + return ('','',) + + + self.pipeline = joy_pipeline + files = [] + captions = [] + + for img_file in image_files: + img_path = os.path.join(image_dir, img_file) + try: + image = Image.open(img_path) + tensor_image = pil2tensor(image) + caption = self.generate_caption( + tensor_image, + prompt, + max_new_tokens, + temperature + ) + # ensure caption to be in oneline + caption = ' '.join(caption.split()) + + files.append(img_path) + captions.append(caption) + print(f"Processed: {img_file}") + except Exception as e: + # captions.append(f"{img_file}: Error - {str(e)}") + print(f"{img_file}: Error - {str(e)}") + continue + if not cache: + self.pipeline.clear_cache() + + return ( + "\n".join(files), + "\n".join(captions) + ) diff --git a/lib/__init__.py b/lib/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/lib/xfile.py b/lib/xfile.py new file mode 100644 index 0000000..d4a9f16 --- /dev/null +++ b/lib/xfile.py @@ -0,0 +1,32 @@ +# import folder_paths +# import os +# import base64 +# import numpy as np +# from PIL import Image,ImageOps, ImageFilter + +# import io + +# comfy_path = os.path.dirname(folder_paths.__file__) +# custom_nodes_path = os.path.join(comfy_path, "custom_nodes") +# # current_folder = os.path.dirname(os.path.abspath(__file__)) + +# # 节点路径 +# def node_path(node_name): +# return os.path.join(custom_nodes_path,node_name) + +# # 创建文件夹 +# def mkdir(path): +# folder = os.path.exists(path) +# if not folder: #判断是否存在文件夹如果不存在则创建为文件夹 +# os.makedirs(path) #makedirs 创建文件时如果路径不存在会创建这个路径 + +# # 获取所有图片文件路径 +# def get_all_image_paths(directory): +# image_paths = [] +# for root, dirs, files in os.walk(directory): +# for file in files: +# if file.lower().endswith(('.png', '.jpg', '.jpeg')): +# image_paths.append(os.path.join(root, file)) +# return image_paths + + diff --git a/lib/ximg.py b/lib/ximg.py new file mode 100644 index 0000000..1a22543 --- /dev/null +++ b/lib/ximg.py @@ -0,0 +1,167 @@ +# -*- encoding: utf-8 -*- +''' +@File :ximg.py +@Description :图片操作类封装 +@Time :2024/04/30 09:46:01 +@Author :ChenXingHua +@Version :1.0 +''' + +# Standard library imports +import base64 +import io +import os +from datetime import datetime, timedelta +from io import BytesIO + +# Third-party imports +import cv2 as cv +import numpy as np +import torch +from PIL import Image, ImageFile, ImageOps, ImageSequence, UnidentifiedImageError +import requests + + +def tensor2pil(t_image: torch.Tensor) -> Image: + return Image.fromarray(np.clip(255.0 * t_image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) + +def pil2tensor(image:Image) -> torch.Tensor: + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + +def tensor2cv2(image:torch.Tensor) -> np.array: + if image.dim() == 4: + image = image.squeeze() + npimage = image.numpy() + cv2image = np.uint8(npimage * 255 / npimage.max()) + return cv.cvtColor(cv2image, cv.COLOR_RGB2BGR) + +def cv22pil(cv2_img:np.ndarray) -> Image: + cv2_img = cv.cvtColor(cv2_img, cv.COLOR_BGR2RGB) + return Image.fromarray(cv2_img) + +# pil转io +def pil2iobyte(pil_image,format='PNG'): + byte_arr = io.BytesIO() + pil_image.save(byte_arr, format=format) + byte_arr = byte_arr.getvalue() + return byte_arr + +# pil转64 +def pilTobase64(pil_image,format='PNG'): + byte_arr = pil2iobyte(pil_image,format) + image_base64 = base64.b64encode(byte_arr).decode('utf-8') + return image_base64 + +def ioBytes2tensor(bytes): + image = Image.open(bytes) + return pil2tensor(image) + +def getImageSize(image): + if image.shape[0] > 0: + image = torch.unsqueeze(image[0], 0) + _image = tensor2pil(image) + + return (_image.width, _image.height) + +# 转成mask +def imageToMask(img): + i = img + if i.mode == 'I': + i = i.point(lambda i: i * (1 / 255)) + image = i.convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + + if 'A' in i.getbands(): + mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) + else: + mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") + return tensor2pil(mask) + +# ret_masks.append(image2mask(_mask)) +def image2mask(image:Image) -> torch.Tensor: + _image = image.convert('RGBA') + alpha = _image.split() [0] + bg = Image.new("L", _image.size) + _image = Image.merge('RGBA', (bg, bg, bg, alpha)) + ret_mask = torch.tensor([pil2tensor(_image)[0, :, :, 3].tolist()]) + return ret_mask + +# 图像回帖 +def croppImg(original_image,cropped_avatar,left_x,top_y): + # 获取原始图像的大小 + original_width, original_height = original_image.size + return croppImageBySize(cropped_avatar,left_x,top_y,original_width,original_height) + +def croppImageBySize(cropped_avatar,left_x,top_y,original_w,original_h): + # 获取原始图像的大小 + original_width, original_height = original_w,original_h + # 获取头像的大小 + avatar_width, avatar_height = cropped_avatar.size + # 创建一个与原始图像相同大小的透明图像 + extended_image = Image.new("RGBA", (original_width, original_height), (0, 0, 0, 0)) + # 将裁剪后的头像粘贴到新图像 + extended_image.paste(cropped_avatar, (left_x, top_y), cropped_avatar) + + return extended_image + + +# 将图片转换为Base64编码 +def image_to_base64(image_path): + with open(image_path, 'rb') as image_file: + return base64.b64encode(image_file.read()).decode('utf-8') + +# 获取网络图片 +def img_from_url(url): + # 发送HTTP请求获取图片 + response = requests.get(url) + response.raise_for_status() # 如果请求失败，这会抛出异常 + # 将响应内容作为BytesIO对象打开，以便PIL可以读取它 + image = Image.open(BytesIO(response.content)) + return image + +def open_image(path): + prev_value = None + + try: + img = Image.open(path) + except (UnidentifiedImageError, ValueError): #PIL issues #4472 and #2445 + prev_value = ImageFile.LOAD_TRUNCATED_IMAGES + ImageFile.LOAD_TRUNCATED_IMAGES = True + img = Image.open(path) + finally: + if prev_value is not None: + ImageFile.LOAD_TRUNCATED_IMAGES = prev_value + return img + +# 批量读取 +def batch_image(directory): + if not os.path.isdir(directory): + raise FileNotFoundError(f"Directory '{directory}' cannot be found.") + dir_files = os.listdir(directory) + if len(dir_files) == 0: + raise FileNotFoundError(f"No files in directory '{directory}'.") + + valid_extensions = ['.jpg', '.jpeg', '.png', '.webp'] + dir_files = [f for f in dir_files if any(f.lower().endswith(ext) for ext in valid_extensions)] + + dir_files = sorted(dir_files) + dir_files = [os.path.join(directory, x) for x in dir_files] + return dir_files + +def calculate_seconds_difference(start_time, end_time): + """ + 计算两个时间点之间的秒数差异 + + :param start_time: 开始时间（可以是时间戳或datetime对象） + :param end_time: 结束时间（可以是时间戳或datetime对象） + :return: 秒数差异（浮点数） + """ + # 如果输入是datetime对象，转换为时间戳 + if isinstance(start_time, datetime): + start_time = start_time.timestamp() + if isinstance(end_time, datetime): + end_time = end_time.timestamp() + + return end_time - start_time \ No newline at end of file diff --git a/lib/xmodel.py b/lib/xmodel.py new file mode 100644 index 0000000..afa06c2 --- /dev/null +++ b/lib/xmodel.py @@ -0,0 +1,34 @@ +# Standard library imports +import os +import json + +# Third-party imports +import folder_paths +import torch +from transformers import AutoProcessor + +# def get_torch_device(): +# """ +# 返回PyTorch模型应该运行的设备（CPU或GPU） +# 如果系统支持CUDA并且至少有一个GPU可用，则返回GPU设备；否则返回CPU设备。 +# """ +# if torch.cuda.is_available(): +# # 选择第一个可用的GPU +# device = torch.device("cuda:0") +# print(f"There are {torch.cuda.device_count()} GPU(s) available.") +# print(f"We will use the GPU: {device}") +# else: +# # 如果没有GPU可用，则使用CPU +# device = torch.device("cpu") +# print("No GPU available, using the CPU instead.") +# return device + +# 下载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 \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..f3bc1ec --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,15 @@ +[project] +name = "ymc_node_joy" +description = "comfyui custom nodes to caption image with joy" +version = "1.0.0" + + [project.license] + file = "LICENSE" + + [project.urls] + Repository = "https://github.com/ymc-github/ymc_node_joy" + +[tool.comfy] +PublisherId = "ymc-github" +DisplayName = "ymc_node_joy" +Icon = "" diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..bb06e79 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,23 @@ +# Core dependencies +# torch==2.5.0+cu124 +# transformers==4.47.1 +# numpy==1.26.3 +# Pillow==10.1.0 +# torch +# transformers +# numpy +# Pillow + +# ComfyUI specific +# + +# Custom node utilities +# yors_comfyui_node_setup==0.10.1 +# yors_pano_ansi_color==1.2.1 +yors_comfyui_node_setup +yors_pano_ansi_color + +# Optional dependencies (for advanced features) +# bitsandbytes==0.45.0 +# bitsandbytes +