Merge pull request #24 from silveroxides/fix/mixed-format-loading

Fix mixed dtype loading and fix loading of text encoders(especially i…

Author: silveroxides

bnb4bit_ops.py

@@ -22,6 +22,7 @@
 import torch.nn.functional as F
 import logging
 from comfy.ops import manual_cast, cast_bias_weight, uncast_bias_weight
+from unifiedefficientloader import tensor_to_dict
 
 
 # NF4 (Normal Float 4-bit) quantization table
@@ -68,13 +69,6 @@
 ], dtype=torch.float32)
 
 
-def tensor_to_dict(tensor_data: torch.Tensor) -> dict:
-    """Convert a uint8 tensor containing JSON bytes back to a dictionary."""
-    byte_data = bytes(tensor_data.cpu().tolist())
-    json_str = byte_data.decode("utf-8")
-    return json.loads(json_str)
-
-
 def get_quant_map(quant_type: str, device: torch.device) -> torch.Tensor:
     """Get the quantization codebook for NF4 or FP4."""
     if quant_type == "nf4":

comfy_quant_helpers.py

@@ -20,6 +20,9 @@ classifiers = [
     "Programming Language :: Python :: 3.12",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
+dependencies = [
+    "unifiedefficientloader>=0.2.0"
+]
 
 [project.urls]
 Repository = "https://github.com/silveroxides/ComfyUI-QuantOps"

fp8_ops.py

@@ -409,6 +409,12 @@ def _int8_gemm_pytorch_fallback(
     b_fp32 = b_blocked.to(torch.float32) * b_scale_broadcast
     b_fp32 = b_fp32.permute(0, 2, 1, 3).reshape(N, K)
 
+    # Bias may arrive in bfloat16 when previous layers (e.g. TensorWiseINT8Layout)
+    # output bfloat16 and cast_bias_weight matches the input dtype.
+    # The INT8 fallback computes in float32, so bias must match.
+    if bias is not None and bias.dtype != a_fp32.dtype:
+        bias = bias.to(a_fp32.dtype)
+
     output = torch.nn.functional.linear(a_fp32, b_fp32, bias)
     return output
 

int8_ops.py

@@ -14,14 +14,19 @@
 from comfy.ops import manual_cast, cast_bias_weight, uncast_bias_weight
 from comfy.quant_ops import QuantizedTensor, QUANT_ALGOS, get_layout_class
 from comfy.model_patcher import LowVramPatch
+from unifiedefficientloader import tensor_to_dict
 
 # Try to import INT8 layouts
 try:
-    from .quant_layouts.int8_layout import BlockWiseINT8Layout
+    from comfy_kitchen.tensor.int8 import BlockWiseINT8Layout
     _HAS_INT8_LAYOUT = True
 except ImportError:
-    _HAS_INT8_LAYOUT = False
-    logging.warning("INT8 blockwise layout not available")
+    try:
+        from .quant_layouts.int8_layout import BlockWiseINT8Layout
+        _HAS_INT8_LAYOUT = True
+    except ImportError:
+        _HAS_INT8_LAYOUT = False
+        logging.warning("INT8 blockwise layout not available")
 
 try:
     from comfy_kitchen.tensor.int8 import TensorWiseINT8Layout
@@ -31,22 +36,6 @@
     logging.warning("INT8 tensorwise layout not available from comfy_kitchen")
 
 
-def tensor_to_dict(tensor_data: torch.Tensor) -> dict:
-    """
-    Convert a torch.uint8 tensor containing JSON bytes to a dictionary.
-    """
-    try:
-        if tensor_data.dtype == torch.uint8:
-            byte_data = bytes(tensor_data.tolist())
-            json_str = byte_data.decode("utf-8")
-            return json.loads(json_str)
-        else:
-            return {}
-    except Exception as e:
-        logging.debug(f"Failed to parse comfy_quant metadata using tensor_to_dict: {e}")
-        return {}
-
-
 class UnifiedQuantOps(manual_cast):
     """
     Unified operations class that handles INT8, FP8, MXFP8, and NVFP4 formats.

nodes/loader_nodes.py

@@ -1,12 +1,11 @@
 """QuantOps utilities."""
-from .safetensors_loader import (
-    MemoryEfficientSafeOpen,
-    load_fp8_state_dict,
-    get_layer_metadata,
-)
+from unifiedefficientloader import UnifiedSafetensorsLoader, tensor_to_dict
+from .safetensors_loader import extract_quantization_metadata, detect_quant_format, _is_scale_tensor
 
 __all__ = [
-    "MemoryEfficientSafeOpen",
-    "load_fp8_state_dict",
-    "get_layer_metadata",
+    "UnifiedSafetensorsLoader",
+    "tensor_to_dict",
+    "extract_quantization_metadata",
+    "detect_quant_format",
+    "_is_scale_tensor",
 ]