73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
import torch
|
|
import folder_paths
|
|
import comfy.utils
|
|
import comfy.lora
|
|
import logging
|
|
|
|
class ZImageDiffSynthLoader:
|
|
"""
|
|
A specialized LoRA loader for Z-Image Turbo that automatically fixes
|
|
key mismatches from DiffSynth-trained LoRAs.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"clip": ("CLIP",),
|
|
"lora_name": (folder_paths.get_filename_list("loras"), ),
|
|
"strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}),
|
|
"strength_clip": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP")
|
|
FUNCTION = "load_lora"
|
|
CATEGORY = "Z-Image/Loaders"
|
|
|
|
def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
|
|
lora_path = folder_paths.get_full_path("loras", lora_name)
|
|
lora_sd = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
|
|
|
# --- KEY MAPPING LOGIC ---
|
|
# DiffSynth keys are "flat" (e.g. "layers.0.attention...")
|
|
# ComfyUI Z-Image expects "diffusion_model." prefix for these.
|
|
|
|
fixed_sd = {}
|
|
for k, v in lora_sd.items():
|
|
new_key = k
|
|
|
|
# 1. Handle specialized Z-Image Refiners
|
|
if k.startswith("context_refiner."):
|
|
new_key = f"diffusion_model.{k}"
|
|
elif k.startswith("noise_refiner."):
|
|
new_key = f"diffusion_model.{k}"
|
|
|
|
# 2. Handle standard Llama-3 Backbone Layers
|
|
# If it starts with "layers." it belongs to the DiT backbone
|
|
elif k.startswith("layers."):
|
|
new_key = f"diffusion_model.{k}"
|
|
|
|
# 3. Handle Input/Output blocks if they are bare
|
|
elif k.startswith("final_layer.") or k.startswith("label_emb."):
|
|
new_key = f"diffusion_model.{k}"
|
|
|
|
fixed_sd[new_key] = v
|
|
|
|
# --- APPLY PATCHES ---
|
|
# We pass the fixed dictionary to ComfyUI's standard loader
|
|
model_lora, clip_lora = comfy.sd.load_lora_for_models(
|
|
model, clip, fixed_sd, strength_model, strength_clip
|
|
)
|
|
|
|
return (model_lora, clip_lora)
|
|
|
|
# Node Registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ZImageDiffSynthLoader": ZImageDiffSynthLoader
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ZImageDiffSynthLoader": "Z-Image DiffSynth LoRA Loader"
|
|
} |