l33chKing
8e772b2283
Bug 1 — Save after INT8 Grouped LoRA produces corrupted/unloadable checkpoint Root cause: ModelPatcher.clone() constructs a fresh patcher and does not carry over arbitrary attributes that the loader had attached to the source patcher (_safetensors_metadata). After INT8GroupedLora.apply_loras() calls model.clone(), the metadata stash is gone, so INT8ModelSave writes a checkpoint with empty safetensors metadata. Some model types (notably LTX2 — which keeps its transformer config in metadata) then fail to redetect on reload, and the file appears smaller because the metadata header is empty. Fixes: int8_unet_loader.py:225-235 — Stash metadata in two places: on the patcher (legacy) and on the inner model.model object (_int8_source_metadata). Inner-model attributes survive clone() because it shares the inner model by reference. int8_lora.py:31-45 — Best-effort copy of _safetensors_metadata / _int8_source_metadata from the source patcher onto the cloned one in INT8GroupedLora.apply_loras. int8_save.py:12-40 — New _resolve_source_metadata() that walks the patcher → parent chain and the inner-model fallback, plus a warning when nothing is found. Bug 2 — Save node was collecting extra_keys BEFORE LoRA bake comfy.sd.save_checkpoint internally triggers load_models_gpu, which is when the INT8ModelPatcher actually bakes LoRA patches into the int8 weights. The previous save node iterated modules before that, so on the LoRA-stacked path: mark_module_for_direct_save was being called against the pre-bake module identity. comfy_quant extra_keys were emitted against pre-bake state. After save_checkpoint baked patches via set_attr(module, "weight", new_int8_param), the saved bytes were the patched int8 weights but with metadata that may or may not have observed the post-bake state cleanly. Fix (int8_save.py:120-154): Pre-call comfy.model_management.load_models_gpu([model], force_full_load=True) so all LoRA patches are baked and comfy_patched_weights = True is set on every module before we walk it. Then re-finalize any aimdo-deferred layers, then collect extra_keys. Bug 3 — OTF quantization round-trip was fragile OTF-quantized layers were never stamping _convrot_groupsize on the module, so the save side fell back to the loader-time default (currently CONVROT_GROUP_SIZE = 256). It happened to match on reload, but any future change to the default would silently corrupt round-trips, and the saved comfy_quant JSON had no convrot_groupsize field at all. Fixes: int8_quant.py:254-257 and int8_quant.py:443-454 — Set self._convrot_groupsize = CONVROT_GROUP_SIZE whenever OTF (or the deferred aimdo finalize) actually applies a Hadamard rotation. int8_save.py:166-187 — Always emit convrot_groupsize in comfy_quant whenever convrot=True, falling back to the module-level constant if _convrot_groupsize is somehow missing. Also emit per_row so the loader can pick the matching forward kernel without re-inspecting tensor shapes. int8_quant.py:340-396 — Parse the new per_row hint in comfy_quant and use it to override the shape-based heuristic in _load_from_state_dict. Existing files (without per_row) continue to work via the original shape heuristic. Files changed int8_save.py — Restructured save flow, added metadata resolver, save convrot_groupsize and per_row deterministically, pre-load patches before extra_keys collection, robust fallback path. int8_unet_loader.py:225-235 — Stash _int8_source_metadata on the inner model so it survives clone(). int8_lora.py:31-45 — Copy metadata attributes onto the cloned patcher in INT8GroupedLora.apply_loras. int8_quant.py:240-256, 340-352, 371-401, 443-454 — Stamp _convrot_groupsize during OTF and aimdo-deferred finalize, parse per_row hint from comfy_quant. Notes on the OTF "worse quality than pre-quantized" report After re-tracing the OTF → save → reload path carefully (Hadamard rotation, per-row scale, bias preservation, dtype tracking) I was not able to find a save-side bug that would degrade quality. The OTF and pre-quantized paths use the same quantize_int8_axiswise(weight, dim=1) code, the same forward kernels, and now write the same comfy_quant payload. The most plausible source of the perceived quality gap is that pre-quantized HF checkpoints were calibrated with a different excluded_names list per model than what Int8TensorwiseOps.excluded_names uses for the same model in this loader; that is a calibration choice, not a bug in the save/reload path. If a concrete reproduction (same workflow, same seed, OTF-then-save vs HF-prequant) shows a measurable delta, the fix would belong in the per-model excluded_names table in int8_unet_loader.py, not in the save node.
76 lines
2.9 KiB
Python
76 lines
2.9 KiB
Python
import torch
|
|
import folder_paths
|
|
import comfy.utils
|
|
import comfy.lora
|
|
import logging
|
|
|
|
class INT8GroupedLora:
|
|
"""
|
|
A simple node to stack multiple LoRAs onto a model.
|
|
The actual application logic is intercepted by INT8ModelPatcher in int8_quant.py.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
inputs = {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
},
|
|
"optional": {}
|
|
}
|
|
lora_list = ["None"] + folder_paths.get_filename_list("loras")
|
|
for i in range(1, 11):
|
|
inputs["optional"][f"lora_{i}"] = (lora_list,)
|
|
inputs["optional"][f"strength_{i}"] = ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01})
|
|
return inputs
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "apply_loras"
|
|
CATEGORY = "loaders"
|
|
DESCRIPTION = "Stacks multiple LoRAs onto an INT8 model. Actual patching is handled by INT8ModelPatcher."
|
|
|
|
def apply_loras(self, model, **kwargs):
|
|
model_patcher = model.clone()
|
|
|
|
# ComfyUI's ModelPatcher.clone() builds a fresh patcher object and does
|
|
# NOT carry over arbitrary attributes set on the source patcher (e.g.
|
|
# the _safetensors_metadata stash from UNetLoaderINTW8A8). Without this,
|
|
# downstream nodes like INT8ModelSave lose the source safetensors
|
|
# metadata (int8_quantized flag, model_type, LTX2 'config', etc.) and
|
|
# produce a corrupted/unloadable checkpoint.
|
|
for attr in ("_safetensors_metadata", "_int8_source_metadata"):
|
|
if hasattr(model, attr) and not hasattr(model_patcher, attr):
|
|
try:
|
|
setattr(model_patcher, attr, getattr(model, attr))
|
|
except Exception:
|
|
pass
|
|
|
|
# Get key mappings from ComfyUI's framework
|
|
key_map = {}
|
|
if model_patcher.model.model_type.name != "ModelType.CLIP":
|
|
key_map = comfy.lora.model_lora_keys_unet(model_patcher.model, key_map)
|
|
|
|
applied_loras = []
|
|
for i in range(1, 11):
|
|
name = kwargs.get(f"lora_{i}")
|
|
strength = kwargs.get(f"strength_{i}", 0)
|
|
|
|
if name and name != "None" and strength != 0:
|
|
lora_path = folder_paths.get_full_path("loras", name)
|
|
lora_data = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
|
patch_dict = comfy.lora.load_lora(lora_data, key_map)
|
|
model_patcher.add_patches(patch_dict, strength)
|
|
applied_loras.append(name)
|
|
del lora_data
|
|
|
|
if applied_loras:
|
|
logging.info(f"INT8 Grouped LoRA: Stacked {len(applied_loras)} LoRAs: {', '.join(applied_loras)}")
|
|
|
|
return (model_patcher,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"INT8GroupedLora": INT8GroupedLora,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"INT8GroupedLora": "INT8 Grouped LoRA",
|
|
} |