Rename downloader nodes and add nodes for displaying previews, metadata of model and workflows extracted from images/videos

2026-04-08 16:46:35 +02:00
parent 39380f533d
commit dc43040562
4 changed files with 263 additions and 17 deletions
@@ -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,
        ]


@@ -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())
@@ -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.
@@ -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])