custom_nodes/ComfyUI-INT8-Fast/int8_quant.py

import torch
from torch import Tensor, nn
import torch.nn.functional as F
import logging
import comfy.model_patcher
import comfy.memory_management
import comfy.model_management
import comfy.lora
import comfy.utils

try:
    import comfy_aimdo.host_buffer
    import comfy_aimdo.torch
    _AIMDO_FILE_SLICE_LOAD = True
except Exception:
    _AIMDO_FILE_SLICE_LOAD = False

# Add this at the top of your file
try:
    from .int8_fused_kernel import triton_int8_linear
    from .int8_fused_kernel import triton_int8_linear_per_row
    from .int8_fused_kernel import triton_quantize_rowwise
    _TRITON_AVAILABLE = True
except ImportError:
    _TRITON_AVAILABLE = False
    print("Triton not found, falling back to torch._int_mm")

# Runtime toggle — set by Int8TensorwiseOps.use_triton via the loader node
_use_triton = True

# ConvRot Configuration
CONVROT_GROUP_SIZE = 256  # Must be a power of 4 for Regular Hadamard (e.g. 16, 64, 256)

# --- Quantization Utils ---

def quantize_int8(x: Tensor, scale: float | Tensor) -> Tensor:
    return x.float().mul(1.0 / scale).round_().clamp_(-128.0, 127.0).to(torch.int8)

def quantize_int8_tensorwise(x: Tensor) -> tuple[Tensor, Tensor]:
    abs_max = x.abs().max()
    scale = (abs_max.float() / 127.0).clamp(min=1e-30)
    return quantize_int8(x, scale), scale

def quantize_int8_axiswise(x: Tensor, dim: int) -> tuple[Tensor, Tensor]:
    abs_max = x.abs().amax(dim=dim, keepdim=True)
    scale = (abs_max.float() / 127.0).clamp(min=1e-30)
    return quantize_int8(x, scale), scale

def dequantize(q: Tensor, scale: float | Tensor) -> Tensor:
    return q.float() * scale

def tensor_to_device_file_slice(tensor: Tensor, device: torch.device) -> Tensor:
    if (
        not _AIMDO_FILE_SLICE_LOAD
        or tensor.device.type != "cpu"
        or device is None
        or device.type != "cuda"
    ):
        return tensor.to(device, non_blocking=True)

    size = tensor.numel() * tensor.element_size()
    if size == 0:
        return tensor.to(device, non_blocking=True)

    hostbuf = comfy_aimdo.host_buffer.HostBuffer(size)
    host_tensor = comfy_aimdo.torch.hostbuf_to_tensor(hostbuf)
    host_view = host_tensor[:size].view(dtype=tensor.dtype).view(tensor.shape)
    if comfy.memory_management.read_tensor_file_slice_into(tensor, host_view):
        out = torch.empty_like(tensor, device=device)
        out.copy_(host_view, non_blocking=False)
        return out

    return tensor.to(device, non_blocking=True)

def stochastic_round_int8_delta(x: Tensor, scale: float | Tensor, seed: int = 0) -> Tensor:
    """
    Quantize a delta tensor to INT8 using stochastic rounding.
    Used for LoRA deltas to minimize quantization error.
    """
    generator = torch.Generator(device=x.device)
    generator.manual_seed(seed)
    
    # Scale to INT8 range — move scale to x's device to handle CPU-stored scales
    if isinstance(scale, torch.Tensor):
        scale = scale.to(x.device)
    x_scaled = x / scale
    
    # Stochastic rounding
    x_floor = torch.floor(x_scaled)
    fraction = x_scaled - x_floor
    del x_scaled # High-precision input no longer needed
    
    # Speed optimization: Create random values directly on the target device
    random_vals = torch.rand(x_floor.shape, generator=generator, device=x.device, dtype=x_floor.dtype)
    x_rounded = torch.where(random_vals < fraction, x_floor + 1, x_floor)
    
    del random_vals
    del fraction
    del x_floor
    
    return torch.clamp(x_rounded, -128, 127).to(torch.int8)



# --- LinearW8A8 Core ---

@torch.no_grad()
def int8_forward_dynamic(x: Tensor, weight: Tensor, weight_scale: float | Tensor, bias: Tensor | None, compute_dtype: torch.dtype) -> Tensor:
    """Forward with dynamic per-token activation quantization."""
    
    # --- FAST PATH: Triton Fused Kernel ---
    if _TRITON_AVAILABLE and _use_triton and x.is_cuda:
        return triton_int8_linear(x, weight, weight_scale, bias, compute_dtype)

    # --- SLOW PATH: Standard PyTorch ---
    # Quantize activations per row (dynamic)
    x_8, x_scale = quantize_int8_axiswise(x, dim=-1)
    
    # INT8 Matmul (Outputs Int32)
    res = torch._int_mm(x_8, weight.T)
    
    # Dequantize: (res * weight_scale * x_scale)
    # Note: Creating intermediate Float tensors here is VRAM heavy
    res_scaled = res.float().mul_(weight_scale * x_scale).to(compute_dtype)
    
    if bias is not None:
        res_scaled = res_scaled + bias.to(compute_dtype)
    return res_scaled


@torch.no_grad()
def int8_forward_dynamic_per_row(x: Tensor, weight: Tensor, weight_scale: Tensor, bias: Tensor | None, compute_dtype: torch.dtype) -> Tensor:
    """Forward with dynamic per-token activation quantization and per-row weight quantization.
    
    Args:
        x: Input activations [batch, in_features]
        weight: INT8 weight matrix [out_features, in_features]
        weight_scale: Per-row weight scales [out_features, 1]
        bias: Optional bias
        compute_dtype: Output dtype
    """
    # --- FAST PATH: Triton Fused Kernel (per-row) ---
    if _TRITON_AVAILABLE and _use_triton and x.is_cuda:
        return triton_int8_linear_per_row(x, weight, weight_scale, bias, compute_dtype)

    # --- SLOW PATH: Standard PyTorch ---
    x_8, x_scale = quantize_int8_axiswise(x, dim=-1)

    # INT8 Matmul (Outputs Int32)
    res = torch._int_mm(x_8, weight.T)  # [batch, out_features]
    
    # Dequantize with per-row weight scales
    # res[i,j] = sum_k(x_8[i,k] * weight[j,k]) * x_scale[i] * weight_scale[j]
    # Broadcasting: res * x_scale * weight_scale.T
    res_scaled = res.float().mul_(x_scale).mul_(weight_scale.T).to(compute_dtype)
    
    if bias is not None:
        res_scaled = res_scaled + bias.to(compute_dtype)
    return res_scaled

# =============================================================================
# Int8TensorwiseOps - ComfyUI Custom Operations
# =============================================================================

try:
    from comfy.ops import manual_cast, cast_bias_weight, uncast_bias_weight
    _COMFY_OPS_AVAILABLE = True
except ImportError:
    _COMFY_OPS_AVAILABLE = False


if _COMFY_OPS_AVAILABLE:
    class Int8TensorwiseOps(manual_cast):
        """
        Custom ComfyUI operations for INT8 tensorwise quantization.
        """
        excluded_names = []
        dynamic_quantize = False # Manual toggle for on-the-fly quantization
        enable_convrot = False # Toggle for ConvRot Hadamard rotation
        use_triton = True  # Toggle for Triton fused kernel (mirrors _use_triton)
        compute_dtype = None # Optional override for INT8 activation/output compute dtype
        _is_prequantized = False # Keep this as a status flag, but don't use for detection
        lora_mode = "None" # None/Stochastic bake into INT8 weights; Dynamic applies LoRA at inference
        dynamic_lora = False # If True, apply LoRA dynamically at inference; if False, bake into INT8 weights at load time
        lora_patches = {} # Map of model_key -> patch list (from load_lora)
        lora_strength = 1.0
        dynamic_load_device = None # Set by the loader when Aimdo should avoid a full CPU staging copy
        skeleton_meta_init = False # Temporary mode for LoRA key-map discovery
        _auto_compute_dtype_by_device = {}

        @staticmethod
        def _default_compute_dtype(x: Tensor) -> torch.dtype:
            if x.dtype in (torch.float16, torch.bfloat16):
                return x.dtype

            if x.dtype == torch.float32 and x.is_cuda:
                device_index = x.device.index
                if device_index is None:
                    device_index = torch.cuda.current_device()
                cached = Int8TensorwiseOps._auto_compute_dtype_by_device.get(device_index)
                if cached is not None:
                    return cached

                compute_dtype = torch.float32
                try:
                    capability = torch.cuda.get_device_capability(device_index)
                    name = torch.cuda.get_device_name(device_index).lower()
                    if capability == (7, 5) and ("rtx" in name or "t4" in name):
                        compute_dtype = torch.float16
                except Exception:
                    pass

                Int8TensorwiseOps._auto_compute_dtype_by_device[device_index] = compute_dtype
                return compute_dtype

            if x.dtype == torch.float32:
                return torch.float32
            return torch.float16
        
        class Linear(manual_cast.Linear):
            def __init__(self, in_features, out_features, bias=True, device=None, dtype=None):
                if getattr(Int8TensorwiseOps, "skeleton_meta_init", False):
                    nn.Module.__init__(self)
                    self.in_features = in_features
                    self.out_features = out_features
                    tensor_kwargs = {"device": "meta"}
                    if dtype is not None:
                        tensor_kwargs["dtype"] = dtype
                    self.weight = nn.Parameter(torch.empty((out_features, in_features), **tensor_kwargs), requires_grad=False)
                    self.bias = nn.Parameter(torch.empty((out_features,), **tensor_kwargs), requires_grad=False) if bias else None
                    self.weight_comfy_model_dtype = dtype
                    self.bias_comfy_model_dtype = dtype
                # Preserve ComfyUI's Windows/Aimdo lazy-init path. The base
                # disable_weight_init.Linear only takes this path for classes
                # that do not override _load_from_state_dict; this INT8 class
                # does override it, so calling super() would allocate full
                # skeleton weights during Pre-LoRA key-map discovery.
                elif comfy.model_management.WINDOWS and comfy.memory_management.aimdo_enabled:
                    nn.Module.__init__(self)
                    self.in_features = in_features
                    self.out_features = out_features
                    self.weight = None
                    self.bias = None
                    self.comfy_need_lazy_init_bias = bias
                    self.weight_comfy_model_dtype = dtype
                    self.bias_comfy_model_dtype = dtype
                else:
                    super().__init__(in_features, out_features, bias, device, dtype)
                self.register_buffer('weight_scale', None)
                self._is_quantized = False
                self._is_per_row = False  # Track quantization granularity
                self._use_convrot = False  # Track if ConvRot was applied
                self._weight_scale_scalar = None  # For scalar (non-tensor) scales
                self.compute_dtype = None
                self.comfy_cast_weights = False
                self.lora_patches = []  # List of (down_scaled, up, start, size) set by INT8ModelPatcher
            
            def reset_parameters(self):
                return None

            @staticmethod
            def _normalize_lora_key(key):
                if not isinstance(key, str):
                    return key
                for p in ["diffusion_model.", "model.diffusion_model.", "model.", "transformer."]:
                    if key.startswith(p):
                        return key[len(p):]
                return key

            @staticmethod
            def _is_bias_key(key):
                return isinstance(key, str) and key.endswith(".bias")

            @staticmethod
            def _format_lora_patches(patches):
                formatted = []
                for patch in patches or []:
                    if len(patch) == 4:
                        v, offset, function, strength = patch
                    else:
                        v, offset, function = patch
                        strength = getattr(Int8TensorwiseOps, "lora_strength", 1.0)
                    formatted.append((strength, v, 1.0, offset, function))
                return formatted

            def _apply_int8_lora_patches(self, tensor, key, patches, device):
                if not patches or tensor.dtype == torch.int8:
                    return tensor

                temp_dtype = comfy.model_management.lora_compute_dtype(device)
                tensor_temp = tensor_to_device_file_slice(tensor, device).to(dtype=temp_dtype)
                return comfy.lora.calculate_weight(self._format_lora_patches(patches), tensor_temp, key)

            def finalize_pending_int8(self):
                pending = getattr(self, "_pending_int8_finalize", None)
                if pending is None:
                    return False

                weight_key = pending["weight_key"]
                device = pending.get("device")
                if device is None:
                    device = torch.device("cuda") if torch.cuda.is_available() else self.weight.device

                weight_tensor = self.weight.detach()
                weight_tensor = self._apply_int8_lora_patches(weight_tensor, weight_key, pending.get("lora_patches"), device)

                if pending["quantize"]:
                    if not hasattr(Int8TensorwiseOps, '_logged_otf'):
                        print(f"INT8 Fast: Quantizing on-the-fly (ConvRot: {pending.get('enable_convrot', False)})")
                        Int8TensorwiseOps._logged_otf = True

                    w_gpu = tensor_to_device_file_slice(weight_tensor, device).float()

                    self._use_convrot = False
                    if pending.get("enable_convrot", False) and self.in_features % CONVROT_GROUP_SIZE == 0:
                        try:
                            from .convrot import build_hadamard, rotate_weight
                            H = build_hadamard(CONVROT_GROUP_SIZE, device=w_gpu.device, dtype=w_gpu.dtype)
                            w_gpu = rotate_weight(w_gpu, H, group_size=CONVROT_GROUP_SIZE)
                            self._use_convrot = True
                            # Stamp the active groupsize so INT8ModelSave can
                            # round-trip it deterministically.
                            self._convrot_groupsize = CONVROT_GROUP_SIZE
                        except ImportError as e:
                            logging.warning(f"INT8 Fast: ConvRot enabled but convrot module error: {e}")

                    q_weight, q_scale = quantize_int8_axiswise(w_gpu, dim=1)
                    self.weight = nn.Parameter(q_weight.cpu(), requires_grad=False)
                    self.register_buffer('weight_scale', q_scale.cpu())
                    self._weight_scale_scalar = None
                    self._is_quantized = True
                    self._is_per_row = True
                    del w_gpu, q_weight, q_scale
                else:
                    self.weight = nn.Parameter(weight_tensor.cpu(), requires_grad=False)

                self.weight_comfy_model_dtype = self.weight.dtype
                if self.weight_scale is not None:
                    self.weight_scale_comfy_model_dtype = self.weight_scale.dtype
                if self.bias is not None:
                    self.bias_comfy_model_dtype = self.bias.dtype

                delattr(self, "_pending_int8_finalize")
                return True
            
            def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
                weight_key = prefix + "weight"
                
                # Utility to normalize keys by stripping common prefixes
                def normalize_key(key):
                    return self._normalize_lora_key(key)

                def apply_lora_patches(tensor, key):
                    if self._is_bias_key(key) or not Int8TensorwiseOps.lora_patches or tensor.dtype == torch.int8:
                        return tensor
                    nk = normalize_key(key)
                    patches = Int8TensorwiseOps.lora_patches.get(nk)
                    if patches:
                        # Track applied patches
                        if not hasattr(Int8TensorwiseOps, 'applied_lora_patches'):
                            Int8TensorwiseOps.applied_lora_patches = set()
                        Int8TensorwiseOps.applied_lora_patches.add(nk)

                        # Print only if multiple sub-patches map to the same layer
                        # if "weight" in key and len(patches) > 1:
                        #     print(f"INT8 Fast: Baking multiple LoRA parts into {nk} ({len(patches)} sub-patches)")
                            
                        # ComfyUI dynamically patches during inference using lora_compute_dtype()
                        # On most modern GPUs, this evaluates to torch.float16. 
                        # We simulate that exact intermediate cast here to achieve a 1:1 binary match.
                        device = getattr(Int8TensorwiseOps, "dynamic_load_device", None)
                        if device is None:
                            device = tensor.device
                        result_temp = self._apply_int8_lora_patches(tensor, key, patches, device)
                        return result_temp.to(tensor.dtype)
                    return tensor

                def source_tensor(tensor):
                    if tensor is not None and getattr(Int8TensorwiseOps, "dynamic_load_device", None) is not None:
                        return tensor.cpu()
                    return tensor

                scale_key = prefix + "weight_scale"
                input_scale_key = prefix + "input_scale"
                bias_key = prefix + "bias"
                
                def pop_metadata(sd, p, k):
                    v = sd.pop(p + k, None)
                    if v is not None: return v
                    v = sd.pop("model." + p + k, None)
                    if v is not None: return v
                    if p.startswith("model."):
                        v = sd.pop(p[6:] + k, None)
                        if v is not None: return v
                    if p.startswith("diffusion_model."):
                        v = sd.pop("diffusion_model." + p + k, None)
                        if v is not None: return v
                    return None

                weight_scale = pop_metadata(state_dict, prefix, "weight_scale")
                comfy_quant_tensor = pop_metadata(state_dict, prefix, "comfy_quant")

                weight_tensor = state_dict.pop(weight_key, None)
                bias_tensor = state_dict.pop(bias_key, None)

                # Pop input_scale to clean state_dict, but ignore it
                _ = state_dict.pop(input_scale_key, None)
                
                quant_conf_parsed = None
                if comfy_quant_tensor is not None:
                    try:
                        import json
                        quant_conf_parsed = json.loads(bytes(comfy_quant_tensor.tolist()).decode('utf-8'))
                        if quant_conf_parsed.get("convrot", False):
                            self._use_convrot = True
                            Int8TensorwiseOps.enable_convrot = True  # Propagate globally for LoRA
                            if "convrot_groupsize" in quant_conf_parsed:
                                self._convrot_groupsize = int(quant_conf_parsed["convrot_groupsize"])
                                Int8TensorwiseOps._global_convrot_groupsize = self._convrot_groupsize
                    except Exception:
                        pass
                
                pending_weight_lora_patches = None
                if weight_tensor is not None and weight_tensor.dtype != torch.int8:
                    pending_weight_lora_patches = Int8TensorwiseOps.lora_patches.get(normalize_key(weight_key))

                defer_weight_lora = (
                    getattr(Int8TensorwiseOps, "dynamic_load_device", None) is not None
                    and pending_weight_lora_patches
                )

                # Apply LoRA patches to weight and bias once. With Aimdo, large
                # weight patches are deferred until KSampler/model load time so
                # the loader node stays cheap and VBAR geometry is finalized once.
                if weight_tensor is not None and not defer_weight_lora:
                    weight_tensor = apply_lora_patches(weight_tensor, weight_key)
                if bias_tensor is not None:
                    bias_tensor = apply_lora_patches(bias_tensor, bias_key)
                
                if weight_tensor is not None:
                    if weight_tensor.dtype == torch.int8 and weight_scale is not None:
                        # Load Quantized
                        self._is_quantized = True
                        self.weight = nn.Parameter(weight_tensor, requires_grad=False)
                        Int8TensorwiseOps._is_prequantized = True # Found a quantized layer

                        # Optional explicit hint from saved comfy_quant
                        per_row_hint = None
                        if isinstance(quant_conf_parsed, dict) and "per_row" in quant_conf_parsed:
                            per_row_hint = bool(quant_conf_parsed["per_row"])

                        if isinstance(weight_scale, torch.Tensor):
                            if weight_scale.numel() == 1:
                                # Scalar scale — store as float for speed
                                self._weight_scale_scalar = weight_scale.float().item()
                                self.register_buffer('weight_scale', weight_scale.float().reshape(1))
                                self._weight_scale_scalar = None
                            elif weight_scale.dim() == 2 and weight_scale.shape[1] == 1:
                                self.register_buffer('weight_scale', weight_scale.float())
                                self._weight_scale_scalar = None
                                self._is_per_row = True if per_row_hint is None else per_row_hint
                            else:
                                self.register_buffer('weight_scale', weight_scale.float())
                                self._weight_scale_scalar = None
                                self._is_per_row = False if per_row_hint is None else per_row_hint
                        else:
                            self.weight_scale = nn.Parameter(
                                torch.tensor(float(weight_scale), dtype=torch.float32), 
                                requires_grad=False
                            )
                            self.weight_scale = None
                            self._is_per_row = False if per_row_hint is None else per_row_hint
                            
                    elif weight_tensor.dtype in (torch.float16, torch.bfloat16, torch.float32):
                        # Load High-Precision
                        is_excluded = any(ex in prefix for ex in Int8TensorwiseOps.excluded_names)
                        is_dim1 = self.in_features == 1 or self.out_features == 1 or weight_tensor.ndim == 1
                        should_quantize = not (is_excluded or is_dim1 or not Int8TensorwiseOps.dynamic_quantize)
                        defer_finalize = (
                            getattr(Int8TensorwiseOps, "dynamic_load_device", None) is not None
                            and (should_quantize or pending_weight_lora_patches)
                        )
                        
                        if defer_finalize:
                            self._is_quantized = False
                            self.weight = nn.Parameter(source_tensor(weight_tensor), requires_grad=False)
                            self._pending_int8_finalize = {
                                "weight_key": weight_key,
                                "quantize": should_quantize,
                                "lora_patches": pending_weight_lora_patches,
                                "device": getattr(Int8TensorwiseOps, "dynamic_load_device", None),
                                "enable_convrot": getattr(Int8TensorwiseOps, "enable_convrot", False),
                            }
                            if pending_weight_lora_patches:
                                if not hasattr(Int8TensorwiseOps, 'applied_lora_patches'):
                                    Int8TensorwiseOps.applied_lora_patches = set()
                                Int8TensorwiseOps.applied_lora_patches.add(normalize_key(weight_key))
                        elif not should_quantize:
                            self._is_quantized = False
                            self.weight = nn.Parameter(source_tensor(weight_tensor), requires_grad=False)
                        else:
                            # Quantize on the fly
                            device = getattr(Int8TensorwiseOps, "dynamic_load_device", None)
                            if device is None:
                                device = torch.device("cuda") if torch.cuda.is_available() else weight_tensor.device
                            
                            # Log the first time we quantize in this loader pass
                            if not hasattr(Int8TensorwiseOps, '_logged_otf'):
                                print(f"INT8 Fast: Quantizing on-the-fly (ConvRot: {getattr(Int8TensorwiseOps, 'enable_convrot', False)})")
                                Int8TensorwiseOps._logged_otf = True

                            # Cast to float32 before rotation and scale computation
                            w_gpu = weight_tensor.to(device, non_blocking=True).float()
                            
                            self._use_convrot = False
                            if getattr(Int8TensorwiseOps, "enable_convrot", False) and self.in_features % CONVROT_GROUP_SIZE == 0:
                                try:
                                    import logging
                                    from .convrot import build_hadamard, rotate_weight
                                    H = build_hadamard(CONVROT_GROUP_SIZE, device=w_gpu.device, dtype=w_gpu.dtype)
                                    w_gpu = rotate_weight(w_gpu, H, group_size=CONVROT_GROUP_SIZE)
                                    self._use_convrot = True
                                    # Stamp the active groupsize on the module so
                                    # INT8ModelSave can serialize it deterministically
                                    # (instead of relying on the loader's default).
                                    self._convrot_groupsize = CONVROT_GROUP_SIZE
                                except ImportError as e:
                                    import logging
                                    logging.warning(f"INT8 Fast: ConvRot enabled but convrot module error: {e}")
                                    
                            q_weight, q_scale = quantize_int8_axiswise(w_gpu, dim=1)

                            q_weight = q_weight.cpu()
                            q_scale = q_scale.cpu()

                            self.weight = nn.Parameter(q_weight, requires_grad=False)
                            self.register_buffer('weight_scale', q_scale)
                            self._weight_scale_scalar = None
                            self._is_quantized = True
                            self._is_per_row = True
                            del w_gpu, q_weight, q_scale
                    else:
                        self._is_quantized = False
                        self.weight = nn.Parameter(source_tensor(weight_tensor), requires_grad=False)
                else:
                    missing_keys.append(weight_key)
                
                # Assign bias if it exists (already patched if needed)
                if bias_tensor is not None:
                    self.bias = nn.Parameter(source_tensor(bias_tensor), requires_grad=False)
                else:
                    self.bias = None

                # Update archived model dtypes so VBAR geometry uses the correct
                # sizes. archive_model_dtypes runs before state_dict loading, so
                # weight_comfy_model_dtype is stale (e.g. bfloat16 instead of int8).
                # Without this, VBAR allocates 2x the needed memory and the cast
                # buffer path misinterprets int8 data as bfloat16.
                if self.weight is not None:
                    self.weight_comfy_model_dtype = self.weight.dtype
                if self.weight_scale is not None:
                    self.weight_scale_comfy_model_dtype = self.weight_scale.dtype
                if self.bias is not None:
                    self.bias_comfy_model_dtype = self.bias.dtype

            def _get_weight_scale(self):
                return self.weight_scale

            def convert_weight(self, _weight, inplace=False):
                if not self._is_quantized:
                    return _weight
                return self.weight

            def set_weight(self, out_weight, inplace_update=False, seed=0, return_weight=False, **kwargs):
                if not self._is_quantized:
                    new_weight = out_weight.to(self.weight.dtype)
                    if return_weight:
                        return new_weight

                    if inplace_update:
                        self.weight.data.copy_(new_weight)
                    else:
                        self.weight = nn.Parameter(new_weight, requires_grad=False)
                    return

                if out_weight.dtype == torch.int8:
                    if return_weight:
                        return out_weight

                    if inplace_update:
                        self.weight.data.copy_(out_weight)
                    else:
                        self.weight = nn.Parameter(out_weight, requires_grad=False)
                    return

                # Re-quantize if fallback occurred
                new_weight = quantize_int8(out_weight, self._get_weight_scale())
                
                if return_weight:
                    return new_weight

                if inplace_update:
                    self.weight.data.copy_(new_weight)
                else:
                    self.weight = nn.Parameter(new_weight, requires_grad=False)

            def set_bias(self, out_bias, inplace_update=False, seed=0, return_weight=False, **kwargs):
                if out_bias is None: return None
                
                new_bias = out_bias
                if return_weight:
                    return new_bias

                if inplace_update:
                    if self.bias is not None:
                        self.bias.data.copy_(new_bias)
                else:
                    self.bias = nn.Parameter(new_bias, requires_grad=False)

            def forward(self, x: Tensor) -> Tensor:
                """Fast forward using torch._int_mm for quantized weights."""
                
                # Check if ComfyUI needs to manage weight transfer (VBAR, offloading, LoRA patches, etc.)
                # This mirrors the base class check in disable_weight_init.Linear.forward()
                need_cast = self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0
                
                if not self._is_quantized:
                    if need_cast:
                        weight, bias, offload_stream = cast_bias_weight(self, x, offloadable=True)
                        out = F.linear(x, weight, bias)
                        uncast_bias_weight(self, weight, bias, offload_stream)
                        return out
                    else:
                        if x.device != self.weight.device or x.dtype != self.weight.dtype:
                            weight = self.weight.to(device=x.device, dtype=x.dtype)
                            bias = self.bias.to(device=x.device, dtype=x.dtype) if self.bias is not None else None
                            return F.linear(x, weight, bias)
                        return F.linear(x, self.weight, self.bias)
                
                # INT8 quantized path
                if need_cast:
                    # VBAR / offload / lowvram path
                    weight, bias, offload_stream = cast_bias_weight(
                        self, input=None, dtype=torch.int8, device=x.device,
                        bias_dtype=x.dtype, offloadable=True
                    )
                else:
                    # Fast path: weights already on GPU, no functions to apply
                    weight = self.weight
                    bias = self.bias
                    offload_stream = None
                
                w_scale = self._get_weight_scale()
                if isinstance(w_scale, torch.Tensor) and w_scale.device != x.device:
                    w_scale = w_scale.to(x.device, non_blocking=True)
                
                compute_dtype = Int8TensorwiseOps.compute_dtype
                if compute_dtype is None:
                    compute_dtype = Int8TensorwiseOps._default_compute_dtype(x)

                x_shape = x.shape
                x_2d = x.reshape(-1, x_shape[-1])
                if x_2d.dtype != compute_dtype:
                    x_2d = x_2d.to(compute_dtype)
                
                if getattr(self, "_use_convrot", False):
                    from .convrot import build_hadamard, rotate_activation
                    group_size = getattr(self, "_convrot_groupsize", CONVROT_GROUP_SIZE)
                    H = build_hadamard(group_size, device=x.device, dtype=x_2d.dtype)
                    x_2d = rotate_activation(x_2d, H, group_size=group_size)
                
                # Sync the loader toggle to the module-level flag read by the forward fns
                import sys as _sys
                _mod = _sys.modules[__name__]
                _mod._use_triton = Int8TensorwiseOps.use_triton

                if x_2d.shape[0] > 16:
                    if self._is_per_row:
                        y = int8_forward_dynamic_per_row(x_2d, weight, w_scale, bias, compute_dtype)
                    else:
                        y = int8_forward_dynamic(x_2d, weight, w_scale, bias, compute_dtype)
                else:
                    # Small batch fallback
                    w_float = dequantize(weight, w_scale).to(x_2d.dtype)
                    bias_typed = bias.to(x_2d.dtype) if bias is not None else None
                    y = F.linear(x_2d, w_float, bias_typed)
                
                # Dynamic LoRA Path — handles split QKV via per-patch offsets
                for lora_down, lora_up, lora_start, lora_size in self.lora_patches:
                    lD = lora_down.to(x.device, non_blocking=True)
                    lU = lora_up.to(x.device, non_blocking=True)
                    lora_x = F.linear(x_2d.to(lD.dtype), lD)
                    lora_y = F.linear(lora_x, lU)  # [batch, slice_size or full_out]
                    if lora_start is not None:
                        y[:, lora_start:lora_start + lora_size] = (
                            y[:, lora_start:lora_start + lora_size] + lora_y.to(y.dtype)
                        )
                    else:
                        y = y + lora_y.to(y.dtype)
                
                if need_cast:
                    uncast_bias_weight(self, weight, bias, offload_stream)
                return y.reshape(*x_shape[:-1], y.shape[-1])
        
        # Pass-through for other layers
        class GroupNorm(manual_cast.GroupNorm): pass
        class LayerNorm(manual_cast.LayerNorm): pass
        class Conv2d(manual_cast.Conv2d): pass
        class Conv3d(manual_cast.Conv3d): pass
        class ConvTranspose2d(manual_cast.ConvTranspose2d): pass
        class Embedding(manual_cast.Embedding): pass
        
        @classmethod
        def conv_nd(cls, dims, *args, **kwargs):
            if dims == 2: return cls.Conv2d(*args, **kwargs)
            elif dims == 3: return cls.Conv3d(*args, **kwargs)
            else: raise ValueError(f"unsupported dimensions: {dims}")

# =============================================================================
# INT8 Model Patcher - Unified LoRA Handling
# =============================================================================

import inspect
try:
    _prefetch_sig = inspect.signature(comfy.lora.prefetch_prepared_value)
    _use_new_prefetch = len(_prefetch_sig.parameters) == 5
except Exception:
    _use_new_prefetch = False


class INT8LowVramPatch:
    is_lowvram_patch = True

    def __init__(self, key, patches, module, lora_mode):
        self.key = key
        self.patches = patches
        self.module = module
        self.lora_mode = lora_mode
        self.prepared_patches = None

    def memory_required(self):
        if not _use_new_prefetch:
            return 0
        counter = [0]
        for patch in self.patches[self.key]:
            comfy.lora.prefetch_prepared_value(patch[1], counter, None, None, False)
        return counter[0]

    def prepare(self, *args, **kwargs):
        if _use_new_prefetch:
            # 0.22.0+ signature: prepare(self, destination, stream, copy=True, commit=True)
            destination = args[0] if len(args) > 0 else kwargs.get("destination")
            stream = args[1] if len(args) > 1 else kwargs.get("stream")
            copy = args[2] if len(args) > 2 else kwargs.get("copy", True)
            commit = args[3] if len(args) > 3 else kwargs.get("commit", True)

            counter = [0]
            prepared_patches = [
                (patch[0], comfy.lora.prefetch_prepared_value(patch[1], counter, destination, stream, copy), patch[2], patch[3], patch[4])
                for patch in self.patches[self.key]
            ]
            if commit:
                self.prepared_patches = prepared_patches
            return prepared_patches
        else:
            # 0.21.1- signature: prepare(self, allocate_buffer, stream)
            allocate_buffer = args[0] if len(args) > 0 else kwargs.get("allocate_buffer")
            stream = args[1] if len(args) > 1 else kwargs.get("stream")

            self.prepared_patches = [
                (patch[0], comfy.lora.prefetch_prepared_value(patch[1], allocate_buffer, stream), patch[2], patch[3], patch[4])
                for patch in self.patches[self.key]
            ]
            return self.prepared_patches

    def clear_prepared(self):
        self.prepared_patches = None

    def __call__(self, weight):
        patches = self.prepared_patches if self.prepared_patches is not None else self.patches[self.key]
        scale = self.module._get_weight_scale()
        if isinstance(scale, torch.Tensor):
            scale = scale.to(weight.device)

        weight_float = dequantize(weight, scale)

        use_convrot = getattr(self.module, "_use_convrot", False)
        if use_convrot:
            group_size = getattr(self.module, "_convrot_groupsize", CONVROT_GROUP_SIZE)
            try:
                from .convrot import build_hadamard, rotate_weight
                H = build_hadamard(group_size, device=weight.device, dtype=weight_float.dtype)
                weight_float = rotate_weight(weight_float, H, group_size=group_size)
            except ImportError:
                use_convrot = False

        patched_weight_float = comfy.lora.calculate_weight(
            patches,
            weight_float,
            self.key,
            intermediate_dtype=weight_float.dtype,
        )

        if use_convrot:
            patched_weight_float = rotate_weight(patched_weight_float, H, group_size=group_size)

        if self.lora_mode == "Stochastic":
            return stochastic_round_int8_delta(
                patched_weight_float,
                scale,
                seed=comfy.utils.string_to_seed(self.key),
            )
        return quantize_int8(patched_weight_float, scale)


class INT8ModelPatcher(comfy.model_patcher.ModelPatcher):
    """
    Custom ModelPatcher that intercepts patching for INT8 layers.
    Routes patching through either a bake-in path (dequant-patch-requant)
    or a dynamic path (runtime injection), depending on the dynamic_lora toggle.
    """
    def finalize_pending_int8(self):
        finalized = 0
        for module in self.model.modules():
            finalize = getattr(module, "finalize_pending_int8", None)
            if finalize is not None and finalize():
                finalized += 1
        if finalized > 0:
            self.size = 0
            #logging.info(f"INT8 Fast: Finalized {finalized} deferred INT8 layer(s) at model load time.")
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    def patch_weight_to_device(self, key, device_to=None, inplace_update=False, return_weight=False, force_cast=False):
        if key not in self.patches and not force_cast:
            return super().patch_weight_to_device(key, device_to, inplace_update, return_weight, force_cast)

        # Check if this is one of our INT8 modules
        module_path = key.rsplit('.', 1)[0]
        try:
            module = comfy.utils.get_attr(self.model, module_path)
        except AttributeError:
            module = None

        is_int8_module = hasattr(module, "_is_quantized") and module._is_quantized
        patches = self.patches.get(key, [])

        if is_int8_module and Int8TensorwiseOps.Linear._is_bias_key(key):
            return comfy.utils.get_attr(self.model, key) if return_weight else None

        if is_int8_module:
            if not Int8TensorwiseOps.dynamic_lora:
                # --- BAKE-IN LORA PATH (Dequant → Patch → Quant) ---
                # Works with the native ComfyUI LoRA Loader (and also INT8LoraLoader).
                # All patches are applied in float space via ComfyUI's standard mechanism,
                # then the result is re-quantized back to INT8.

                # Identify current weight in the model
                current_weight = comfy.utils.get_attr(self.model, key)
                scale = module._get_weight_scale()

                if device_to is None:
                    device_to = current_weight.device

                # ALWAYS use the weight from backup as the source if it exists to prevent additive stacking.
                # If it doesn't exist, this is the first patch, so create it from the current model weight.
                if key not in self.backup:
                    import collections
                    BackupEntry = collections.namedtuple('Dimension', ['weight', 'inplace_update'])
                    self.backup[key] = BackupEntry(
                        weight=current_weight.to(device=self.offload_device, copy=inplace_update),
                        inplace_update=inplace_update,
                    )
                    source_weight = current_weight
                else:
                    # Use existing backup as source
                    source_weight = self.backup[key].weight

                # 1. Dequantize to float (move scale to device_to since it lives on CPU)
                if isinstance(scale, torch.Tensor):
                    scale = scale.to(device_to)
                weight_float = dequantize(source_weight.to(device_to), scale)

                # 2. Handle ConvRot: de-rotate into weight space before patching
                use_convrot = getattr(module, "_use_convrot", False)
                if use_convrot:
                    group_size = getattr(module, "_convrot_groupsize", CONVROT_GROUP_SIZE)
                    try:
                        from .convrot import build_hadamard, rotate_weight
                        H = build_hadamard(group_size, device=device_to, dtype=weight_float.dtype)
                        weight_float = rotate_weight(weight_float, H, group_size=group_size)
                    except ImportError:
                        pass

                # 3. Patch in float space using ComfyUI's standard mechanism.
                # calculate_weight handles LoRA, LoHA, LoKR, DoRA, etc.
                patches_list = self.patches.get(key, [])
                patched_weight_float = comfy.lora.calculate_weight(patches_list, weight_float, key)

                # 4. Handle ConvRot: re-rotate
                if use_convrot:
                    patched_weight_float = rotate_weight(patched_weight_float, H, group_size=group_size)

                # 5. Re-quantize back to INT8 using the original scale
                if getattr(Int8TensorwiseOps, "lora_mode", "None") == "Stochastic":
                    patched_weight_int8 = stochastic_round_int8_delta(patched_weight_float, scale)
                else:
                    patched_weight_int8 = quantize_int8(patched_weight_float, scale)

                # 6. Move back to original device and store
                patched_weight_int8 = patched_weight_int8.to(current_weight.device)

                if return_weight:
                    return patched_weight_int8

                if inplace_update:
                    current_weight.data.copy_(patched_weight_int8)
                else:
                    comfy.utils.set_attr(self.model, key, nn.Parameter(patched_weight_int8, requires_grad=False))
                return

            else:
                # --- DYNAMIC LORA PATH ---
                # Build a list of (down_scaled, up, start, size) per patch.
                # Keeping patches separate preserves the offset info needed for
                # fused QKV layers where each of Q/K/V targets a different output slice.
                weight = comfy.utils.get_attr(self.model, key)
                device = weight.device if weight is not None else self.offload_device
                lora_patches = []
                for p in patches:
                    strength_patch = p[0]  # float
                    adapter = p[1]         # the LoRA adapter object
                    strength_model = p[2]  # float
                    offset = p[3] if len(p) > 3 else None  # (dim, start, size) or None

                    if not hasattr(adapter, "weights"):
                        continue

                    strength = strength_patch * strength_model
                    weights = adapter.weights
                    # Standard LoRA: (up, down, alpha, mid, dora_scale, reshape)
                    if len(weights) == 6:
                        up, down, alpha, mid, dora, reshape = weights
                        rank = down.shape[0] if down.ndim >= 2 else 1
                        scale = (alpha / rank) * strength if alpha is not None else strength

                        down_scaled = down.flatten(1) * scale
                        if mid is not None:
                            down_scaled = torch.mm(mid.flatten(1), down.flatten(1)) * scale

                        # If this layer has ConvRot applied, rotate the 'down' matrix
                        # so the LoRA delta is coherent with the rotated weight basis:
                        #   W_rot = W @ H^T  =>  ΔW_rot = ΔW @ H^T  =>  rotate down only
                        if getattr(module, "_use_convrot", False) and down_scaled.shape[1] % CONVROT_GROUP_SIZE == 0:
                            try:
                                from .convrot import build_hadamard, rotate_weight
                                group_size = getattr(module, "_convrot_groupsize", CONVROT_GROUP_SIZE)
                                H = build_hadamard(group_size, device=down_scaled.device, dtype=down_scaled.dtype)
                                down_scaled = rotate_weight(down_scaled, H, group_size=group_size)
                            except ImportError:
                                pass

                        # Extract offset: which output rows this patch targets
                        start, size = None, None
                        if offset is not None:
                            _dim, start, size = offset  # dim is always 0 for linear weights

                        lora_patches.append((down_scaled.to(device), up.flatten(1).to(device), start, size))

                module.lora_patches = lora_patches
                if return_weight:
                    return weight
                return  # Skip standard weight-merging path

        # --- NON-INT8 MODULE PATH ---
        return super().patch_weight_to_device(key, device_to, inplace_update, return_weight, force_cast)

    def load(self, *args, **kwargs):
        self.finalize_pending_int8()

        save_materialized = bool(getattr(self, "_int8_save_materialized_lora", False))

        if not Int8TensorwiseOps.dynamic_lora and not save_materialized:
            for k in list(self.backup):
                if k in self.patches:
                    try:
                        module = comfy.utils.get_attr(self.model, k.rsplit('.', 1)[0])
                    except AttributeError:
                        module = None
                    if hasattr(module, "_is_quantized") and module._is_quantized:
                        bk = self.backup.pop(k)
                        if bk.inplace_update:
                            dest = comfy.utils.get_attr(self.model, k)
                            dest.data.copy_(bk.weight)
                        else:
                            comfy.utils.set_attr(self.model, k, bk.weight)

        # Cleanup: Revert any keys that are in backup but no longer in patches (stale patches)
        # This ensures that when a LoRA is disabled, the model returns to its base state.
        stale_keys = [k for k in self.backup if k not in self.patches]
        for k in stale_keys:
            bk = self.backup.pop(k)
            if bk.inplace_update:
                dest = comfy.utils.get_attr(self.model, k)
                dest.data.copy_(bk.weight)
            else:
                comfy.utils.set_attr(self.model, k, bk.weight)
        
        # Cleanup: Clear stale dynamic LoRA patches.
        # This prevents LoRA from "sticking" when dynamic_lora is toggled or LoRAs are disabled.
        for name, module in self.model.named_modules():
            if hasattr(module, "lora_patches") and module.lora_patches:
                # If dynamic LoRA is disabled globally, or if this module has no active patches, clear them.
                if not Int8TensorwiseOps.dynamic_lora or (name + ".weight") not in self.patches:
                    module.lora_patches = []

        res = super().load(*args, **kwargs) if hasattr(super(), "load") else None
        
        device_to = kwargs.get("device_to", args[0] if len(args) > 0 else self.model.device)
        
        for name, module in self.model.named_modules():
            if hasattr(module, "_is_quantized") and module._is_quantized:
                weight_key = name + ".weight"
                
                if weight_key in self.patches:
                    if save_materialized:
                        if hasattr(module, "weight_lowvram_function"):
                            module.weight_lowvram_function = None
                        if hasattr(module, "weight_function"):
                            module.weight_function = [f for f in getattr(module, "weight_function", []) if type(f).__name__ != "LowVramPatch"]
                    elif Int8TensorwiseOps.dynamic_lora:
                        if hasattr(module, "weight_lowvram_function"):
                            module.weight_lowvram_function = None
                        if hasattr(module, "weight_function"):
                            module.weight_function = [f for f in getattr(module, "weight_function", []) if type(f).__name__ != "LowVramPatch"]
                        self.patch_weight_to_device(weight_key, device_to=device_to)
                    else:
                        if hasattr(module, "weight_function"):
                            module.weight_function = [f for f in getattr(module, "weight_function", []) if type(f).__name__ != "LowVramPatch"]
                            
                        lowvram_patch = INT8LowVramPatch(
                            weight_key,
                            self.patches,
                            module,
                            getattr(Int8TensorwiseOps, "lora_mode", "None"),
                        )
                        pin_state = getattr(self.model, "dynamic_pins", {}).get(self.load_device, None)
                        if pin_state is not None:
                            lowvram_patch._pin_state = pin_state
                        module.weight_lowvram_function = lowvram_patch
                    
        return res

    def unpatch_model(self, device_to=None, unpatch_weights=True):
        if unpatch_weights:
            for name, module in self.model.named_modules():
                if hasattr(module, "lora_patches"):
                    module.lora_patches = []
        return super().unpatch_model(device_to, unpatch_weights)

    def clone(self, *args, **kwargs):
        src_cls = self.__class__
        
        if src_cls is INT8ModelPatcher:
            return super().clone(*args, **kwargs)
            
        if not issubclass(src_cls, INT8ModelPatcher):
            name = f"INT8_{src_cls.__name__}"
            dynamic_cls = type(name, (INT8ModelPatcher, src_cls), {})
        else:
            dynamic_cls = src_cls
            
        self.__class__ = dynamic_cls
        
        # Static clones do not need a disk reload factory. Dynamic-to-static
        # delegates do: sharing the dynamic model object makes ComfyUI treat
        # the static copy as a replacement instead of an independent model.
        if not self.is_dynamic() and getattr(self, "cached_patcher_init", None) is None:
            self.cached_patcher_init = (lambda *a, **kw: self, ())
            
        n = super().clone(*args, **kwargs)
        
        # If disable_dynamic is True, the core strips dynamic wrappers. We must re-apply INT8!
        disable_dyn = kwargs.get("disable_dynamic", False)
        if len(args) > 0:
            disable_dyn = args[0]
            
        if disable_dyn and not issubclass(n.__class__, INT8ModelPatcher):
            new_cls = type(f"INT8_{n.__class__.__name__}", (INT8ModelPatcher, n.__class__), {})
            n.__class__ = new_cls

        self.__class__ = src_cls
        return n