86 lines
4.5 KiB
Python
86 lines
4.5 KiB
Python
"""The LoRA loader.
|
|
|
|
One LoRA per node, chained for more. v4 carried N LoRAs in a single node behind a
|
|
JSON list widget and a `FlexibleOptionalInputType` -- a dict subclass that lied to
|
|
ComfyUI's validator so arbitrary keys would pass. Chaining needs neither: it is
|
|
what ComfyUI's own LoraLoader does, every strength is a real socket, and the graph
|
|
shows how many LoRAs are actually applied instead of hiding the count in a string.
|
|
"""
|
|
|
|
import comfy.utils
|
|
import folder_paths
|
|
from comfy_api.latest import io
|
|
|
|
from ..._core import log
|
|
from .keys import match
|
|
|
|
|
|
NONE = "None"
|
|
|
|
|
|
class FunPackLoraLoader(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="FunPackLoraLoader",
|
|
display_name="FunPack LoRA Loader",
|
|
category="FunPack/Loaders",
|
|
description="Apply one LoRA to a model, and optionally its text encoder. "
|
|
"Chain these for more than one.",
|
|
inputs=[
|
|
io.Model.Input("model"),
|
|
io.Combo.Input("lora_name", options=[NONE] + folder_paths.get_filename_list("loras"),
|
|
default=NONE, tooltip="None leaves the model untouched."),
|
|
io.Float.Input("strength_model", default=1.0, min=-100.0, max=100.0, step=0.01,
|
|
tooltip="How strongly the LoRA changes the model. 0 = not at all, negative = the opposite."),
|
|
io.Clip.Input("clip", optional=True,
|
|
tooltip="Leave unwired for a model-only LoRA."),
|
|
io.Float.Input("strength_clip", default=1.0, min=-100.0, max=100.0, step=0.01,
|
|
optional=True, tooltip="How strongly the LoRA changes the text encoder. Only used when a clip is wired."),
|
|
],
|
|
outputs=[
|
|
io.Model.Output(display_name="model"),
|
|
io.Clip.Output(display_name="clip"),
|
|
io.String.Output(display_name="status"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model, lora_name: str, strength_model: float,
|
|
clip=None, strength_clip: float = 1.0) -> io.NodeOutput:
|
|
if lora_name == NONE: # kept in the pipeline so a LoRA is one pick away; nothing picked = no effect
|
|
return io.NodeOutput(model, clip, "no LoRA picked")
|
|
# Doing nothing is a valid request, but it must return the SAME objects
|
|
# rather than clones: a clone here would silently drop patches a later
|
|
# loader in the chain applied to the original.
|
|
if strength_model == 0 and (clip is None or strength_clip == 0):
|
|
# Loud, because "applied a LoRA at strength 0" is indistinguishable
|
|
# from "forgot to set the strength" by looking at the picture.
|
|
log.alert("FunPack LoRA Loader",
|
|
f"{lora_name} has both strengths at 0 and was not applied")
|
|
return io.NodeOutput(model, clip, f"{lora_name}: both strengths are 0, not applied")
|
|
|
|
path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
|
lora = comfy.utils.load_torch_file(path, safe_load=True)
|
|
patches, how, dropped = match(model, clip, lora)
|
|
if not patches:
|
|
log.alert("FunPack LoRA Loader", f"{lora_name} matched NOTHING in this model, so it does nothing: "
|
|
"it is for another model, or names its layers in a way not known here")
|
|
return io.NodeOutput(model, clip, f"{lora_name}: matched nothing, not applied")
|
|
patched_model = model.clone()
|
|
patched_model.add_patches(patches, strength_model)
|
|
patched_clip = None
|
|
if clip is not None:
|
|
patched_clip = clip.clone()
|
|
patched_clip.add_patches(patches, strength_clip)
|
|
if dropped:
|
|
log.alert("FunPack LoRA Loader", f"{lora_name}: {dropped} of {len(patches) + dropped} weights do not fit this "
|
|
"model's shapes and were left out (trained on another variant of it)")
|
|
|
|
applied = [f"model={strength_model}"]
|
|
applied.append(f"clip={strength_clip}" if clip is not None else "clip=not wired")
|
|
found = f"{len(patches)} weights" + (f", keys {how}" if how != "as-is" else "") + (f", {dropped} dropped" if dropped else "")
|
|
log.info("FunPack LoRA Loader", f"{lora_name} applied ({', '.join(applied)}; {found})")
|
|
return io.NodeOutput(patched_model, patched_clip,
|
|
f"FunPack LoRA Loader | {lora_name} ({', '.join(applied)}; {found})")
|