diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..3dc5da3 --- /dev/null +++ b/__init__.py @@ -0,0 +1,117 @@ +# __init__.py +# This file is necessary to make Python treat the directory as a package. +# It's also where ComfyUI looks for node mappings. + +import sys +import traceback + +# Import node classes from nodes.py +# Renamed QuantizeScaled to QuantizeModel +try: + from .nodes import ModelToStateDict, QuantizeFP8Format, QuantizeModel, SaveAsSafeTensor +except ImportError: + # Fallback for direct execution or when relative imports fail + from nodes import ModelToStateDict, QuantizeFP8Format, QuantizeModel, SaveAsSafeTensor + +# Import ControlNet FP8 quantization nodes +ControlNetFP8QuantizeNode = None +ControlNetMetadataViewerNode = None +CONTROLNET_NODES_AVAILABLE = False + +try: + try: + from .controlnet_fp8_node import ControlNetFP8QuantizeNode, ControlNetMetadataViewerNode + except ImportError: + from controlnet_fp8_node import ControlNetFP8QuantizeNode, ControlNetMetadataViewerNode + CONTROLNET_NODES_AVAILABLE = True + print("✅ ControlNet FP8 nodes imported successfully") +except Exception as e: + print(f"❌ Failed to import ControlNet FP8 nodes: {e}") + print(f"❌ ControlNet error details: {traceback.format_exc()}") + +# Import GGUF quantization nodes +GGUFQuantizerNode = None # For the new GGUF node +GGUF_NODES_AVAILABLE = False + +try: + try: + from .gguf_quantizer_node import GGUFQuantizerNode # Updated filename + except ImportError: + from gguf_quantizer_node import GGUFQuantizerNode + GGUF_NODES_AVAILABLE = True + print("✅ GGUF quantization nodes imported successfully") +except Exception as e: + print(f"❌ Failed to import GGUF quantization nodes: {e}") + print(f"❌ GGUF error details: {traceback.format_exc()}") + GGUF_NODES_AVAILABLE = False + + +# A dictionary that ComfyUI uses to map node class names to node objects +NODE_CLASS_MAPPINGS = { + # Node Utils + "ModelToStateDict": ModelToStateDict, + # Direct FP8 Conversion + "QuantizeFP8Format": QuantizeFP8Format, + # Scaled Quantization + Casting (FP16/BF16) + "QuantizeModel": QuantizeModel, # Renamed class QuantizeScaled -> QuantizeModel + # Saving Node + "SaveAsSafeTensor": SaveAsSafeTensor, +} + +# Add ControlNet nodes if available +if CONTROLNET_NODES_AVAILABLE and ControlNetFP8QuantizeNode is not None: + NODE_CLASS_MAPPINGS.update({ + # ControlNet FP8 Quantization Nodes + "ControlNetFP8QuantizeNode": ControlNetFP8QuantizeNode, + "ControlNetMetadataViewerNode": ControlNetMetadataViewerNode, + }) + print("✅ ControlNet FP8 nodes registered in NODE_CLASS_MAPPINGS") + +# Add GGUF nodes if available +if GGUF_NODES_AVAILABLE and GGUFQuantizerNode is not None: + NODE_CLASS_MAPPINGS["GGUFQuantizerNode"] = GGUFQuantizerNode # New GGUF node + print("✅ GGUF quantization nodes registered in NODE_CLASS_MAPPINGS") +else: + print(f"❌ GGUF nodes NOT registered. Available: {GGUF_NODES_AVAILABLE}, Node: {GGUFQuantizerNode}") + + +# A dictionary that ComfyUI uses to map node class names to display names +NODE_DISPLAY_NAME_MAPPINGS = { + "ModelToStateDict": "Model To State Dict", + "QuantizeFP8Format": "Quantize Model to FP8 Format", + "QuantizeModel": "Quantize Model Scaled", # Display name for the renamed node + "SaveAsSafeTensor": "Save Model as SafeTensor", +} + +# Add ControlNet display names if available +if CONTROLNET_NODES_AVAILABLE: + NODE_DISPLAY_NAME_MAPPINGS.update({ + # ControlNet FP8 Quantization Node Display Names + "ControlNetFP8QuantizeNode": "ControlNet FP8 Quantizer", + "ControlNetMetadataViewerNode": "ControlNet Metadata Viewer", + }) + print("✅ ControlNet FP8 display names registered") + +# Add GGUF display names if available +if GGUF_NODES_AVAILABLE and GGUFQuantizerNode is not None: # Check new node for display name + NODE_DISPLAY_NAME_MAPPINGS.update({ + "GGUFQuantizerNode": "GGUF Quantizer 👾", # Display name for GGUF + }) + print("✅ GGUF quantization display names registered") + + +# Optional: Print a message to the console when the extension is loaded +print("----------------------------------------------------") +print("--- ComfyUI Quantization Node Pack Loaded ---") +print("--- Renamed QuantizeScaled to QuantizeModel ---") +# ... (other existing print messages you want to keep) ... +print("--- NEW: ControlNet FP8 Quantization Nodes ---") +print("--- NEW: GGUF Model Quantization ---") +print("--- Developed by [Lum3on] ---") # Remember to change this! +print("--- Version 0.8.2 ---") # Incremented version +print("----------------------------------------------------") + +# Tell ComfyUI where to find web files (for appearance.js) +WEB_DIRECTORY = "./web" + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] \ No newline at end of file