diff --git a/__init__.py b/__init__.py index f78e99f..8896830 100644 --- a/__init__.py +++ b/__init__.py @@ -30,8 +30,12 @@ from .nodes.lora_resize import ( LoRAMergeToModel ) from .nodes.downloader_nodes import ( - CheckpointDownloader, DiffusionModelDownloader, LoRADownloader, EmbeddingDownloader, - VAEDownloader, ControlNetDownloader, ManualPathDownloader + CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader, + VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader +) +from .nodes.model_info_nodes import ( + CheckpointInfoLoader, LoRAInfoLoader, EmbeddingInfoLoader, + VAEInfoLoader, ControlNetInfoLoader, DiffusionModelInfoLoader ) @@ -65,8 +69,11 @@ class ModelUtilsExtension(ComfyExtension): # LoRA Merge To Model LoRAMergeToModel, # Downloaders - CheckpointDownloader, DiffusionModelDownloader, LoRADownloader, EmbeddingDownloader, - VAEDownloader, ControlNetDownloader, ManualPathDownloader, + CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader, + VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader, + # Info Loaders + CheckpointInfoLoader, LoRAInfoLoader, EmbeddingInfoLoader, + VAEInfoLoader, ControlNetInfoLoader, DiffusionModelInfoLoader, ] diff --git a/nodes/downloader_nodes.py b/nodes/downloader_nodes.py index 439fc3f..f6f15ef 100644 --- a/nodes/downloader_nodes.py +++ b/nodes/downloader_nodes.py @@ -26,18 +26,18 @@ def get_subdirectories(folder_name: str) -> list[str]: return sorted(list(subdirs)) -class BaseDownloaderNode(io.ComfyNode): +class BaseInfoMetaDownloaderNode(io.ComfyNode): CATEGORY = "ModelUtils/Downloader" OUTPUT_NODE = True @classmethod def get_common_inputs(cls): return { - "recursive": io.Boolean.Input("recursive", default=True), - "nsfw_level": io.Combo.Input("nsfw_level", options=["None", "Soft", "Mature", "X", "XXX", "All"], default="All"), - "max_examples": io.Int.Input("max_examples", default=0, min=0, max=100), - "threads": io.Int.Input("threads", default=4, min=1, max=32), - "api_key": io.String.Input("api_key", default="") + "recursive": io.Boolean.Input("recursive", default=True, tooltip="Whether to scan subdirectories recursively."), + "nsfw_level": io.Combo.Input("nsfw_level", options=["None", "Soft", "Mature", "X", "XXX", "All"], default="All", tooltip="Filter models based on NSFW level in metadata."), + "max_examples": io.Int.Input("max_examples", default=0, min=0, max=100, tooltip="Maximum number of examples to process."), + "threads": io.Int.Input("threads", default=4, min=1, max=32, tooltip="Number of threads to use for processing."), + "api_key": io.String.Input("api_key", default="", tooltip="API key for accessing the model repository.") } @classmethod @@ -49,6 +49,7 @@ class BaseDownloaderNode(io.ComfyNode): inputs=list(cls.get_common_inputs().values()), outputs=[io.String.Output(display_name="status")], is_output_node=True, + description=("Scans specified directories for models, extracts metadata, generates previews, and optionally compresses files.") ) @classmethod @@ -84,7 +85,7 @@ class BaseDownloaderNode(io.ComfyNode): def get_scan_dirs(cls, **kwargs) -> list[str]: return [] -class CheckpointDownloader(BaseDownloaderNode): +class CheckpointInfoMetaDownloader(BaseInfoMetaDownloaderNode): @classmethod def define_schema(cls): inputs = list(cls.get_common_inputs().values()) @@ -121,7 +122,7 @@ class CheckpointDownloader(BaseDownloaderNode): target_paths.append(target_path) return target_paths -class DiffusionModelDownloader(BaseDownloaderNode): +class DiffusionModelInfoMetaDownloader(BaseInfoMetaDownloaderNode): @classmethod def define_schema(cls): inputs = list(cls.get_common_inputs().values()) @@ -158,7 +159,7 @@ class DiffusionModelDownloader(BaseDownloaderNode): target_paths.append(target_path) return target_paths -class LoRADownloader(BaseDownloaderNode): +class LoRAInfoMetaDownloader(BaseInfoMetaDownloaderNode): @classmethod def define_schema(cls): inputs = list(cls.get_common_inputs().values()) @@ -195,7 +196,7 @@ class LoRADownloader(BaseDownloaderNode): target_paths.append(target_path) return target_paths -class EmbeddingDownloader(BaseDownloaderNode): +class EmbeddingInfoMetaDownloader(BaseInfoMetaDownloaderNode): @classmethod def define_schema(cls): inputs = list(cls.get_common_inputs().values()) @@ -232,7 +233,7 @@ class EmbeddingDownloader(BaseDownloaderNode): target_paths.append(target_path) return target_paths -class VAEDownloader(BaseDownloaderNode): +class VAEInfoMetaDownloader(BaseInfoMetaDownloaderNode): @classmethod def define_schema(cls): inputs = list(cls.get_common_inputs().values()) @@ -269,7 +270,7 @@ class VAEDownloader(BaseDownloaderNode): target_paths.append(target_path) return target_paths -class ControlNetDownloader(BaseDownloaderNode): +class ControlNetInfoMetaDownloader(BaseInfoMetaDownloaderNode): @classmethod def define_schema(cls): inputs = list(cls.get_common_inputs().values()) @@ -306,7 +307,7 @@ class ControlNetDownloader(BaseDownloaderNode): target_paths.append(target_path) return target_paths -class ManualPathDownloader(BaseDownloaderNode): +class ManualPathInfoMetaDownloader(BaseInfoMetaDownloaderNode): @classmethod def define_schema(cls): inputs = list(cls.get_common_inputs().values()) diff --git a/nodes/downloader_utils.py b/nodes/downloader_utils.py index 8917afa..a27b072 100644 --- a/nodes/downloader_utils.py +++ b/nodes/downloader_utils.py @@ -576,6 +576,60 @@ 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): """ Scans directories for models and fetches their metadata/previews. diff --git a/nodes/model_info_nodes.py b/nodes/model_info_nodes.py new file mode 100644 index 0000000..a9a4335 --- /dev/null +++ b/nodes/model_info_nodes.py @@ -0,0 +1,184 @@ +import os +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 + +class BaseModelInfoLoader(io.ComfyNode): + CATEGORY = "ModelUtils/Info" + + @classmethod + def get_info(cls, category, name, workflow_index): + full_path = folder_paths.get_full_path(category, name) + if not full_path or not os.path.exists(full_path): + return (None, "{}", "{}") + + base_path = os.path.splitext(full_path)[0] + + # 1. Image Preview + image_tensor = None + previews = get_potential_preview_files(base_path) + for p in previews: + if os.path.exists(p) and not p.lower().endswith('.mp4'): + image_tensor = load_image_tensor(p) + if image_tensor is not None: + break + + # 2. Workflow + workflow_json = "{}" + workflows = get_model_workflows(base_path) + if workflows and 0 <= workflow_index < len(workflows): + try: + with open(workflows[workflow_index], 'r', encoding='utf-8') as f: + workflow_json = f.read() + except Exception as e: + print(f"Error reading workflow {workflows[workflow_index]}: {e}") + + # 3. Metadata + metadata_json = "{}" + metadata_file = get_model_metadata_file(base_path) + if metadata_file: + try: + with open(metadata_file, 'r', encoding='utf-8') as f: + metadata_json = f.read() + except Exception as e: + print(f"Error reading metadata {metadata_file}: {e}") + + return (image_tensor, workflow_json, metadata_json) + +class CheckpointInfoLoader(BaseModelInfoLoader): + @classmethod + def define_schema(cls): + return io.Schema( + node_id=cls.__name__, + display_name="Checkpoint Info Loader", + category=cls.CATEGORY, + inputs=[ + io.Combo.Input("checkpoint", options=folder_paths.get_filename_list("checkpoints")), + io.Int.Input("workflow_index", default=0, min=0, max=100) + ], + outputs=[ + io.Image.Output(display_name="preview"), + io.String.Output(display_name="workflow_json"), + io.String.Output(display_name="metadata_json") + ] + ) + + @classmethod + def execute(cls, checkpoint, workflow_index): + res = cls.get_info("checkpoints", checkpoint, workflow_index) + return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + +class LoRAInfoLoader(BaseModelInfoLoader): + @classmethod + def define_schema(cls): + return io.Schema( + node_id=cls.__name__, + display_name="LoRA Info Loader", + category=cls.CATEGORY, + inputs=[ + io.Combo.Input("lora", options=folder_paths.get_filename_list("loras")), + io.Int.Input("workflow_index", default=0, min=0, max=100) + ], + outputs=[ + io.Image.Output(display_name="preview"), + io.String.Output(display_name="workflow_json"), + io.String.Output(display_name="metadata_json") + ] + ) + + @classmethod + def execute(cls, lora, workflow_index): + res = cls.get_info("loras", lora, workflow_index) + return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + +class EmbeddingInfoLoader(BaseModelInfoLoader): + @classmethod + def define_schema(cls): + return io.Schema( + node_id=cls.__name__, + display_name="Embedding Info Loader", + category=cls.CATEGORY, + inputs=[ + io.Combo.Input("embedding", options=folder_paths.get_filename_list("embeddings")), + io.Int.Input("workflow_index", default=0, min=0, max=100) + ], + outputs=[ + io.Image.Output(display_name="preview"), + io.String.Output(display_name="workflow_json"), + io.String.Output(display_name="metadata_json") + ] + ) + + @classmethod + def execute(cls, embedding, workflow_index): + res = cls.get_info("embeddings", embedding, workflow_index) + return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + +class VAEInfoLoader(BaseModelInfoLoader): + @classmethod + def define_schema(cls): + return io.Schema( + node_id=cls.__name__, + display_name="VAE Info Loader", + category=cls.CATEGORY, + inputs=[ + io.Combo.Input("vae", options=folder_paths.get_filename_list("vae")), + io.Int.Input("workflow_index", default=0, min=0, max=100) + ], + outputs=[ + io.Image.Output(display_name="preview"), + io.String.Output(display_name="workflow_json"), + io.String.Output(display_name="metadata_json") + ] + ) + + @classmethod + def execute(cls, vae, workflow_index): + res = cls.get_info("vae", vae, workflow_index) + return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + +class ControlNetInfoLoader(BaseModelInfoLoader): + @classmethod + def define_schema(cls): + return io.Schema( + node_id=cls.__name__, + display_name="ControlNet Info Loader", + category=cls.CATEGORY, + inputs=[ + io.Combo.Input("controlnet", options=folder_paths.get_filename_list("controlnet")), + io.Int.Input("workflow_index", default=0, min=0, max=100) + ], + outputs=[ + io.Image.Output(display_name="preview"), + io.String.Output(display_name="workflow_json"), + io.String.Output(display_name="metadata_json") + ] + ) + + @classmethod + def execute(cls, controlnet, workflow_index): + res = cls.get_info("controlnet", controlnet, workflow_index) + return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2]) + +class DiffusionModelInfoLoader(BaseModelInfoLoader): + @classmethod + def define_schema(cls): + return io.Schema( + node_id=cls.__name__, + display_name="Diffusion Model Info Loader", + category=cls.CATEGORY, + inputs=[ + io.Combo.Input("diffusion_model", options=folder_paths.get_filename_list("diffusion_models")), + io.Int.Input("workflow_index", default=0, min=0, max=100) + ], + outputs=[ + io.Image.Output(display_name="preview"), + io.String.Output(display_name="workflow_json"), + io.String.Output(display_name="metadata_json") + ] + ) + + @classmethod + def execute(cls, diffusion_model, workflow_index): + res = cls.get_info("diffusion_models", diffusion_model, workflow_index) + return io.NodeOutput(preview=res[0], workflow_json=res[1], metadata_json=res[2])