248 lines
11 KiB
Python
248 lines
11 KiB
Python
import os
|
|
import torch
|
|
import folder_paths
|
|
import comfy.sd
|
|
import comfy.utils
|
|
import comfy.model_detection
|
|
import comfy.lora
|
|
import comfy.lora_convert
|
|
import comfy.memory_management
|
|
import comfy.model_management
|
|
import logging
|
|
|
|
from .int4_quant import Int4Ops, INT4ModelPatcher
|
|
|
|
|
|
def _load_int4_unet_cached_patcher(unet_name, weight_dtype, model_type, on_the_fly_quantization, enable_convrot, lora_mode, pre_lora=None, disable_dynamic=False):
|
|
return UNetLoaderINTW4A4().load_unet(
|
|
unet_name,
|
|
weight_dtype,
|
|
model_type,
|
|
on_the_fly_quantization,
|
|
enable_convrot=enable_convrot,
|
|
lora_mode=lora_mode,
|
|
pre_lora=pre_lora,
|
|
disable_dynamic=disable_dynamic,
|
|
)[0]
|
|
|
|
|
|
class UNetLoaderINTW4A4:
|
|
"""
|
|
Load INT4 convrot quantized diffusion models.
|
|
|
|
Uses Int4Ops for direct INT4 loading and comfy-kitchen execution.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"unet_name": (folder_paths.get_filename_list("diffusion_models"),),
|
|
"weight_dtype": (["default", "fp16", "bf16", "fp32"], {"tooltip": "Compute dtype. Default follows the model dtype."}),
|
|
"model_type": (["flux2", "z-image", "ideogram4", "chroma", "krea2", "wan", "ltx2", "qwen", "ernie", "anima", "hidream o1", "boogu"], {"tooltip": "Only used for on the fly quantization, to filter sensitive layers."}),
|
|
"on_the_fly_quantization": ("BOOLEAN", {"default": False, "tooltip": "Quantize a higher precision model to INT4. If the selected model is already INT4 keep unchecked."}),
|
|
"enable_convrot": ("BOOLEAN", {"default": True, "tooltip": "Enable ConvRot for better quantization."}),
|
|
"lora_mode": (["None", "Stochastic", "Dynamic"], {"default": "None", "tooltip": "None bakes LoRA patches with normal rounding which is the default behavior. Stochastic bakes with stochastic rounding. Dynamic applies LoRA at inference time."}),
|
|
},
|
|
"optional": {
|
|
"pre_lora": ("PRE_LORA",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "load_unet"
|
|
CATEGORY = "loaders"
|
|
DESCRIPTION = "Load and Quantize INT4 models with fast comfy-kitchen GPU execution."
|
|
|
|
def load_unet(self, unet_name, weight_dtype, model_type, on_the_fly_quantization, enable_convrot=True, lora_mode="None", pre_lora=None, disable_dynamic=False):
|
|
unet_path = folder_paths.get_full_path("diffusion_models", unet_name)
|
|
|
|
if isinstance(lora_mode, bool):
|
|
lora_mode = "Dynamic" if lora_mode else "None"
|
|
lora_mode = str(lora_mode)
|
|
if lora_mode not in {"None", "Stochastic", "Dynamic"}:
|
|
lora_mode = "None"
|
|
|
|
if pre_lora is not None:
|
|
loras_to_load = pre_lora if isinstance(pre_lora, list) else [pre_lora]
|
|
else:
|
|
loras_to_load = []
|
|
|
|
model_options = {"custom_operations": Int4Ops}
|
|
|
|
Int4Ops.excluded_names = []
|
|
Int4Ops.dynamic_quantize = on_the_fly_quantization
|
|
Int4Ops.enable_convrot = enable_convrot
|
|
Int4Ops.convrot_groupsize = 256
|
|
Int4Ops.quant_group_size = 64
|
|
Int4Ops.linear_dtype = "int4"
|
|
Int4Ops._is_prequantized = False
|
|
Int4Ops.lora_mode = lora_mode
|
|
Int4Ops.dynamic_lora = lora_mode == "Dynamic"
|
|
Int4Ops.dynamic_load_device = None
|
|
|
|
dtype_map = {
|
|
"fp16": torch.float16,
|
|
"bf16": torch.bfloat16,
|
|
"fp32": torch.float32,
|
|
}
|
|
Int4Ops.compute_dtype = dtype_map.get(str(weight_dtype), None)
|
|
if comfy.memory_management.aimdo_enabled and (on_the_fly_quantization or len(loras_to_load) > 0):
|
|
Int4Ops.dynamic_load_device = comfy.model_management.get_torch_device()
|
|
logging.info(f"INT4 Fast: Aimdo dynamic loading active, using {Int4Ops.dynamic_load_device} as work device.")
|
|
if hasattr(Int4Ops, "_logged_otf"):
|
|
delattr(Int4Ops, "_logged_otf")
|
|
|
|
if model_type == "flux2":
|
|
Int4Ops.excluded_names = [
|
|
'img_in', 'time_in', 'guidance_in', 'txt_in',
|
|
'double_stream_modulation_img', 'double_stream_modulation_txt',
|
|
'single_stream_modulation',
|
|
]
|
|
elif model_type == "z-image":
|
|
Int4Ops.excluded_names = [
|
|
'cap_embedder', 't_embedder', 'x_embedder', 'cap_pad_token', 'context_refiner',
|
|
'final_layer', 'noise_refiner', 'adaLN',
|
|
'x_pad_token', 'layers.0.',
|
|
]
|
|
elif model_type == "chroma":
|
|
Int4Ops.excluded_names = [
|
|
'distilled_guidance_layer', 'final_layer', 'img_in', 'txt_in', 'nerf_image_embedder',
|
|
'nerf_blocks', 'nerf_final_layer_conv', '__x0__', 'nerf_final_layer_conv',
|
|
]
|
|
elif model_type == "qwen":
|
|
Int4Ops.excluded_names = [
|
|
'time_text_embed', 'img_in', 'norm_out', 'proj_out', 'txt_in'
|
|
]
|
|
elif model_type == "ernie":
|
|
Int4Ops.excluded_names = [
|
|
'time', 'x_embedder', 'text_proj', 'adaLN',
|
|
]
|
|
elif model_type == "anima":
|
|
Int4Ops.excluded_names = [
|
|
'embed', 'llm', 'adaln',
|
|
]
|
|
elif model_type == "krea2":
|
|
Int4Ops.excluded_names = [
|
|
'first', 'last', 'tmlp', 'tproj', 'txtfusion', 'txtmlp'
|
|
]
|
|
elif model_type == "hidream o1":
|
|
Int4Ops.excluded_names = [
|
|
'embed', 'language_model.layers.35.mlp',
|
|
]
|
|
elif model_type == "boogu":
|
|
Int4Ops.excluded_names = [
|
|
'embed', 'refine', 'norm_out',
|
|
]
|
|
elif model_type == "ideogram4":
|
|
Int4Ops.excluded_names = [
|
|
'embed_image_indicator', 't_embedding', 'proj',
|
|
]
|
|
elif model_type == "wan":
|
|
Int4Ops.excluded_names = [
|
|
'patch_embedding', 'text_embedding', 'time_embedding', 'time_projection', 'head',
|
|
'img_emb', 'face_adapter', 'face_encoder', 'motion_encoder', 'pose_patch_embedding',
|
|
]
|
|
elif model_type == "ltx2":
|
|
Int4Ops.excluded_names = [
|
|
'adaln', 'embedding', 'patchify', 'to_gate_logits', 'proj_out', 'model.audio', 'model.video', 'model.av', 'model.patch', 'model.proj', 'shift',
|
|
]
|
|
|
|
sd, metadata = comfy.utils.load_torch_file(unet_path, return_metadata=True)
|
|
|
|
Int4Ops.lora_patches = {}
|
|
if len(loras_to_load) > 0:
|
|
grouped_patches = {}
|
|
for lora in loras_to_load:
|
|
lora_name = lora.get("lora_name", "None")
|
|
lora_strength = lora.get("lora_strength", 1.0)
|
|
|
|
if lora_name == "None":
|
|
continue
|
|
|
|
lora_path = folder_paths.get_full_path("loras", lora_name)
|
|
lora_data = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
|
lora_data = comfy.lora_convert.convert_lora(lora_data)
|
|
|
|
unet_prefix = comfy.model_detection.unet_prefix_from_state_dict(sd)
|
|
m_config = comfy.model_detection.model_config_from_unet(sd, unet_prefix, metadata=metadata)
|
|
|
|
if m_config is None and unet_prefix != "":
|
|
m_config = comfy.model_detection.model_config_from_unet(sd, "", metadata=metadata)
|
|
if m_config is not None:
|
|
unet_prefix = ""
|
|
|
|
if m_config is not None:
|
|
m_config.custom_operations = Int4Ops
|
|
Int4Ops.skeleton_meta_init = True
|
|
try:
|
|
skeleton_model = m_config.get_model(sd, unet_prefix)
|
|
finally:
|
|
Int4Ops.skeleton_meta_init = False
|
|
key_map = comfy.lora.model_lora_keys_unet(skeleton_model, {})
|
|
|
|
patch_dict = comfy.lora.load_lora(lora_data, key_map)
|
|
|
|
def normalize_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
|
|
|
|
for k, v in patch_dict.items():
|
|
target_key = k
|
|
offset = None
|
|
function = None
|
|
if isinstance(k, tuple):
|
|
target_key = k[0]
|
|
if len(k) > 1: offset = k[1]
|
|
if len(k) > 2: function = k[2]
|
|
|
|
nk = normalize_key(target_key)
|
|
if nk not in grouped_patches:
|
|
grouped_patches[nk] = []
|
|
grouped_patches[nk].append((v, offset, function, lora_strength))
|
|
else:
|
|
logging.warning(f"INT4 Fast: Could not detect model type for LoRA mapping.")
|
|
|
|
del lora_data
|
|
|
|
if grouped_patches:
|
|
Int4Ops.lora_patches = grouped_patches
|
|
logging.info(f"INT4 Fast: Prepared {len(grouped_patches)} layer patches for baking.")
|
|
|
|
try:
|
|
Int4Ops.applied_lora_patches = set()
|
|
model = comfy.sd.load_diffusion_model_state_dict(sd, model_options=model_options, metadata=metadata, disable_dynamic=disable_dynamic)
|
|
|
|
if Int4Ops.lora_patches:
|
|
unmatched = set(Int4Ops.lora_patches.keys()) - Int4Ops.applied_lora_patches
|
|
if unmatched:
|
|
print(f"INT4 Fast: {len(unmatched)} LoRA keys were NOT matched:")
|
|
for k in sorted(unmatched):
|
|
print(f" unmatched: {k}")
|
|
finally:
|
|
dynamic_load_device = Int4Ops.dynamic_load_device
|
|
Int4Ops.lora_patches = {}
|
|
Int4Ops.dynamic_load_device = None
|
|
if dynamic_load_device is not None and torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
if hasattr(Int4Ops, 'applied_lora_patches'):
|
|
delattr(Int4Ops, 'applied_lora_patches')
|
|
|
|
model = INT4ModelPatcher.clone(model)
|
|
model.cached_patcher_init = (
|
|
_load_int4_unet_cached_patcher,
|
|
(unet_name, weight_dtype, model_type, on_the_fly_quantization, enable_convrot, lora_mode, pre_lora),
|
|
)
|
|
model._safetensors_metadata = metadata
|
|
try:
|
|
if model.model is not None:
|
|
model.model._int4_source_metadata = metadata
|
|
except Exception:
|
|
pass
|
|
|
|
return (model,)
|