128 lines
4.2 KiB
Python
128 lines
4.2 KiB
Python
import os
|
|
import folder_paths
|
|
import torch
|
|
from typing import Dict
|
|
|
|
class HunyuanVideoLoraLoader:
|
|
"""
|
|
混元视频LoRA加载器,支持选择性加载blocks
|
|
|
|
这个节点允许您:
|
|
1. 从下拉列表选择LoRA文件
|
|
2. 调整LoRA的强度
|
|
3. 选择要加载的blocks类型(all/single/double)
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.blocks_type = ["all", "single_blocks", "double_blocks"]
|
|
self.loaded_lora = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"lora_name": (folder_paths.get_filename_list("loras"),),
|
|
"strength": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": -10.0,
|
|
"max": 10.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"blocks_type": (["all", "single_blocks", "double_blocks"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load_lora"
|
|
CATEGORY = "loaders/hunyuan"
|
|
OUTPUT_NODE = False
|
|
DESCRIPTION = "加载并应用混元视频LoRA,支持选择性加载single blocks或double blocks"
|
|
|
|
def convert_key_format(self, key: str) -> str:
|
|
"""转换LoRA key格式,支持多种命名方式"""
|
|
# 移除可能的前缀
|
|
prefixes = ["diffusion_model.", "transformer."]
|
|
for prefix in prefixes:
|
|
if key.startswith(prefix):
|
|
key = key[len(prefix):]
|
|
break
|
|
|
|
return key
|
|
|
|
def filter_lora_keys(self, lora: Dict[str, torch.Tensor], blocks_type: str) -> Dict[str, torch.Tensor]:
|
|
"""根据blocks类型过滤LoRA权重"""
|
|
if blocks_type == "all":
|
|
return lora
|
|
|
|
filtered_lora = {}
|
|
for key, value in lora.items():
|
|
base_key = self.convert_key_format(key)
|
|
|
|
# 检查是否包含目标block
|
|
if blocks_type == "single_blocks" in base_key:
|
|
filtered_lora[key] = value
|
|
elif blocks_type == "double_blocks" in base_key:
|
|
filtered_lora[key] = value
|
|
|
|
return filtered_lora
|
|
|
|
def load_lora(self, model, lora_name: str, strength: float, blocks_type: str):
|
|
"""
|
|
加载并应用LoRA到模型
|
|
|
|
Parameters
|
|
----------
|
|
model : ModelPatcher
|
|
要应用LoRA的基础模型
|
|
lora_name : str
|
|
LoRA文件名
|
|
strength : float
|
|
LoRA权重强度
|
|
blocks_type : str
|
|
要加载的blocks类型: "all", "single_blocks" 或 "double_blocks"
|
|
|
|
Returns
|
|
-------
|
|
tuple
|
|
包含应用了LoRA的模型的元组
|
|
"""
|
|
if not lora_name:
|
|
return (model,)
|
|
|
|
from comfy.utils import load_torch_file
|
|
from comfy.sd import load_lora_for_models
|
|
from comfy.lora import load_lora
|
|
|
|
# 获取LoRA文件路径
|
|
lora_path = folder_paths.get_full_path("loras", lora_name)
|
|
if not os.path.exists(lora_path):
|
|
raise Exception(f"Lora {lora_name} not found at {lora_path}")
|
|
|
|
# 缓存LoRA加载
|
|
if self.loaded_lora is not None:
|
|
if self.loaded_lora[0] == lora_path:
|
|
lora = self.loaded_lora[1]
|
|
else:
|
|
self.loaded_lora = None
|
|
|
|
if self.loaded_lora is None:
|
|
lora = load_torch_file(lora_path)
|
|
self.loaded_lora = (lora_path, lora)
|
|
|
|
# 过滤并转换LoRA权重
|
|
filtered_lora = self.filter_lora_keys(lora, blocks_type)
|
|
|
|
# 应用LoRA
|
|
new_model, _ = load_lora_for_models(model, None, filtered_lora, strength, 0)
|
|
if new_model is not None:
|
|
return (new_model,)
|
|
|
|
return (model,)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s, model, lora_name, strength, blocks_type):
|
|
"""当LoRA的配置发生变化时重新执行"""
|
|
return f"{lora_name}_{strength}_{blocks_type}" |