diff --git a/.comfyignore b/.comfyignore index 2bfa6a4..762cc57 100644 --- a/.comfyignore +++ b/.comfyignore @@ -1 +1,4 @@ tests/ +nodes/downloader_nodes.py +nodes/downloader_utils.py + diff --git a/__init__.py b/__init__.py index 57986da..5366272 100644 --- a/__init__.py +++ b/__init__.py @@ -49,10 +49,18 @@ from .nodes.lora_resize import ( LoRAResizeFrobenius, LoRAResizeCumulative, LoRAMergeToModel ) -from .nodes.downloader_nodes import ( - CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader, - VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader -) +try: + from .nodes.downloader_nodes import ( + CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader, + VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader + ) + DOWNLOADER_NODES: list[type[io.ComfyNode]] = [ + CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader, + VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader, + ] +except ImportError: + DOWNLOADER_NODES = [] + from .nodes.model_info_nodes import ( CheckpointInfoLoader, LoRAInfoLoader, EmbeddingInfoLoader, VAEInfoLoader, ControlNetInfoLoader, DiffusionModelInfoLoader @@ -125,9 +133,8 @@ class ModelUtilsExtension(ComfyExtension): DiffusionModelDtypeConversion, MiniMaxH3FoldAdaLN, MiniMaxH3UnfoldAdaLN, - # Downloaders - CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader, - VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader, + # Downloaders (available in manual git/source installations) + *DOWNLOADER_NODES, # Info Loaders CheckpointInfoLoader, LoRAInfoLoader, EmbeddingInfoLoader, VAEInfoLoader, ControlNetInfoLoader, DiffusionModelInfoLoader, diff --git a/nodes/downloader_utils.py b/nodes/downloader_utils.py index ba6500a..7c5b1a3 100644 --- a/nodes/downloader_utils.py +++ b/nodes/downloader_utils.py @@ -32,11 +32,15 @@ except ImportError: HAS_AV = False print("Warning: av (PyAV) not found. MP4 metadata extraction fallback disabled.") -# Add the reference module to the Python path # Assuming downloader_utils.py is inside nodes/, the root is one level up root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.append(os.path.join(root_dir, 'unifiedefficientloader-reference')) from unifiedefficientloader import UnifiedSafetensorsLoader +from .model_info_nodes import ( + get_potential_preview_files, + load_image_tensor, + get_model_workflows, + get_model_metadata_file, +) # --- Database Path Logic --- try: @@ -381,14 +385,6 @@ class Downloader: os.rename(temp_path, target_path) return True -def get_potential_preview_files(base_path): - exts = ["png", "jpg", "jpeg", "webp", "mp4"] - files = [] - for ext in exts: - files.append(f"{base_path}.preview.{ext}") - files.append(f"{base_path}.{ext}") - return files - def preview_exists(base_path): previews = get_potential_preview_files(base_path) for p in previews: @@ -569,59 +565,6 @@ def process_file(filepath, cache, api, nsfw_threshold, max_examples): for k in stats: stats[k] += res.get(k, 0) return stats -import numpy as np -from PIL import Image, ImageOps -import torch - -def load_image_tensor(path): - if not os.path.exists(path): - return None - try: - img = Image.open(path) - img = ImageOps.exif_transpose(img) - image = img.convert("RGB") - image = np.array(image).astype(np.float32) / 255.0 - image = torch.from_numpy(image)[None,] - return image - except Exception as e: - print(f"Error loading image {path}: {e}") - return None - -def get_model_workflows(base_path): - workflows = [] - # Main workflow - main_json = f"{base_path}.json" - if os.path.exists(main_json): - workflows.append(main_json) - - # Example workflows - # They are named like base_path.example.N.json - # or base_path.preview.json (if extracted from preview) - preview_json = f"{base_path}.preview.json" - if os.path.exists(preview_json) and preview_json not in workflows: - workflows.append(preview_json) - - example_jsons = glob.glob(f"{base_path}.example.*.json") - # Sort them by index - example_jsons.sort() - for ej in example_jsons: - if ej not in workflows: - workflows.append(ej) - - return workflows - -def get_model_metadata_file(base_path): - # Prefer civitai.info if available - civitai_info = f"{base_path}.civitai.info" - if os.path.exists(civitai_info): - return civitai_info - - # Fallback to metadata.json (safetensors metadata) - metadata_json = f"{base_path}.metadata.json" - if os.path.exists(metadata_json): - return metadata_json - - return None def scan_and_process(scan_dirs, recursive=True, db_path=None, nsfw_level="All", max_examples=0, api_key=None, threads=4): """ diff --git a/nodes/model_info_nodes.py b/nodes/model_info_nodes.py index a22b107..2e91a91 100644 --- a/nodes/model_info_nodes.py +++ b/nodes/model_info_nodes.py @@ -1,7 +1,64 @@ +import glob import os +import numpy as np +from PIL import Image, ImageOps +import torch import folder_paths from comfy_api.latest import io -from .downloader_utils import get_potential_preview_files, load_image_tensor, get_model_workflows, get_model_metadata_file + + +def get_potential_preview_files(base_path: str) -> list[str]: + exts = ["png", "jpg", "jpeg", "webp", "mp4"] + files = [] + for ext in exts: + files.append(f"{base_path}.preview.{ext}") + files.append(f"{base_path}.{ext}") + return files + + +def load_image_tensor(path: str): + if not os.path.exists(path): + return None + try: + img = Image.open(path) + img = ImageOps.exif_transpose(img) + image = img.convert("RGB") + image_np = np.array(image).astype(np.float32) / 255.0 + return torch.from_numpy(image_np)[None,] + except Exception as e: + print(f"Error loading image {path}: {e}") + return None + + +def get_model_workflows(base_path: str) -> list[str]: + workflows = [] + main_json = f"{base_path}.json" + if os.path.exists(main_json): + workflows.append(main_json) + + preview_json = f"{base_path}.preview.json" + if os.path.exists(preview_json) and preview_json not in workflows: + workflows.append(preview_json) + + example_jsons = glob.glob(f"{base_path}.example.*.json") + example_jsons.sort() + for ej in example_jsons: + if ej not in workflows: + workflows.append(ej) + + return workflows + + +def get_model_metadata_file(base_path: str) -> str | None: + civitai_info = f"{base_path}.civitai.info" + if os.path.exists(civitai_info): + return civitai_info + + metadata_json = f"{base_path}.metadata.json" + if os.path.exists(metadata_json): + return metadata_json + + return None class BaseModelInfoLoader(io.ComfyNode): CATEGORY = "ModelUtils/Info" diff --git a/tests/test_conditional_registration.py b/tests/test_conditional_registration.py new file mode 100644 index 0000000..9a61944 --- /dev/null +++ b/tests/test_conditional_registration.py @@ -0,0 +1,49 @@ +import asyncio +import importlib.util +import sys +from pathlib import Path +from comfy_api.latest import io + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def _load_extension_module(): + package_name = "ComfyUI_ModelUtils_init_test" + if package_name not in sys.modules: + spec = importlib.util.spec_from_file_location( + package_name, + REPO_ROOT / "__init__.py", + submodule_search_locations=[str(REPO_ROOT)], + ) + module = importlib.util.module_from_spec(spec) + sys.modules[package_name] = module + spec.loader.exec_module(module) + return sys.modules[package_name] + + +def test_startup_node_registration_when_downloaders_absent(monkeypatch): + ext_module = _load_extension_module() + + ext = ext_module.ModelUtilsExtension() + nodes_with_downloader = asyncio.run(ext.get_node_list()) + assert len(nodes_with_downloader) > 0 + + # Simulate registry install by setting DOWNLOADER_NODES to empty + monkeypatch.setattr(ext_module, "DOWNLOADER_NODES", []) + nodes_registry = asyncio.run(ext.get_node_list()) + assert len(nodes_registry) < len(nodes_with_downloader) + assert all(issubclass(n, io.ComfyNode) for n in nodes_registry) + + +def test_model_info_nodes_importable_standalone(): + spec = importlib.util.spec_from_file_location( + "modelutils_model_info_nodes_test", + REPO_ROOT / "nodes" / "model_info_nodes.py", + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + assert hasattr(module, "get_potential_preview_files") + assert hasattr(module, "load_image_tensor") + assert hasattr(module, "get_model_workflows") + assert hasattr(module, "get_model_metadata_file") diff --git a/tests/test_groups.toml b/tests/test_groups.toml index 72f4fbb..72a2d76 100644 --- a/tests/test_groups.toml +++ b/tests/test_groups.toml @@ -1,6 +1,7 @@ [groups.repository_contract] paths = ["__init__.py", "nodes/*.py"] python_tests = [ + "tests/test_conditional_registration.py", "tests/test_schema_tooltips.py", "tests/test_uel_repository_contract.py", ]