import folder_paths import comfy.sd import comfy.model_patcher import comfy.model_management import json import os import logging import torch from comfy.cli_args import args def _is_dynamic_lora_enabled(): try: from .int4_quant import Int4Ops return bool(getattr(Int4Ops, "dynamic_lora", False)) except Exception: return False def _resolve_source_metadata(model): seen = set() def _walk(m): if m is None or id(m) in seen: return None seen.add(id(m)) meta = getattr(m, "_safetensors_metadata", None) if isinstance(meta, dict) and meta: return meta inner = getattr(m, "model", None) if inner is not None: inner_meta = getattr(inner, "_int4_source_metadata", None) if not inner_meta: inner_meta = getattr(inner, "_int8_source_metadata", None) if isinstance(inner_meta, dict) and inner_meta: return inner_meta parent = getattr(m, "parent", None) return _walk(parent) return _walk(model) class INT4ModelSave: def __init__(self): self.output_dir = folder_paths.get_output_directory() @classmethod def INPUT_TYPES(s): return {"required": { "model": ("MODEL",), "filename_prefix": ("STRING", {"default": "int4_models/INT4_Model"}),}, "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},} RETURN_TYPES = () FUNCTION = "save" OUTPUT_NODE = True CATEGORY = "loaders" def save(self, model, filename_prefix, prompt=None, extra_pnginfo=None): full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) metadata = {} src_meta = _resolve_source_metadata(model) if isinstance(src_meta, dict): metadata.update(src_meta) if not src_meta: logging.warning( "INT4 Save: source safetensors metadata could not be located on the patcher chain. " "The output checkpoint will be saved without int4_quantized/int4_model_type/config metadata." ) output_checkpoint = f"{filename}_{counter:05}_.safetensors" output_checkpoint = os.path.join(full_output_folder, output_checkpoint) extra_keys = {} patched_modules = [] patched_module_ids = set() def mark_module_for_direct_save(module): module_id = id(module) if module_id in patched_module_ids: return had_flag = hasattr(module, "comfy_patched_weights") old_flag = getattr(module, "comfy_patched_weights", False) patched_modules.append((module, had_flag, old_flag)) patched_module_ids.add(module_id) module.comfy_patched_weights = True def module_has_quantized_param(module): from comfy.quant_ops import QuantizedTensor for attr in ("weight", "bias"): tensor = getattr(module, attr, None) if isinstance(tensor, QuantizedTensor): return True if isinstance(tensor, torch.Tensor) and tensor.dtype == torch.int8: return True return False def iter_model_modules(model_patcher): if hasattr(model_patcher, "model") and hasattr(model_patcher.model, "named_modules"): yield from model_patcher.model.named_modules() def materialize_int4_lora_patches(model_patcher): if _is_dynamic_lora_enabled() or not hasattr(model_patcher, "patch_weight_to_device"): return patches = getattr(model_patcher, "patches", None) if not patches: return load_device = getattr(model_patcher, "load_device", None) materialized = 0 for name, module in iter_model_modules(model_patcher): if not getattr(module, "_is_quantized", False): continue weight_key = name + ".weight" if name else "weight" if weight_key not in patches: continue try: current_weight = getattr(module, "weight", None) device_to = load_device if load_device is not None else getattr(current_weight, "device", None) model_patcher.patch_weight_to_device(weight_key, device_to=device_to) if hasattr(module, "weight_lowvram_function"): module.weight_lowvram_function = None materialized += 1 except Exception as e: logging.warning( f"INT4 Save: failed to materialize LoRA patch for {weight_key}: {e}." ) if materialized > 0: logging.info(f"INT4 Save: materialized {materialized} INT4 LoRA patched weight(s) before saving.") finalize_fn = getattr(model, "finalize_pending_int8", None) if finalize_fn is not None: finalize_fn() try: comfy.model_management.load_models_gpu([model], force_full_load=True) except Exception as e: logging.warning( f"INT4 Save: full-load pre-pass failed ({e}); falling back to default load_models_gpu." ) try: comfy.model_management.load_models_gpu([model]) except Exception as e2: logging.warning( f"INT4 Save: load_models_gpu fallback also failed ({e2}); continuing best-effort." ) if finalize_fn is not None: finalize_fn() materialize_int4_lora_patches(model) if hasattr(model, "model"): for name, module in iter_model_modules(model): if module_has_quantized_param(module): mark_module_for_direct_save(module) if getattr(module, "_is_quantized", False): quant_conf = { "format": "convrot_w4a4", "convrot_groupsize": int(getattr(module, "_convrot_groupsize", 256)), "quant_group_size": int(getattr(module, "_quant_group_size", 64)), "linear_dtype": getattr(module, "_linear_dtype", "int4"), } prefix = "model." + name + "." if name else "model." extra_keys[prefix + "comfy_quant"] = torch.tensor( list(json.dumps(quant_conf).encode('utf-8')), dtype=torch.uint8 ) mark_module_for_direct_save(module) original_lazy_new = comfy.model_patcher.LazyCastingParam.__new__ original_lazy_piece_new = comfy.model_patcher.LazyCastingParamPiece.__new__ def lazy_casting_param_new(cls, model, key, tensor): requires_grad = tensor.is_floating_point() or tensor.is_complex() return torch.nn.Parameter.__new__(cls, tensor, requires_grad=requires_grad) def lazy_casting_param_piece_new(cls, caster, state_dict_key, tensor): requires_grad = tensor.is_floating_point() or tensor.is_complex() return torch.nn.Parameter.__new__(cls, tensor, requires_grad=requires_grad) original_state_dict_for_saving = model.state_dict_for_saving def custom_state_dict_for_saving(*args, **kwargs): from comfy.quant_ops import QuantizedTensor sd = original_state_dict_for_saving(*args, **kwargs) new_sd = {} for k, v in sd.items(): if isinstance(v, QuantizedTensor) or (hasattr(v, "_qdata") and hasattr(v, "_params")): new_sd[k] = v._qdata.cpu() scale_key = k.rsplit('.', 1)[0] + ".weight_scale" scale_val = v._params.scale if isinstance(scale_val, torch.Tensor): scale_val = scale_val.reshape(-1).cpu() new_sd[scale_key] = scale_val elif hasattr(v, "weight") and (isinstance(v.weight, QuantizedTensor) or (hasattr(v.weight, "_qdata") and hasattr(v.weight, "_params"))): new_sd[k] = v.weight._qdata.cpu() scale_key = k.rsplit('.', 1)[0] + ".weight_scale" scale_val = v.weight._params.scale if isinstance(scale_val, torch.Tensor): scale_val = scale_val.reshape(-1).cpu() new_sd[scale_key] = scale_val else: new_sd[k] = v return new_sd had_save_flag = hasattr(model, "_int4_save_materialized_lora") old_save_flag = getattr(model, "_int4_save_materialized_lora", False) try: model._int4_save_materialized_lora = True comfy.model_patcher.LazyCastingParam.__new__ = staticmethod(lazy_casting_param_new) comfy.model_patcher.LazyCastingParamPiece.__new__ = staticmethod(lazy_casting_param_piece_new) model.state_dict_for_saving = custom_state_dict_for_saving comfy.sd.save_checkpoint(output_checkpoint, model, metadata=metadata, extra_keys=extra_keys) finally: model.state_dict_for_saving = original_state_dict_for_saving comfy.model_patcher.LazyCastingParam.__new__ = original_lazy_new comfy.model_patcher.LazyCastingParamPiece.__new__ = original_lazy_piece_new if had_save_flag: model._int4_save_materialized_lora = old_save_flag else: try: delattr(model, "_int4_save_materialized_lora") except AttributeError: pass for module, had_flag, old_flag in patched_modules: if had_flag: module.comfy_patched_weights = old_flag else: try: delattr(module, "comfy_patched_weights") except AttributeError: pass return {}