523 lines
20 KiB
Python
523 lines
20 KiB
Python
import json
|
|
import logging
|
|
import threading
|
|
from contextlib import contextmanager
|
|
|
|
import torch
|
|
import comfy
|
|
import comfy.ldm.flux.layers
|
|
import comfy.model_detection
|
|
import comfy.sd
|
|
from comfy.model_detection import unet_prefix_from_state_dict, convert_diffusers_mmdit, detect_unet_config
|
|
|
|
|
|
_MODEL_CONFIG_FACTORY_KEY = "comfyui_piflow_model_config_factory"
|
|
_MODEL_DETECTION_LOCK = threading.RLock()
|
|
|
|
_LORA_REQUIRED_PAIRS = (
|
|
(".lora_A.default.weight", ".lora_B.default.weight"),
|
|
(".lora_down.weight", ".lora_up.weight"),
|
|
(".lora_A.weight", ".lora_B.weight"),
|
|
(".lora_A", ".lora_B"),
|
|
)
|
|
_LORA_PRIMARY_SUFFIXES = tuple(suffix for pair in _LORA_REQUIRED_PAIRS for suffix in pair)
|
|
_LORA_COMPANION_SUFFIXES = _LORA_PRIMARY_SUFFIXES + (
|
|
".alpha",
|
|
".dora_scale",
|
|
".lora_A.bias",
|
|
".lora_B.bias",
|
|
".lora_mid.weight",
|
|
".reshape_weight",
|
|
)
|
|
_LORA_NATIVE_PREFIXES = (
|
|
"base_model.",
|
|
"diffusion_model.",
|
|
"lora_transformer_",
|
|
"lora_unet_",
|
|
"lycoris_",
|
|
"transformer.",
|
|
"unet.",
|
|
)
|
|
_QUANTIZATION_POSTFIXES = (
|
|
"scale_input",
|
|
"scale_weight",
|
|
"input_scale",
|
|
"weight_scale",
|
|
"weight_scale_2",
|
|
"comfy_quant",
|
|
)
|
|
|
|
|
|
def flux_norm_target_suffix():
|
|
"""Match the active ComfyUI Flux RMSNorm state-dict name."""
|
|
if getattr(comfy.ldm.flux.layers, "RMSNorm", None) is None:
|
|
return "weight"
|
|
return "scale"
|
|
|
|
|
|
def normalize_flux_norm_keys(state_dict, model_config):
|
|
if model_config.unet_config.get("image_model") not in ("flux", "flux2", "gm_flux", "gm_flux2", "asym_flux2"):
|
|
return
|
|
|
|
target_suffix = flux_norm_target_suffix()
|
|
source_suffix = "scale" if target_suffix == "weight" else "weight"
|
|
source_ending = f".{source_suffix}"
|
|
|
|
for key in list(state_dict.keys()):
|
|
if not key.endswith(source_ending):
|
|
continue
|
|
if ".norm.query_norm." not in key and ".norm.key_norm." not in key:
|
|
continue
|
|
target_key = f"{key[:-len(source_ending)]}.{target_suffix}"
|
|
if target_key not in state_dict:
|
|
state_dict[target_key] = state_dict.pop(key)
|
|
|
|
|
|
def _key_without_suffix(key, suffixes):
|
|
for suffix in suffixes:
|
|
if key.endswith(suffix):
|
|
return key[:-len(suffix)]
|
|
return None
|
|
|
|
|
|
def split_adapter_state_dict(adapter_sd, base_image_model):
|
|
"""Separate complete LoRA groups from tensors merged into the base model."""
|
|
lora_roots = set()
|
|
for key in adapter_sd:
|
|
root = _key_without_suffix(key, _LORA_PRIMARY_SUFFIXES)
|
|
if root is not None:
|
|
lora_roots.add(root)
|
|
|
|
for root in lora_roots:
|
|
matching_pairs = [
|
|
pair for pair in _LORA_REQUIRED_PAIRS
|
|
if root + pair[0] in adapter_sd or root + pair[1] in adapter_sd
|
|
]
|
|
if not any(root + left in adapter_sd and root + right in adapter_sd
|
|
for left, right in matching_pairs):
|
|
raise ValueError("Incomplete LoRA pair for adapter layer: {}".format(root))
|
|
|
|
lora_sd = {}
|
|
full_sd = {}
|
|
|
|
for key, value in adapter_sd.items():
|
|
root = _key_without_suffix(key, _LORA_COMPANION_SUFFIXES)
|
|
if root not in lora_roots:
|
|
full_sd[key] = value
|
|
continue
|
|
|
|
lora_key = key
|
|
if base_image_model in ("flux", "flux2") and not key.startswith(_LORA_NATIVE_PREFIXES):
|
|
lora_key = "transformer." + key
|
|
lora_sd[lora_key] = value
|
|
|
|
return full_sd, lora_sd
|
|
|
|
|
|
def _quantization_layer_from_key(key):
|
|
for postfix in _QUANTIZATION_POSTFIXES:
|
|
suffix = "." + postfix
|
|
if key.endswith(suffix):
|
|
return key[:-len(suffix)], postfix
|
|
return None, None
|
|
|
|
|
|
def _mapped_weight_key(key_mapping, weight_key):
|
|
mapped = key_mapping.get(weight_key)
|
|
if mapped is None:
|
|
return weight_key
|
|
if isinstance(mapped, str):
|
|
return mapped
|
|
return mapped[0]
|
|
|
|
|
|
def _mapped_weight_target(key_mapping, weight_key):
|
|
mapped = key_mapping.get(weight_key)
|
|
if not isinstance(mapped, tuple):
|
|
return None, None
|
|
return mapped[0], mapped[1]
|
|
|
|
|
|
def _quantized_layers(state_dict, metadata):
|
|
layers = set()
|
|
for key in state_dict:
|
|
layer, postfix = _quantization_layer_from_key(key)
|
|
if postfix is not None:
|
|
layers.add(layer)
|
|
|
|
quant_metadata = metadata.get("_quantization_metadata")
|
|
if isinstance(quant_metadata, str):
|
|
quant_metadata = json.loads(quant_metadata)
|
|
if quant_metadata is not None:
|
|
layers.update(quant_metadata.get("layers", {}))
|
|
return layers
|
|
|
|
|
|
def _fully_replaced_mapped_weights(full_sd, key_mapping, base_model_sd):
|
|
mapped_slices = {}
|
|
for key in full_sd:
|
|
target_key, offset = _mapped_weight_target(key_mapping, key)
|
|
if target_key is None or offset is None or target_key not in base_model_sd:
|
|
continue
|
|
dim, start, length = offset
|
|
group = mapped_slices.setdefault((target_key, dim), [])
|
|
group.append((start, start + length))
|
|
|
|
fully_replaced = set()
|
|
for (target_key, dim), intervals in mapped_slices.items():
|
|
cursor = 0
|
|
for start, end in sorted(intervals):
|
|
if start > cursor:
|
|
break
|
|
cursor = max(cursor, end)
|
|
if cursor >= base_model_sd[target_key].shape[dim]:
|
|
fully_replaced.add(target_key)
|
|
return fully_replaced
|
|
|
|
|
|
def merge_model_metadata(base_metadata, adapter_metadata, updated_weight_layers, key_mapping):
|
|
metadata = base_metadata.copy()
|
|
metadata.update({k: v for k, v in adapter_metadata.items() if k != "_quantization_metadata"})
|
|
|
|
base_quant = base_metadata.get("_quantization_metadata")
|
|
adapter_quant = adapter_metadata.get("_quantization_metadata")
|
|
if base_quant is None and adapter_quant is None:
|
|
return metadata
|
|
|
|
if isinstance(base_quant, str):
|
|
base_quant = json.loads(base_quant)
|
|
if isinstance(adapter_quant, str):
|
|
adapter_quant = json.loads(adapter_quant)
|
|
|
|
quant_config = (base_quant or {}).copy()
|
|
quant_config.update({k: v for k, v in (adapter_quant or {}).items() if k != "layers"})
|
|
quant_layers = (base_quant or {}).get("layers", {}).copy()
|
|
for layer in updated_weight_layers:
|
|
quant_layers.pop(layer, None)
|
|
|
|
for layer, config in (adapter_quant or {}).get("layers", {}).items():
|
|
mapped_weight = _mapped_weight_key(key_mapping, layer + ".weight")
|
|
mapped_layer = mapped_weight[:-len(".weight")] if mapped_weight.endswith(".weight") else layer
|
|
quant_layers[mapped_layer] = config
|
|
|
|
if quant_layers:
|
|
quant_config["layers"] = quant_layers
|
|
metadata["_quantization_metadata"] = json.dumps(quant_config)
|
|
else:
|
|
metadata.pop("_quantization_metadata", None)
|
|
return metadata
|
|
|
|
|
|
@contextmanager
|
|
def use_model_config_factory(model_config_factory):
|
|
"""Route one marked ComfyUI load through piFlow's model detection."""
|
|
with _MODEL_DETECTION_LOCK:
|
|
original = comfy.model_detection.model_config_from_unet
|
|
|
|
def model_config_from_unet(
|
|
state_dict, key_prefix, use_base_if_no_match=False, metadata=None):
|
|
if metadata is None or metadata.get(_MODEL_CONFIG_FACTORY_KEY) is not model_config_factory:
|
|
return original(
|
|
state_dict, key_prefix,
|
|
use_base_if_no_match=use_base_if_no_match, metadata=metadata)
|
|
|
|
original_from_config = comfy.model_detection.model_config_from_unet_config
|
|
|
|
def model_config_from_unet_config(
|
|
unet_config, candidate_state_dict=None, unet_key_prefix=""):
|
|
if candidate_state_dict is not state_dict:
|
|
return original_from_config(
|
|
unet_config, candidate_state_dict, unet_key_prefix=unet_key_prefix)
|
|
return model_config_factory(
|
|
candidate_state_dict, unet_key_prefix, metadata=metadata)
|
|
|
|
comfy.model_detection.model_config_from_unet_config = model_config_from_unet_config
|
|
try:
|
|
model_config = original(
|
|
state_dict, key_prefix,
|
|
use_base_if_no_match=use_base_if_no_match, metadata=metadata)
|
|
finally:
|
|
comfy.model_detection.model_config_from_unet_config = original_from_config
|
|
if model_config is not None:
|
|
normalize_flux_norm_keys(state_dict, model_config)
|
|
return model_config
|
|
|
|
comfy.model_detection.model_config_from_unet = model_config_from_unet
|
|
try:
|
|
yield
|
|
finally:
|
|
comfy.model_detection.model_config_from_unet = original
|
|
|
|
|
|
def convert_diffusers_to_comfyui(state_dict, diffusers_weight, comfy_weight_map, cloned_weight_keys=None):
|
|
"""Modified from convert_diffusers_mmdit.
|
|
|
|
This updates state_dict in place. Source tensors are never modified.
|
|
"""
|
|
if cloned_weight_keys is None:
|
|
cloned_weight_keys = set()
|
|
|
|
if isinstance(comfy_weight_map, str):
|
|
comfy_weight_key = comfy_weight_map
|
|
state_dict[comfy_weight_key] = diffusers_weight
|
|
else:
|
|
comfy_weight_key = comfy_weight_map[0]
|
|
if len(comfy_weight_map) > 2:
|
|
weight_convert_fun = comfy_weight_map[2]
|
|
else:
|
|
weight_convert_fun = lambda a: a
|
|
offset = comfy_weight_map[1]
|
|
converted_weight = weight_convert_fun(diffusers_weight)
|
|
if offset is not None:
|
|
updated_weight = state_dict.get(comfy_weight_key, None)
|
|
if updated_weight is None:
|
|
updated_shape = list(diffusers_weight.shape)
|
|
updated_shape[offset[0]] = offset[1] + offset[2]
|
|
updated_weight = torch.empty(
|
|
updated_shape, device=diffusers_weight.device, dtype=diffusers_weight.dtype)
|
|
elif comfy_weight_key not in cloned_weight_keys:
|
|
updated_weight = updated_weight.clone()
|
|
cloned_weight_keys.add(comfy_weight_key)
|
|
if updated_weight.shape[offset[0]] < offset[1] + offset[2]:
|
|
expanded_shape = list(diffusers_weight.shape)
|
|
expanded_shape[offset[0]] = offset[1] + offset[2]
|
|
expanded_weight = torch.empty(
|
|
expanded_shape, device=diffusers_weight.device, dtype=diffusers_weight.dtype)
|
|
_updated_weight = expanded_weight.narrow(offset[0], 0, updated_weight.shape[offset[0]])
|
|
_updated_weight[:] = updated_weight
|
|
updated_weight = expanded_weight
|
|
cloned_weight_keys.add(comfy_weight_key)
|
|
target_slice = updated_weight.narrow(offset[0], offset[1], offset[2])
|
|
target_slice[:] = converted_weight
|
|
else:
|
|
updated_weight = converted_weight
|
|
state_dict[comfy_weight_key] = updated_weight
|
|
return comfy_weight_key
|
|
|
|
|
|
def prepare_base_model_state_dict(base_model_sd, base_metadata=None):
|
|
if base_metadata is None:
|
|
base_metadata = {}
|
|
|
|
diffusion_model_prefix = unet_prefix_from_state_dict(base_model_sd)
|
|
temp_sd = comfy.utils.state_dict_prefix_replace(base_model_sd, {diffusion_model_prefix: ""}, filter_keys=True)
|
|
if len(temp_sd) > 0:
|
|
base_model_sd = temp_sd
|
|
base_unet_config = detect_unet_config(base_model_sd, "", metadata=base_metadata)
|
|
if base_unet_config is None:
|
|
base_model_sd = convert_diffusers_mmdit(base_model_sd, "")
|
|
base_unet_config = detect_unet_config(base_model_sd, "", metadata=base_metadata)
|
|
|
|
return base_model_sd, base_unet_config
|
|
|
|
|
|
def merge_adapter_state_dict(
|
|
base_model_sd, base_unet_config, adapter_sd=None,
|
|
base_metadata=None, adapter_metadata=None):
|
|
if base_metadata is None:
|
|
base_metadata = {}
|
|
if adapter_metadata is None:
|
|
adapter_metadata = {}
|
|
|
|
new_sd = base_model_sd.copy()
|
|
|
|
if adapter_sd is None:
|
|
return new_sd, {}, base_metadata.copy()
|
|
|
|
updated_weight_layers = set()
|
|
updated_keys = set()
|
|
cloned_weight_keys = set()
|
|
|
|
key_mapping = {}
|
|
base_image_model = base_unet_config["image_model"]
|
|
if base_image_model in ("flux", "flux2"):
|
|
key_mapping = comfy.utils.flux_to_diffusers(base_unet_config, output_prefix="")
|
|
|
|
full_sd, lora_sd = split_adapter_state_dict(adapter_sd, base_image_model)
|
|
fully_replaced = _fully_replaced_mapped_weights(full_sd, key_mapping, base_model_sd)
|
|
quantized_layers = _quantized_layers(base_model_sd, base_metadata)
|
|
for key in full_sd:
|
|
target_key, offset = _mapped_weight_target(key_mapping, key)
|
|
if target_key is None or offset is None:
|
|
continue
|
|
target_layer = target_key[:-len(".weight")] if target_key.endswith(".weight") else target_key
|
|
if target_layer in quantized_layers and target_key not in fully_replaced:
|
|
raise ValueError(
|
|
"Cannot partially replace quantized weight {} from adapter key {}. "
|
|
"The adapter must replace every mapped slice of that weight.".format(target_key, key))
|
|
|
|
for target_key in fully_replaced:
|
|
new_sd.pop(target_key, None)
|
|
|
|
for key, value in full_sd.items():
|
|
if key in key_mapping:
|
|
comfy_weight_key = convert_diffusers_to_comfyui(
|
|
new_sd, value, key_mapping[key], cloned_weight_keys=cloned_weight_keys)
|
|
else:
|
|
source_layer, quant_postfix = _quantization_layer_from_key(key)
|
|
if quant_postfix is not None:
|
|
source_weight_key = source_layer + ".weight"
|
|
mapped_weight_key = _mapped_weight_key(key_mapping, source_weight_key)
|
|
if mapped_weight_key != source_weight_key:
|
|
comfy_layer = mapped_weight_key[:-len(".weight")]
|
|
comfy_weight_key = ".".join([comfy_layer, quant_postfix])
|
|
new_sd[comfy_weight_key] = value
|
|
else:
|
|
new_sd[key] = value
|
|
comfy_weight_key = key
|
|
else:
|
|
new_sd[key] = value
|
|
comfy_weight_key = key
|
|
updated_keys.add(comfy_weight_key)
|
|
if comfy_weight_key.endswith(".weight"):
|
|
updated_weight_layers.add(comfy_weight_key[:-len(".weight")])
|
|
|
|
for layer in updated_weight_layers:
|
|
for postfix in _QUANTIZATION_POSTFIXES:
|
|
auxiliary_key = ".".join([layer, postfix])
|
|
if auxiliary_key in new_sd and auxiliary_key not in updated_keys:
|
|
del new_sd[auxiliary_key]
|
|
|
|
metadata = merge_model_metadata(
|
|
base_metadata, adapter_metadata, updated_weight_layers, key_mapping)
|
|
return new_sd, lora_sd, metadata
|
|
|
|
|
|
def build_model_from_state_dict(
|
|
new_sd, metadata, model_options, model_config_factory, disable_dynamic=False):
|
|
metadata = metadata.copy() if metadata is not None else {}
|
|
metadata[_MODEL_CONFIG_FACTORY_KEY] = model_config_factory
|
|
|
|
with use_model_config_factory(model_config_factory):
|
|
return comfy.sd.load_diffusion_model_state_dict(
|
|
new_sd, model_options=model_options, metadata=metadata,
|
|
disable_dynamic=disable_dynamic)
|
|
|
|
|
|
def load_lakonlab_model_state_dict(
|
|
base_model_sd, adapter_sd=None, model_options=None,
|
|
base_metadata=None, adapter_metadata=None, model_config_factory=None,
|
|
disable_dynamic=False):
|
|
if model_options is None:
|
|
model_options = {}
|
|
if base_metadata is None:
|
|
base_metadata = {}
|
|
if adapter_metadata is None:
|
|
adapter_metadata = {}
|
|
|
|
base_model_sd, base_unet_config = prepare_base_model_state_dict(base_model_sd, base_metadata)
|
|
if base_unet_config is None:
|
|
return None, None
|
|
|
|
new_sd, lora_sd, metadata = merge_adapter_state_dict(
|
|
base_model_sd, base_unet_config, adapter_sd,
|
|
base_metadata=base_metadata, adapter_metadata=adapter_metadata)
|
|
|
|
model = build_model_from_state_dict(
|
|
new_sd, metadata, model_options, model_config_factory,
|
|
disable_dynamic=disable_dynamic)
|
|
if model is None:
|
|
return None, None
|
|
return model, lora_sd
|
|
|
|
|
|
def load_lakonlab_model_from_files(
|
|
base_model_path, adapter_path, model_options=None, adapter_strength=1.0,
|
|
model_config_factory=None, error_label="model", disable_dynamic=False):
|
|
if model_options is None:
|
|
model_options = {}
|
|
|
|
base_model_sd, base_metadata = comfy.utils.load_torch_file(base_model_path, return_metadata=True)
|
|
adapter_sd = adapter_metadata = None
|
|
if adapter_path is not None:
|
|
adapter_sd, adapter_metadata = comfy.utils.load_torch_file(adapter_path, return_metadata=True)
|
|
|
|
model, lora_sd = load_lakonlab_model_state_dict(
|
|
base_model_sd, adapter_sd=adapter_sd, model_options=model_options,
|
|
base_metadata=base_metadata, adapter_metadata=adapter_metadata,
|
|
model_config_factory=model_config_factory, disable_dynamic=disable_dynamic)
|
|
if model is None:
|
|
logging.error("ERROR UNSUPPORTED %s MODEL", error_label.upper())
|
|
raise RuntimeError("ERROR: Could not detect {} model type of: {}\n".format(error_label, base_model_path))
|
|
if len(lora_sd) > 0:
|
|
model, _ = comfy.sd.load_lora_for_models(model, None, lora_sd, adapter_strength, None)
|
|
model.cached_patcher_init = (
|
|
load_lakonlab_model_from_files,
|
|
(base_model_path, adapter_path, model_options, adapter_strength,
|
|
model_config_factory, error_label),
|
|
)
|
|
return model
|
|
|
|
|
|
def set_gguf_linear_dtypes(ops, dequant_dtype=None, patch_dtype=None):
|
|
if dequant_dtype in ("default", None):
|
|
ops.Linear.dequant_dtype = None
|
|
elif dequant_dtype in ["target"]:
|
|
ops.Linear.dequant_dtype = dequant_dtype
|
|
else:
|
|
ops.Linear.dequant_dtype = getattr(torch, dequant_dtype)
|
|
|
|
if patch_dtype in ("default", None):
|
|
ops.Linear.patch_dtype = None
|
|
elif patch_dtype in ["target"]:
|
|
ops.Linear.patch_dtype = patch_dtype
|
|
else:
|
|
ops.Linear.patch_dtype = getattr(torch, patch_dtype)
|
|
|
|
|
|
def load_lakonlab_model_from_gguf(
|
|
base_model_path, adapter_path, model_options=None, adapter_strength=1.0,
|
|
dequant_dtype=None, patch_dtype=None, patch_on_device=None,
|
|
model_config_factory=None, error_label="model",
|
|
gguf_model_patcher=None, gguf_sd_loader=None, ggml_ops_class=None,
|
|
disable_dynamic=False):
|
|
if gguf_model_patcher is None:
|
|
raise RuntimeError(
|
|
"ComfyUI-GGUF not found. Please install the ComfyUI-GGUF custom nodes to enable GGUF loading.")
|
|
if model_options is None:
|
|
model_options = {}
|
|
reload_model_options = model_options.copy()
|
|
|
|
ops = ggml_ops_class()
|
|
set_gguf_linear_dtypes(ops, dequant_dtype=dequant_dtype, patch_dtype=patch_dtype)
|
|
|
|
loaded_gguf_data = gguf_sd_loader(base_model_path)
|
|
if isinstance(loaded_gguf_data, tuple):
|
|
base_model_sd, extra = loaded_gguf_data
|
|
base_metadata = extra.get("metadata", None)
|
|
else:
|
|
base_model_sd = loaded_gguf_data
|
|
base_metadata = None
|
|
|
|
model_options = reload_model_options.copy()
|
|
model_options.update(custom_operations=ops)
|
|
|
|
adapter_sd = adapter_metadata = None
|
|
if adapter_path is not None:
|
|
adapter_sd, adapter_metadata = comfy.utils.load_torch_file(adapter_path, return_metadata=True)
|
|
|
|
model, lora_sd = load_lakonlab_model_state_dict(
|
|
base_model_sd, adapter_sd=adapter_sd, model_options=model_options,
|
|
base_metadata=base_metadata, adapter_metadata=adapter_metadata,
|
|
model_config_factory=model_config_factory,
|
|
disable_dynamic=disable_dynamic)
|
|
if model is None:
|
|
logging.error("ERROR UNSUPPORTED %s MODEL", error_label.upper())
|
|
raise RuntimeError("ERROR: Could not detect {} model type of: {}\n".format(error_label, base_model_path))
|
|
|
|
model = gguf_model_patcher.clone(model)
|
|
model.patch_on_device = patch_on_device
|
|
|
|
if len(lora_sd) > 0:
|
|
model, _ = comfy.sd.load_lora_for_models(model, None, lora_sd, adapter_strength, None)
|
|
|
|
model.cached_patcher_init = (
|
|
load_lakonlab_model_from_gguf,
|
|
(base_model_path, adapter_path, reload_model_options, adapter_strength,
|
|
dequant_dtype, patch_dtype, patch_on_device,
|
|
model_config_factory, error_label,
|
|
gguf_model_patcher, gguf_sd_loader, ggml_ops_class),
|
|
)
|
|
|
|
return model
|