From 9adee6a1da6153fa462759e33dfb167caf99c424 Mon Sep 17 00:00:00 2001 From: yemiancheng Date: Mon, 16 Jun 2025 08:25:57 +0800 Subject: [PATCH] build(core): put diretory caption --- __init__.py | 95 +++++++++------------------------------------------ joy.py | 15 +++++--- lib/xmodel.py | 18 ++++++++++ 3 files changed, 46 insertions(+), 82 deletions(-) diff --git a/__init__.py b/__init__.py index b22fbe1..a40ef8d 100644 --- a/__init__.py +++ b/__init__.py @@ -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,"=")) \ No newline at end of file diff --git a/joy.py b/joy.py index 44d60a5..2959594 100644 --- a/joy.py +++ b/joy.py @@ -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() diff --git a/lib/xmodel.py b/lib/xmodel.py index afa06c2..5450678 100644 --- a/lib/xmodel.py +++ b/lib/xmodel.py @@ -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=''): # 下载本地