build(core): put diretory caption

This commit is contained in:
yemiancheng
2025-06-16 08:25:57 +08:00
parent a45b8424b3
commit 9adee6a1da
3 changed files with 46 additions and 82 deletions
+17 -78
View File
@@ -1,80 +1,11 @@
# 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)
from yors_pano_path_util import path_resolve,path_dirname,path_parse,path_comfy_get
__all__,NODE_CLASS_MAPPINGS,NODE_DISPLAY_NAME_MAPPINGS,NODE_MENU_NAMES = entry(__name__,__file__,False)
@@ -82,19 +13,26 @@ 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 root path with __file
root = path_resolve("../",__file__)
# build(core): get name from root path
name, stem, suffix, parent=path_parse(root)
comfy_root_path = path_comfy_get(__file__,2,'./')
# - 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)
# comfy_root_path = path_resolve('../../../',__file__)
# log_msg(f'comfy_root_path: {comfy_root_path}')
# comfy_model_path = path_comfy_get(__file__,2,'models')
# log_msg(f'comfy_model_path: {comfy_model_path}')
# comfy_model_path = path_comfy_get(__file__,'../../models')
# log_msg(f'comfy_model_path: {comfy_model_path}')
# comfy_model_path = path_resolve('models',comfy_root_path)
# log_msg(f'comfy_model_path: {comfy_model_path}')
# name='ymc_node_joy'
# version='1.0.0'
@@ -107,4 +45,5 @@ 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(f'comfy_root_path: {comfy_root_path}')
log_msg(msg_padd("=",60,"="))
+11 -4
View File
@@ -1,6 +1,7 @@
import os
from pathlib import Path
from typing import Optional, Tuple, List
import time
import torch
import torch.nn as nn
@@ -352,9 +353,12 @@ class JoyCaptionFromDir(JoyCaptionBase):
files = []
captions = []
for img_file in image_files:
total_images = len(image_files)
for index, img_file in enumerate(image_files, 1):
img_path = os.path.join(image_dir, img_file)
try:
start_time = time.time()
image = Image.open(img_path)
tensor_image = pil2tensor(image)
caption = self.generate_caption(
@@ -368,10 +372,13 @@ class JoyCaptionFromDir(JoyCaptionBase):
files.append(img_path)
captions.append(caption)
print(f"Processed: {img_file}")
elapsed = time.time() - start_time
print(f"{index}/{total_images}:{elapsed:.2f} s - Processed: {img_file}")
except Exception as e:
# captions.append(f"{img_file}: Error - {str(e)}")
print(f"{img_file}: Error - {str(e)}")
elapsed = time.time() - start_time if 'start_time' in locals() else 0
print(f"{index}/{total_images}:{elapsed:.2f} s - Error: {img_file} - {str(e)}")
continue
if not cache:
self.pipeline.clear_cache()
+18
View File
@@ -23,6 +23,24 @@ from transformers import AutoProcessor
# print("No GPU available, using the CPU instead.")
# return device
def get_comfyui_models_dir() -> str:
"""
Get ComfyUI models directory path by traversing up from current custom nodes directory
Returns:
str: Absolute path to ComfyUI models directory
"""
current_dir = os.path.dirname(os.path.abspath(__file__))
# Traverse up 3 levels from custom_nodes to reach ComfyUI root
comfyui_root = os.path.dirname(os.path.dirname(os.path.dirname(current_dir)))
models_dir = os.path.join(comfyui_root, 'models')
if not os.path.exists(models_dir):
raise FileNotFoundError(f"ComfyUI models directory not found at: {models_dir}")
return models_dir
# 下载hg 模型到本地
def download_hg_model(model_id:str,exDir:str=''):
# 下载本地