311 lines
12 KiB
Python
311 lines
12 KiB
Python
import logging
|
|
|
|
import torch
|
|
import comfy
|
|
import comfy.ldm.flux.layers
|
|
from comfy import model_management
|
|
from comfy.model_detection import unet_prefix_from_state_dict, convert_diffusers_mmdit, detect_unet_config
|
|
|
|
try:
|
|
from comfy.utils import convert_old_quants
|
|
except ImportError:
|
|
convert_old_quants = None
|
|
|
|
|
|
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 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, adapter_metadata=None):
|
|
metadata = {}
|
|
if adapter_metadata is None:
|
|
adapter_metadata = {}
|
|
|
|
new_sd = base_model_sd.copy()
|
|
lora_sd = {}
|
|
|
|
if adapter_sd is None:
|
|
return new_sd, lora_sd, metadata
|
|
|
|
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="")
|
|
|
|
for k in adapter_sd.keys():
|
|
if "lora" in k:
|
|
if base_image_model in ("flux", "flux2") and not k.startswith("transformer."):
|
|
lora_sd["transformer." + k] = adapter_sd[k]
|
|
else:
|
|
lora_sd[k] = adapter_sd[k]
|
|
else:
|
|
if k in key_mapping:
|
|
comfy_weight_key = convert_diffusers_to_comfyui(
|
|
new_sd, adapter_sd[k], key_mapping[k], cloned_weight_keys=cloned_weight_keys)
|
|
else:
|
|
new_sd[k] = adapter_sd[k]
|
|
comfy_weight_key = k
|
|
updated_keys.add(comfy_weight_key)
|
|
if comfy_weight_key.endswith(".weight"):
|
|
updated_weight_layers.add(comfy_weight_key[:-7])
|
|
|
|
for layer in updated_weight_layers:
|
|
for scale_postfix in ["scale_input", "scale_weight", "input_scale", "weight_scale"]:
|
|
scale_key = ".".join([layer, scale_postfix])
|
|
if scale_key in new_sd and scale_key not in updated_keys:
|
|
del new_sd[scale_key]
|
|
|
|
metadata.update(adapter_metadata)
|
|
return new_sd, lora_sd, metadata
|
|
|
|
|
|
def build_model_from_state_dict(new_sd, metadata, weight_dtype, model_options, model_config_factory):
|
|
parameters = comfy.utils.calculate_parameters(new_sd)
|
|
|
|
if convert_old_quants is not None:
|
|
if model_options.get("custom_operations", None) is None:
|
|
new_sd, metadata = convert_old_quants(new_sd, "", metadata=metadata)
|
|
|
|
model_config = model_config_factory(new_sd, "", metadata=metadata)
|
|
if model_config is None:
|
|
return None
|
|
normalize_flux_norm_keys(new_sd, model_config)
|
|
|
|
offload_device = model_management.unet_offload_device()
|
|
unet_weight_dtype = list(model_config.supported_inference_dtypes)
|
|
if hasattr(model_config, "quant_config"):
|
|
has_quant_config = model_config.quant_config is not None
|
|
if has_quant_config:
|
|
weight_dtype = None
|
|
else:
|
|
has_quant_config = getattr(model_config, "layer_quant_config", None) is not None
|
|
if getattr(model_config, "scaled_fp8", None) is not None:
|
|
weight_dtype = None
|
|
|
|
dtype = model_options.get("dtype", None)
|
|
if dtype is None:
|
|
unet_dtype = model_management.unet_dtype(
|
|
model_params=parameters, supported_dtypes=unet_weight_dtype, weight_dtype=weight_dtype)
|
|
else:
|
|
unet_dtype = dtype
|
|
|
|
load_device = model_management.get_torch_device()
|
|
if has_quant_config:
|
|
manual_cast_dtype = model_management.unet_manual_cast(
|
|
None, load_device, model_config.supported_inference_dtypes)
|
|
else:
|
|
manual_cast_dtype = model_management.unet_manual_cast(
|
|
unet_dtype, load_device, model_config.supported_inference_dtypes)
|
|
model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
|
|
model_config.custom_operations = model_options.get("custom_operations", model_config.custom_operations)
|
|
if model_options.get("fp8_optimizations", False):
|
|
model_config.optimizations["fp8"] = True
|
|
|
|
model = model_config.get_model(new_sd, "")
|
|
model = model.to(offload_device)
|
|
model.load_model_weights(new_sd, "")
|
|
left_over = new_sd.keys()
|
|
if len(left_over) > 0:
|
|
logging.info("left over keys in diffusion model: {}".format(left_over))
|
|
return comfy.model_patcher.ModelPatcher(
|
|
model, load_device=load_device, offload_device=offload_device)
|
|
|
|
|
|
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):
|
|
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
|
|
|
|
weight_dtype = comfy.utils.weight_dtype(base_model_sd)
|
|
metadata = base_metadata.copy()
|
|
new_sd, lora_sd, adapter_metadata = merge_adapter_state_dict(
|
|
base_model_sd, base_unet_config, adapter_sd, adapter_metadata)
|
|
metadata.update(adapter_metadata)
|
|
|
|
model = build_model_from_state_dict(
|
|
new_sd, metadata, weight_dtype, model_options, model_config_factory)
|
|
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"):
|
|
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)
|
|
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)
|
|
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):
|
|
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 = {}
|
|
|
|
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 = 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)
|
|
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)
|
|
|
|
return model
|