Add first nodes, toml and publish

.github/workflows/publish_action.yml

@@ -0,0 +1,20 @@
+name: Publish to Comfy registry
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - master
+    paths:
+      - "pyproject.toml"
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+      - name: Publish Custom Node
+        uses: Comfy-Org/publish-node-action@main
+        with:
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} ## Add your own personal access token to your Github Repository secrets and reference it here.

__init__.py

@@ -0,0 +1,15 @@
+from .nodes import *
+
+NODE_CLASS_MAPPINGS = {
+    "UNetMetaKeys": UNetMetaKeys,
+    "CLIPMetaKeys": CLIPMetaKeys,
+    "LoRAMetaKeys": LoRAMetaKeys,
+    "CheckpointMetaKeys": CheckpointMetaKeys,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "UNetMetaKeys": "Get UNet Metadata and Keys",
+    "CLIPMetaKeys": "Get CLIP Metadata and Keys",
+    "LoRAMetaKeys": "Get LoRA Metadata and Keys",
+    "CheckpointMetaKeys": "Get Checkpoint Metadata and Keys",
+}

nodes/__init__.py

@@ -0,0 +1,8 @@
+from .metakeys import *
+
+__all__ = (
+    "UNetMetaKeys",
+    "CLIPMetaKeys",
+    "LoRAMetaKeys",
+    "CheckpointMetaKeys",
+)

nodes/metakeys.py

@@ -0,0 +1,118 @@
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+import os
+import folder_paths
+from comfy.model_management import get_torch_device
+
+
+def convert_pt_to_safetensors(pt_path, safe_path):
+    try:
+        model = torch.load(pt_path, map_location="cpu")
+        if "state_dict" in model:
+            state_dict = model["state_dict"]
+        elif "model" in model:
+            state_dict = model["model"]
+        else:
+            state_dict = model
+        state_dict = {k.replace("state_dict.", ""): v for k, v in state_dict.items()}
+        metadata = {"format": "pt"}
+        save_file(state_dict, safe_path, metadata)
+        return True, ""
+    except Exception as e:
+        return False, str(e)
+
+def load_metadata_from_safetensors(safetensors_file: str) -> dict:
+    try:
+        with safe_open(safetensors_file, framework="pt", device="cpu") as f:
+            model_weights = {}
+            for key in f.keys():
+                model_weights[key] = f.get_tensor(key)
+            metadata = f.metadata()
+        if metadata is None:
+            metadata = {}
+        return model_weights, metadata
+    except Exception as e:
+        print(f"Error loading safetensors file: {e}")
+        return {}, {}
+
+class BaseMetaKeys:
+    RETURN_TYPES = ("STRING", "STRING")
+    RETURN_NAMES = ("metadata", "keys")
+    FUNCTION = "get_metakeys"
+    CATEGORY = "utils"
+    DESCRIPTION = "Loads a model and returns its metadata and layer names."
+
+    def _get_metakeys(self, model_name, model_type):
+        model_path = folder_paths.get_full_path_or_raise(model_type, model_name)
+        model_weights = {}
+        metadata = {}
+        layer_shapes = ""
+
+        if model_path.endswith(('.pt', '.pth', '.bin', '.ckpt')):
+            temp_safe_path = model_path + ".safetensors"
+            if not os.path.exists(temp_safe_path):
+                conversion_successful, error_message = convert_pt_to_safetensors(model_path, temp_safe_path)
+                if not conversion_successful:
+                    return f"Conversion failed: {error_message}", ""
+            model_path = temp_safe_path
+
+        model_weights, metadata = load_metadata_from_safetensors(model_path)
+        metadata_str = str(metadata)
+        for layer_name, tensor in model_weights.items():
+            shape = str(tensor.shape)
+            shape = shape.replace("torch.Size(", "").replace(")", "")
+            name_shape = f"{layer_name}, {shape}\n"
+            layer_shapes += name_shape
+        return metadata_str, layer_shapes
+
+
+
+class UNetMetaKeys(BaseMetaKeys):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "unet_name": (folder_paths.get_filename_list("diffusion_models"), ),
+            }
+        }
+
+    def get_metakeys(self, unet_name):
+        return self._get_metakeys(unet_name, "diffusion_models")
+
+class CLIPMetaKeys(BaseMetaKeys):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "clip_name": (folder_paths.get_filename_list("text_encoders"), ),
+            }
+        }
+
+    def get_metakeys(self, clip_name):
+        return self._get_metakeys(clip_name, "text_encoders")
+
+class LoRAMetaKeys(BaseMetaKeys):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "lora_name": (folder_paths.get_filename_list("loras"), ),
+            }
+        }
+    def get_metakeys(self, lora_name):
+        return self._get_metakeys(lora_name, "loras")
+
+class CheckpointMetaKeys(BaseMetaKeys):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
+            }
+        }
+
+    def get_metakeys(self, ckpt_name):
+        return self._get_metakeys(ckpt_name, "checkpoints")
+
+__all__ = ("UNetMetaKeys", "CLIPMetaKeys", "LoRAMetaKeys", "CheckpointMetaKeys")

pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "comfyui_modelutils"
+description = "Custom nodes for handling, inspecting, modifying and creating various model files."
+version = "0.1.0"
+license = { file = "LICENSE" }
+
+[project.urls]
+Repository = "https://github.com/silveroxides/ComfyUI-ModelUtils"
+#  Used by Comfy Registry https://comfyregistry.org
+
+[tool.comfy]
+PublisherId = "silveroxides"
+DisplayName = "Model Utility Toolkit"
+Icon = "🛠️"