77 lines
2.7 KiB
Python
77 lines
2.7 KiB
Python
import logging
|
|
|
|
# =============================================================================
|
|
# Layout Registration
|
|
# =============================================================================
|
|
def _register_layouts():
|
|
try:
|
|
from comfy.quant_ops import QUANT_ALGOS, QuantizedLayout
|
|
|
|
# Register W8A8 (needed as fallback/mixed-precision in some models)
|
|
if "TensorCoreConvRotW8A8Layout" not in QUANT_ALGOS:
|
|
class TensorCoreConvRotW8A8Layout(QuantizedLayout):
|
|
pass
|
|
QUANT_ALGOS["TensorCoreConvRotW8A8Layout"] = TensorCoreConvRotW8A8Layout
|
|
logging.info("Int4Fast: Registered TensorCoreConvRotW8A8Layout dynamically.")
|
|
|
|
# Register W4A4
|
|
if "TensorCoreConvRotW4A4Layout" not in QUANT_ALGOS:
|
|
class TensorCoreConvRotW4A4Layout(QuantizedLayout):
|
|
pass
|
|
QUANT_ALGOS["TensorCoreConvRotW4A4Layout"] = TensorCoreConvRotW4A4Layout
|
|
logging.info("Int4Fast: Registered TensorCoreConvRotW4A4Layout dynamically.")
|
|
except Exception as e:
|
|
logging.warning(f"Int4Fast: Failed to register layouts dynamically: {e}")
|
|
|
|
_register_layouts()
|
|
|
|
# =============================================================================
|
|
# Export Custom Ops
|
|
# =============================================================================
|
|
try:
|
|
from .int4_quant import Int4Ops
|
|
except ImportError:
|
|
Int4Ops = None
|
|
|
|
# Register the native W4A4 implementation with comfy-kitchen as well as the
|
|
# custom loader path. This makes stock QuantizedTensor checkpoints dispatch to
|
|
# gfx1103 Triton without modifying the comfy-kitchen installation.
|
|
try:
|
|
from .kitchen_backend import register_kitchen_backend
|
|
|
|
KITCHEN_BACKEND_REGISTERED = register_kitchen_backend()
|
|
except Exception as e:
|
|
KITCHEN_BACKEND_REGISTERED = False
|
|
logging.warning(f"INT4 Fast: comfy-kitchen backend registration failed: {e}")
|
|
|
|
# =============================================================================
|
|
# Node Mappings
|
|
# =============================================================================
|
|
try:
|
|
from .int4_unet_loader import UNetLoaderINTW4A4
|
|
from .int4_save import INT4ModelSave
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"OTUNetLoaderW4A4": UNetLoaderINTW4A4,
|
|
"INT4ModelSave": INT4ModelSave,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"OTUNetLoaderW4A4": "Load Diffusion Model INT4 (W4A4)",
|
|
"INT4ModelSave": "Save Int4 Model",
|
|
}
|
|
except ImportError as e:
|
|
logging.error(f"Int4Fast: Failed to import nodes: {e}")
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
WEB_DIRECTORY = "./js"
|
|
|
|
__all__ = [
|
|
"NODE_CLASS_MAPPINGS",
|
|
"NODE_DISPLAY_NAME_MAPPINGS",
|
|
"WEB_DIRECTORY",
|
|
"Int4Ops",
|
|
"KITCHEN_BACKEND_REGISTERED",
|
|
]
|