diff --git a/.github/workflows/publish_action.yml b/.github/workflows/publish_action.yml new file mode 100644 index 0000000..c3112f9 --- /dev/null +++ b/.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. \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..bf2efef --- /dev/null +++ b/__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", +} \ No newline at end of file diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..c59fb29 --- /dev/null +++ b/nodes/__init__.py @@ -0,0 +1,8 @@ +from .metakeys import * + +__all__ = ( + "UNetMetaKeys", + "CLIPMetaKeys", + "LoRAMetaKeys", + "CheckpointMetaKeys", +) \ No newline at end of file diff --git a/nodes/metakeys.py b/nodes/metakeys.py new file mode 100644 index 0000000..e4e7359 --- /dev/null +++ b/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") diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..b2015d2 --- /dev/null +++ b/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 = "🛠️"