137 lines
4.5 KiB
Python
137 lines
4.5 KiB
Python
import torch
|
||
import comfy.utils
|
||
import comfy.sd
|
||
import folder_paths
|
||
from nodes import LoraLoader
|
||
|
||
#代码参考了：https://github.com/PGCRT/CRT-Nodes/blob/main/py/LoraLoaderZImage.py
|
||
#由GPT结合官方LoRA Model-Only Loader节点进行改写和优化。
|
||
|
||
# ---------------------------------------------------------------------
|
||
# Z-Image LoRA Model-Only Loader
|
||
# ---------------------------------------------------------------------
|
||
class ZImageLoraModelOnly(LoraLoader):
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"model": ("MODEL",),
|
||
"lora_name": (folder_paths.get_filename_list("loras"), ),
|
||
"strength_model": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": -100.0,
|
||
"max": 100.0,
|
||
"step": 0.01
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("MODEL",)
|
||
FUNCTION = "load_zimage_lora_model_only"
|
||
CATEGORY = "loaders"
|
||
DESCRIPTION = "Loads Z-Image LoRA (QKV fused) and applies it to MODEL only (no CLIP)."
|
||
|
||
# -------------------------------------------------------
|
||
# Z-IMAGE QKV FUSION
|
||
# -------------------------------------------------------
|
||
def _convert_zimage_lora(self, lora_dict):
|
||
new_lora = {}
|
||
qkv_groups = {}
|
||
|
||
for k, v in lora_dict.items():
|
||
new_k = k
|
||
|
||
# 修正 to_out.0 → out
|
||
if ".attention.to_out.0." in new_k:
|
||
new_k = new_k.replace(".attention.to_out.0.", ".attention.out.")
|
||
new_lora[new_k] = v
|
||
continue
|
||
|
||
# Q/K/V 三分结构
|
||
if ".attention.to_" in new_k:
|
||
parts = new_k.split(".attention.to_")
|
||
base_prefix = parts[0] + ".attention"
|
||
remainder = parts[1] # 例如：q.lora_A.weight
|
||
|
||
qkv_type = remainder[0] # q / k / v
|
||
suffix = remainder[2:] # lora_A.weight
|
||
|
||
if base_prefix not in qkv_groups:
|
||
qkv_groups[base_prefix] = {"q": {}, "k": {}, "v": {}}
|
||
|
||
qkv_groups[base_prefix][qkv_type][suffix] = v
|
||
continue
|
||
|
||
# 其他权重直接 passthrough
|
||
new_lora[new_k] = v
|
||
|
||
# ---- 开始 Fusion ----
|
||
for base_key, group in qkv_groups.items():
|
||
|
||
# A 权重
|
||
ak = "lora_A.weight"
|
||
if ak in group["q"] and ak in group["k"] and ak in group["v"]:
|
||
fused_A = torch.cat(
|
||
[group["q"][ak], group["k"][ak], group["v"][ak]], dim=0
|
||
)
|
||
new_lora[f"{base_key}.qkv.lora_A.weight"] = fused_A
|
||
|
||
# B 权重
|
||
bk = "lora_B.weight"
|
||
if bk in group["q"] and bk in group["k"] and bk in group["v"]:
|
||
q_b = group["q"][bk]
|
||
k_b = group["k"][bk]
|
||
v_b = group["v"][bk]
|
||
|
||
out_dim, rank = q_b.shape
|
||
fused_B = torch.zeros((out_dim * 3, rank * 3),
|
||
dtype=q_b.dtype, device=q_b.device)
|
||
|
||
fused_B[0:out_dim, 0:rank] = q_b
|
||
fused_B[out_dim:2*out_dim, rank:2*rank] = k_b
|
||
fused_B[2*out_dim:3*out_dim, 2*rank:3*rank] = v_b
|
||
|
||
new_lora[f"{base_key}.qkv.lora_B.weight"] = fused_B
|
||
|
||
# alpha
|
||
if "lora_alpha" in group["q"]:
|
||
new_lora[f"{base_key}.qkv.lora_alpha"] = group["q"]["lora_alpha"]
|
||
|
||
return new_lora
|
||
|
||
# -------------------------------------------------------
|
||
# 主函数
|
||
# -------------------------------------------------------
|
||
def load_zimage_lora_model_only(self, model, lora_name, strength_model):
|
||
|
||
# load .safetensors / .pt
|
||
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
||
raw_lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
||
|
||
# 转换为 Z-Image 格式
|
||
fixed_lora = self._convert_zimage_lora(raw_lora)
|
||
|
||
# 注入 LoRA（clip=None）
|
||
new_model, _ = comfy.sd.load_lora_for_models(
|
||
model,
|
||
None,
|
||
fixed_lora,
|
||
strength_model,
|
||
0 # clip 强度 = 0
|
||
)
|
||
|
||
return (new_model,)
|
||
|
||
|
||
# ---------------------------------------------------------------------
|
||
# 节点注册
|
||
# ---------------------------------------------------------------------
|
||
NODE_CLASS_MAPPINGS = {
|
||
"ZImageLoraModelOnly": ZImageLoraModelOnly,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"ZImageLoraModelOnly": "ZImageLoraModelOnly☀",
|
||
}
|