117 lines
4.7 KiB
Python
117 lines
4.7 KiB
Python
# __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'] |