import torch from torch import Tensor, nn import torch.nn.functional as F import logging import json import comfy.model_patcher import comfy.memory_management import comfy.model_management import comfy.lora import comfy.utils from comfy.ops import manual_cast, cast_bias_weight, uncast_bias_weight from comfy.quant_ops import QuantizedTensor, TensorCoreConvRotW4A4Layout from .int8_quant import ( int8_forward_dynamic, int8_forward_dynamic_per_row, dequantize as dequantize_int8, Int8TensorwiseOps ) _use_new_prefetch = False import inspect try: _prefetch_sig = inspect.signature(comfy.lora.prefetch_prepared_value) _use_new_prefetch = len(_prefetch_sig.parameters) == 5 except Exception: pass fp8_types = [] for tname in ("float8_e4m3fn", "float8_e5m2"): if hasattr(torch, tname): fp8_types.append(getattr(torch, tname)) def _copy_tensor_to_device(tensor, device): """Materialize an owned tensor on device without an async source lifetime.""" if tensor is None or tensor.device == device: return tensor result = torch.empty(tensor.shape, dtype=tensor.dtype, device=device) result.copy_(tensor.detach(), non_blocking=False) return result def _copy_tensor_to_host(tensor): """Materialize an owned tensor in pageable CPU memory.""" source = tensor.detach() if source.device.type == "cuda": torch.cuda.synchronize(source.device) host = torch.empty(source.shape, dtype=source.dtype, device="cpu") host.copy_(source, non_blocking=False) torch.cuda.synchronize(source.device) return host return source.to(device="cpu", copy=True) class Int4Ops(manual_cast): excluded_names = [] dynamic_quantize = False enable_convrot = True convrot_groupsize = 256 quant_group_size = 64 linear_dtype = "int4" lora_mode = "None" dynamic_lora = False lora_patches = {} dynamic_load_device = None skeleton_meta_init = False _is_prequantized = False class Linear(manual_cast.Linear): def __init__(self, in_features, out_features, bias=True, device=None, dtype=None): if getattr(Int4Ops, "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 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._is_quantized = False self._quant_format = "convrot_w4a4" self._convrot_groupsize = 256 self._quant_group_size = 64 self._linear_dtype = "int4" self._is_per_row = False self._use_convrot = False self.comfy_cast_weights = False self.lora_patches = [] 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") def finalize_pending_int8(self): pending = getattr(self, "_pending_int4_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 if isinstance(weight_tensor, QuantizedTensor) or (hasattr(weight_tensor, "_qdata") and hasattr(weight_tensor, "_params")): weight_tensor = weight_tensor.dequantize() if hasattr(weight_tensor, "to"): weight_tensor = weight_tensor.to(device=device) else: weight_tensor = weight_tensor.detach() if pending.get("cast_fp8") or weight_tensor.dtype in fp8_types: weight_tensor = weight_tensor.to(dtype=torch.bfloat16) if pending.get("lora_patches"): temp_dtype = comfy.model_management.lora_compute_dtype(device) w_temp = weight_tensor.to(device=device, dtype=temp_dtype) formatted = [] for patch in pending["lora_patches"]: if len(patch) == 4: v, offset, function, strength = patch else: v, offset, function = patch strength = 1.0 formatted.append((strength, v, 1.0, offset, function)) weight_tensor = comfy.lora.calculate_weight(formatted, w_temp, weight_key).cpu() if pending["quantize"]: w_gpu = weight_tensor.to(device).float() stochastic_rounding = 1 if getattr(Int4Ops, "lora_mode", "None") == "Stochastic" else 0 qdata, params = TensorCoreConvRotW4A4Layout.quantize( w_gpu, convrot_groupsize=pending.get("convrot_groupsize", 256), quant_group_size=pending.get("quant_group_size", 64), stochastic_rounding=stochastic_rounding, linear_dtype=pending.get("linear_dtype", "int4"), ) params_cpu = TensorCoreConvRotW4A4Layout.Params( scale=params.scale.cpu() if isinstance(params.scale, torch.Tensor) else params.scale, orig_dtype=params.orig_dtype, orig_shape=params.orig_shape, convrot_groupsize=params.convrot_groupsize, quant_group_size=params.quant_group_size, linear_dtype=params.linear_dtype, ) self.weight = nn.Parameter(QuantizedTensor(qdata.cpu(), "TensorCoreConvRotW4A4Layout", params_cpu), requires_grad=False) self._is_quantized = True self._quant_format = "convrot_w4a4" self._convrot_groupsize = params.convrot_groupsize self._quant_group_size = params.quant_group_size self._linear_dtype = params.linear_dtype del w_gpu, qdata else: self.weight = nn.Parameter(weight_tensor.cpu(), requires_grad=False) self.weight_comfy_model_dtype = self.weight.dtype if self.bias is not None: self.bias_comfy_model_dtype = self.bias.dtype delattr(self, "_pending_int4_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" bias_key = prefix + "bias" def normalize_key(key): return self._normalize_lora_key(key) 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) is_fp8 = weight_tensor is not None and weight_tensor.dtype in fp8_types quant_conf_parsed = None quant_format = "convrot_w4a4" if comfy_quant_tensor is not None: try: quant_conf_parsed = json.loads(bytes(comfy_quant_tensor.tolist()).decode('utf-8')) quant_format = quant_conf_parsed.get("format", "convrot_w4a4") if "convrot_groupsize" in quant_conf_parsed: self._convrot_groupsize = int(quant_conf_parsed["convrot_groupsize"]) if "quant_group_size" in quant_conf_parsed: self._quant_group_size = int(quant_conf_parsed["quant_group_size"]) if "linear_dtype" in quant_conf_parsed: self._linear_dtype = quant_conf_parsed["linear_dtype"] if quant_conf_parsed.get("convrot", False): self._use_convrot = True except Exception: pass self._quant_format = quant_format pending_weight_lora_patches = None if weight_tensor is not None and not isinstance(weight_tensor, QuantizedTensor) and weight_tensor.dtype != torch.int8: pending_weight_lora_patches = Int4Ops.lora_patches.get(normalize_key(weight_key)) defer_weight_lora = ( getattr(Int4Ops, "dynamic_load_device", None) is not None and pending_weight_lora_patches ) if weight_tensor is not None and not defer_weight_lora and pending_weight_lora_patches: device = getattr(Int4Ops, "dynamic_load_device", None) if device is None: device = weight_tensor.device if hasattr(weight_tensor, "to"): weight_tensor_resolved = weight_tensor.to(device=device) else: weight_tensor_resolved = weight_tensor if isinstance(weight_tensor_resolved, QuantizedTensor) or (hasattr(weight_tensor_resolved, "_qdata") and hasattr(weight_tensor_resolved, "_params")): weight_tensor_resolved = weight_tensor_resolved.dequantize() if is_fp8 or weight_tensor_resolved.dtype in fp8_types: weight_tensor_resolved = weight_tensor_resolved.to(dtype=torch.bfloat16) temp_dtype = comfy.model_management.lora_compute_dtype(device) w_temp = weight_tensor_resolved.to(device=device, dtype=temp_dtype) formatted = [] for patch in pending_weight_lora_patches: if len(patch) == 4: v, offset, function, strength = patch else: v, offset, function = patch strength = 1.0 formatted.append((strength, v, 1.0, offset, function)) weight_tensor = comfy.lora.calculate_weight(formatted, w_temp, weight_key).cpu() if not hasattr(Int4Ops, 'applied_lora_patches'): Int4Ops.applied_lora_patches = set() Int4Ops.applied_lora_patches.add(normalize_key(weight_key)) is_fp8 = False if bias_tensor is not None: bias_patches = Int4Ops.lora_patches.get(normalize_key(bias_key)) if bias_patches: device = getattr(Int4Ops, "dynamic_load_device", None) if device is None: device = bias_tensor.device temp_dtype = comfy.model_management.lora_compute_dtype(device) b_temp = bias_tensor.to(device=device, dtype=temp_dtype) formatted = [] for patch in bias_patches: if len(patch) == 4: v, offset, function, strength = patch else: v, offset, function = patch strength = 1.0 formatted.append((strength, v, 1.0, offset, function)) bias_tensor = comfy.lora.calculate_weight(formatted, b_temp, bias_key).cpu() if not hasattr(Int4Ops, 'applied_lora_patches'): Int4Ops.applied_lora_patches = set() Int4Ops.applied_lora_patches.add(normalize_key(bias_key)) def source_tensor(tensor): if tensor is not None and getattr(Int4Ops, "dynamic_load_device", None) is not None: return tensor.cpu() return tensor if weight_tensor is not None: if (weight_tensor.dtype == torch.int8 or isinstance(weight_tensor, QuantizedTensor)) and weight_scale is not None: self._is_quantized = True Int4Ops._is_prequantized = True if quant_format == "int8_tensorwise": self.weight = nn.Parameter(weight_tensor, requires_grad=False) 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: self.register_buffer('weight_scale', weight_scale.float().reshape(1)) self._is_per_row = False elif weight_scale.dim() == 2 and weight_scale.shape[1] == 1: self.register_buffer('weight_scale', weight_scale.float()) 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._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._is_per_row = False if per_row_hint is None else per_row_hint else: orig_dtype = getattr(self, "weight_comfy_model_dtype", torch.bfloat16) if orig_dtype is None: orig_dtype = torch.bfloat16 orig_shape = (self.out_features, self.in_features) if isinstance(weight_scale, torch.Tensor): weight_scale = weight_scale.reshape(-1) params = TensorCoreConvRotW4A4Layout.Params( scale=weight_scale, orig_dtype=orig_dtype, orig_shape=orig_shape, convrot_groupsize=self._convrot_groupsize, quant_group_size=self._quant_group_size, linear_dtype=self._linear_dtype, ) q_data = weight_tensor._qdata if isinstance(weight_tensor, QuantizedTensor) else weight_tensor q_weight = QuantizedTensor(q_data, "TensorCoreConvRotW4A4Layout", params) self.weight = nn.Parameter(q_weight, requires_grad=False) elif weight_tensor.dtype in (torch.float16, torch.bfloat16, torch.float32) or is_fp8: is_excluded = any(ex in prefix for ex in Int4Ops.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 Int4Ops.dynamic_quantize) defer_finalize = ( getattr(Int4Ops, "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_int4_finalize = { "weight_key": weight_key, "quantize": should_quantize, "lora_patches": pending_weight_lora_patches, "device": getattr(Int4Ops, "dynamic_load_device", None), "convrot_groupsize": getattr(Int4Ops, "convrot_groupsize", 256), "quant_group_size": getattr(Int4Ops, "quant_group_size", 64), "linear_dtype": getattr(Int4Ops, "linear_dtype", "int4"), "cast_fp8": is_fp8, } if pending_weight_lora_patches: if not hasattr(Int4Ops, 'applied_lora_patches'): Int4Ops.applied_lora_patches = set() Int4Ops.applied_lora_patches.add(normalize_key(weight_key)) elif not should_quantize: self._is_quantized = False if is_fp8 and hasattr(weight_tensor, "to"): weight_tensor = weight_tensor.to(dtype=torch.bfloat16) self.weight = nn.Parameter(source_tensor(weight_tensor), requires_grad=False) else: device = getattr(Int4Ops, "dynamic_load_device", None) if device is None: device = torch.device("cuda") if torch.cuda.is_available() else weight_tensor.device if hasattr(weight_tensor, "to"): w_gpu = weight_tensor.to(device=device).to(dtype=torch.float32) else: w_gpu = weight_tensor.to(device).float() stochastic_rounding = 1 if getattr(Int4Ops, "lora_mode", "None") == "Stochastic" else 0 qdata, params = TensorCoreConvRotW4A4Layout.quantize( w_gpu, convrot_groupsize=getattr(Int4Ops, "convrot_groupsize", 256), quant_group_size=getattr(Int4Ops, "quant_group_size", 64), stochastic_rounding=stochastic_rounding, linear_dtype=getattr(Int4Ops, "linear_dtype", "int4"), ) params_cpu = TensorCoreConvRotW4A4Layout.Params( scale=params.scale.cpu() if isinstance(params.scale, torch.Tensor) else params.scale, orig_dtype=params.orig_dtype, orig_shape=params.orig_shape, convrot_groupsize=params.convrot_groupsize, quant_group_size=params.quant_group_size, linear_dtype=params.linear_dtype, ) self.weight = nn.Parameter(QuantizedTensor(qdata.cpu(), "TensorCoreConvRotW4A4Layout", params_cpu), requires_grad=False) self._is_quantized = True self._quant_format = "convrot_w4a4" self._convrot_groupsize = params.convrot_groupsize self._quant_group_size = params.quant_group_size self._linear_dtype = params.linear_dtype del w_gpu, qdata else: self._is_quantized = False self.weight = nn.Parameter(source_tensor(weight_tensor), requires_grad=False) else: missing_keys.append(weight_key) if bias_tensor is not None: self.bias = nn.Parameter(source_tensor(bias_tensor), requires_grad=False) else: self.bias = None if self.weight is not None: self.weight_comfy_model_dtype = self.weight.dtype if self.bias is not None: self.bias_comfy_model_dtype = self.bias.dtype def forward(self, x: Tensor) -> Tensor: if not self._is_quantized: # Sensitive/excluded layers are intentionally retained in # BF16 by mixed-precision INT4 checkpoints. Always use # ComfyUI's manual-cast path here so their weight/device dtype # follows the current activation and dynamic LoRA patches are # applied. Quantized layers continue through the native # INT4/INT8 branches below. return self.forward_comfy_cast_weights(x) if getattr(self, "_quant_format", "convrot_w4a4") == "int8_tensorwise": from . import int8_quant int8_quant._use_triton = getattr(Int8TensorwiseOps, "use_triton", True) need_cast = self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0 if need_cast: weight, bias, offload_stream = cast_bias_weight( self, input=None, dtype=torch.int8, device=x.device, bias_dtype=x.dtype, offloadable=True ) else: weight = self.weight bias = self.bias offload_stream = None if weight is not None and weight.device != x.device: weight = weight.to(x.device, non_blocking=True) if bias is not None and bias.device != x.device: bias = bias.to(x.device, non_blocking=True) w_scale = self.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 = getattr(Int4Ops, "compute_dtype", None) 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 getattr(self, "_use_convrot", False): group_size = getattr(self, "_convrot_groupsize", 256) from .native_int4 import is_gfx1103_target, triton_rotate_activation if is_gfx1103_target(): if not x_2d.is_contiguous(): raise ValueError( "gfx1103 native INT8 ConvRot activation must be contiguous; " "the hot path will not make an eager copy" ) x_2d = triton_rotate_activation( x_2d, group_size, output_dtype=compute_dtype, ) else: if x_2d.dtype != compute_dtype: x_2d = x_2d.to(compute_dtype) from .convrot import build_hadamard, rotate_activation H = build_hadamard(group_size, device=x.device, dtype=x_2d.dtype) x_2d = rotate_activation(x_2d, H, group_size=group_size) elif x_2d.dtype != compute_dtype: x_2d = x_2d.to(compute_dtype) if getattr(self, "_is_per_row", False): 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) if self.lora_patches: 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) 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]) else: from .native_int4 import is_gfx1103_target, triton_convrot_w4a4_linear if is_gfx1103_target(): source_weight = self.weight if not isinstance(source_weight, QuantizedTensor): raise RuntimeError("gfx1103 native INT4 path expected a QuantizedTensor weight") # Never send a packed ConvRot QuantizedTensor through # cast_bias_weight(): requesting dtype=int8 changes the # logical output dtype and forces an unsupported dequantize. # Move the two physical tensors independently instead. params = source_weight.params if getattr(params, "transposed", False): raise RuntimeError("gfx1103 native INT4 path does not accept a transposed packed weight") qdata = _copy_tensor_to_device(source_weight._qdata, x.device) weight_scale = _copy_tensor_to_device(params.scale, x.device) bias = _copy_tensor_to_device(self.bias, x.device) for function in self.bias_function: bias = function(bias) compute_dtype = getattr(Int4Ops, "compute_dtype", None) if compute_dtype is None: compute_dtype = x.dtype if x.dtype in (torch.float16, torch.bfloat16) else params.orig_dtype if compute_dtype not in (torch.float16, torch.bfloat16): raise TypeError( "gfx1103 native INT4 path requires FP16 or BF16 compute; " f"got {compute_dtype}" ) y = triton_convrot_w4a4_linear( x, qdata, weight_scale, bias, convrot_groupsize=params.convrot_groupsize, quant_group_size=params.quant_group_size, linear_dtype=params.linear_dtype, output_dtype=compute_dtype, ) else: y = super().forward(x) if self.lora_patches: x_shape = x.shape y_shape = y.shape x_2d = x.reshape(-1, x_shape[-1]) y_2d = y.reshape(-1, y_shape[-1]) for lora_down, lora_up, lora_start, lora_size in self.lora_patches: lD = _copy_tensor_to_device(lora_down, x.device) lU = _copy_tensor_to_device(lora_up, x.device) lora_x = F.linear(x_2d.to(lD.dtype), lD) lora_y = F.linear(lora_x, lU).to(y_2d.dtype) if lora_start is not None: y_2d[:, lora_start:lora_start + lora_size] = ( y_2d[:, lora_start:lora_start + lora_size] + lora_y ) else: y_2d = y_2d + lora_y y = y_2d.reshape(y_shape) return y 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}") class INT4LowVramPatch: 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: 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: 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] if getattr(self.module, "_quant_format", "convrot_w4a4") == "int8_tensorwise": scale = self.module.weight_scale if isinstance(scale, torch.Tensor): scale = scale.to(weight.device) weight_float = dequantize_int8(weight, scale) use_convrot = getattr(self.module, "_use_convrot", False) if use_convrot: group_size = getattr(self.module, "_convrot_groupsize", 256) 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) from .int8_quant import quantize_int8, stochastic_round_int8_delta 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) else: weight_float = weight.dequantize() patched_weight_float = comfy.lora.calculate_weight(patches, weight_float, self.key, intermediate_dtype=weight_float.dtype) stochastic_rounding = 1 if self.lora_mode == "Stochastic" else 0 qdata, params = TensorCoreConvRotW4A4Layout.quantize( patched_weight_float, convrot_groupsize=getattr(self.module, "_convrot_groupsize", 256), quant_group_size=getattr(self.module, "_quant_group_size", 64), stochastic_rounding=stochastic_rounding, linear_dtype=getattr(self.module, "_linear_dtype", "int4"), ) return QuantizedTensor(qdata.to(weight.device), "TensorCoreConvRotW4A4Layout", params) class INT4ModelPatcher(comfy.model_patcher.ModelPatcher): 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 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) module_path = key.rsplit('.', 1)[0] try: module = comfy.utils.get_attr(self.model, module_path) except AttributeError: module = None is_int4_module = hasattr(module, "_is_quantized") and module._is_quantized patches = self.patches.get(key, []) if is_int4_module and Int4Ops.Linear._is_bias_key(key): # Only the weight is packed. LightX2V acceleration LoRAs also # contain diff_b patches, so regular floating-point biases must be # handled by ComfyUI's standard patcher. return super().patch_weight_to_device( key, device_to, inplace_update, return_weight, force_cast ) if is_int4_module: if not Int4Ops.dynamic_lora: current_weight = comfy.utils.get_attr(self.model, key) if device_to is None: device_to = current_weight.device 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) if hasattr(current_weight, "to") else current_weight, inplace_update=inplace_update, ) source_weight = current_weight else: source_weight = self.backup[key].weight if getattr(module, "_quant_format", "convrot_w4a4") == "int8_tensorwise": scale = module.weight_scale if isinstance(scale, torch.Tensor): scale = scale.to(device_to) weight_float = dequantize_int8(source_weight.to(device_to), scale) use_convrot = getattr(module, "_use_convrot", False) if use_convrot: group_size = getattr(module, "_convrot_groupsize", 256) 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 patches_list = self.patches.get(key, []) patched_weight_float = comfy.lora.calculate_weight(patches_list, weight_float, key) if use_convrot: patched_weight_float = rotate_weight(patched_weight_float, H, group_size=group_size) from .int8_quant import quantize_int8, stochastic_round_int8_delta if getattr(Int4Ops, "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) 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: weight_float = source_weight.dequantize().to(device_to) patches_list = self.patches.get(key, []) patched_weight_float = comfy.lora.calculate_weight(patches_list, weight_float, key) stochastic_rounding = 1 if getattr(Int4Ops, "lora_mode", "None") == "Stochastic" else 0 qdata, params = TensorCoreConvRotW4A4Layout.quantize( patched_weight_float, convrot_groupsize=getattr(module, "_convrot_groupsize", 256), quant_group_size=getattr(module, "_quant_group_size", 64), stochastic_rounding=stochastic_rounding, linear_dtype=getattr(module, "_linear_dtype", "int4"), ) patched_weight_quant = QuantizedTensor(qdata.to(current_weight.device), "TensorCoreConvRotW4A4Layout", params) if return_weight: return patched_weight_quant if inplace_update: current_weight.data.copy_(patched_weight_quant) else: comfy.utils.set_attr(self.model, key, nn.Parameter(patched_weight_quant, requires_grad=False)) return else: weight = comfy.utils.get_attr(self.model, key) lora_patches = [] for p in patches: strength_patch = p[0] adapter = p[1] strength_model = p[2] offset = p[3] if len(p) > 3 else None if not hasattr(adapter, "weights"): continue strength = strength_patch * strength_model weights = adapter.weights 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 ( getattr(module, "_quant_format", "convrot_w4a4") == "int8_tensorwise" and getattr(module, "_use_convrot", False) ): group_size = getattr(module, "_convrot_groupsize", 256) if down_scaled.shape[1] % group_size == 0: try: from .convrot import build_hadamard, rotate_weight 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 start, size = None, None if offset is not None: _dim, start, size = offset # Dynamic adapter tensors are not registered Parameters, # so Module.to()/partial unload cannot manage their # lifetime. Keep owned pageable CPU storage until load() # packs the complete patch set into stable GPU arenas. down_owned = down_scaled.detach().to(device="cpu").contiguous().clone() up_owned = up.flatten(1).detach().to(device="cpu").contiguous().clone() lora_patches.append((down_owned, up_owned, start, size)) module.lora_patches = lora_patches if return_weight: return weight return return super().patch_weight_to_device(key, device_to, inplace_update, return_weight, force_cast) def _materialize_dynamic_lora_arena(self, device): if not Int4Ops.dynamic_lora or device is None: return device = torch.device(device) if device.type != "cuda": return module_patches = {} records_by_dtype = {} all_ready = True tensor_count = 0 for module in self.model.modules(): patches = getattr(module, "lora_patches", None) if not patches: continue editable = [list(patch) for patch in patches] module_patches[module] = editable for patch_index, patch in enumerate(patches): for tensor_index in (0, 1): tensor = patch[tensor_index] tensor_count += 1 all_ready = all_ready and tensor.device == device records_by_dtype.setdefault(tensor.dtype, []).append( (module, patch_index, tensor_index, tensor) ) if tensor_count == 0 or all_ready: return arenas = [] total_bytes = 0 for dtype, records in records_by_dtype.items(): total_numel = sum(tensor.numel() for _, _, _, tensor in records) host_arena = torch.empty(total_numel, dtype=dtype, device="cpu") offset = 0 layouts = [] for module, patch_index, tensor_index, tensor in records: count = tensor.numel() host_source = _copy_tensor_to_host(tensor) if tensor.device.type == "cuda" else tensor.detach() host_arena[offset:offset + count].copy_(host_source.reshape(-1), non_blocking=False) layouts.append((module, patch_index, tensor_index, offset, count, tuple(tensor.shape))) offset += count device_arena = _copy_tensor_to_device(host_arena, device) arenas.append(device_arena) total_bytes += device_arena.nbytes for module, patch_index, tensor_index, offset, count, shape in layouts: module_patches[module][patch_index][tensor_index] = device_arena[offset:offset + count].view(shape) for module, patches in module_patches.items(): module.lora_patches = [tuple(patch) for patch in patches] self._int4_lora_device_arenas = arenas logging.info( "INT4 Fast: materialized %d Dynamic LoRA tensors in %d stable GPU arena(s), %.2f MiB total.", tensor_count, len(arenas), total_bytes / (1024 * 1024), ) def load(self, *args, **kwargs): self.finalize_pending_int8() save_materialized = bool(getattr(self, "_int4_save_materialized_lora", False)) dynamic_lora_uuid = getattr(self, "patches_uuid", None) dynamic_lora_changed = ( Int4Ops.dynamic_lora and getattr(self, "_int4_dynamic_lora_patches_uuid", None) != dynamic_lora_uuid ) if dynamic_lora_changed: for module in self.model.modules(): if hasattr(module, "lora_patches"): module.lora_patches = [] self._int4_lora_device_arenas = [] if not Int4Ops.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) 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) for name, module in self.model.named_modules(): if hasattr(module, "lora_patches") and module.lora_patches: if not Int4Ops.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 Int4Ops.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"] if not module.lora_patches: 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 = INT4LowVramPatch( weight_key, self.patches, module, getattr(Int4Ops, "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 self._materialize_dynamic_lora_arena(device_to) if Int4Ops.dynamic_lora: self._int4_dynamic_lora_patches_uuid = dynamic_lora_uuid 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 = [] self._int4_lora_device_arenas = [] self._int4_dynamic_lora_patches_uuid = None return super().unpatch_model(device_to, unpatch_weights) def clone(self, *args, **kwargs): src_cls = self.__class__ if src_cls is INT4ModelPatcher: return super().clone(*args, **kwargs) if not issubclass(src_cls, INT4ModelPatcher): name = f"INT4_{src_cls.__name__}" dynamic_cls = type(name, (INT4ModelPatcher, src_cls), {}) else: dynamic_cls = src_cls self.__class__ = dynamic_cls 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) disable_dyn = kwargs.get("disable_dynamic", False) if len(args) > 0: disable_dyn = args[0] if disable_dyn and not issubclass(n.__class__, INT4ModelPatcher): new_cls = type(f"INT4_{n.__class__.__name__}", (INT4ModelPatcher, n.__class__), {}) n.__class__ = new_cls self.__class__ = src_cls return n