diff --git a/animatediff/nodes.py b/animatediff/nodes.py index 93b086f..ef05d1c 100644 --- a/animatediff/nodes.py +++ b/animatediff/nodes.py @@ -21,6 +21,7 @@ from .nodes_context import (LegacyLoopedUniformContextOptionsNode, LoopedUniform from .nodes_ad_settings import AnimateDiffSettingsNode, ManualAdjustPENode, SweetspotStretchPENode, FullStretchPENode from .nodes_extras import AnimateDiffUnload, EmptyLatentImageLarge, CheckpointLoaderSimpleWithNoiseSelect from .nodes_deprecated import AnimateDiffLoader_Deprecated, AnimateDiffLoaderAdvanced_Deprecated, AnimateDiffCombine_Deprecated +from .nodes_lora import MaskedLoraLoader # override comfy_sample.sample with animatediff-support version comfy_sample.sample = motion_sample_factory(comfy_sample.sample) @@ -107,6 +108,8 @@ NODE_CLASS_MAPPINGS = { "ADE_ApplyAnimateDiffModelSimple": ApplyAnimateDiffModelBasicNode, "ADE_ApplyAnimateDiffModel": ApplyAnimateDiffModelNode, "ADE_LoadAnimateDiffModel": LoadAnimateDiffModelNode, + # MaskedLoraLoader + "ADE_MaskedLoadLora": MaskedLoraLoader, # Deprecated Nodes "AnimateDiffLoaderV1": AnimateDiffLoader_Deprecated, "ADE_AnimateDiffLoaderV1Advanced": AnimateDiffLoaderAdvanced_Deprecated, diff --git a/animatediff/nodes_lora.py b/animatediff/nodes_lora.py new file mode 100644 index 0000000..5300551 --- /dev/null +++ b/animatediff/nodes_lora.py @@ -0,0 +1,51 @@ +import folder_paths +import comfy.utils +import comfy.sd + + +class MaskedLoraLoader: + def __init__(self): + self.loaded_lora = None + + @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 = () + RETURN_TYPES = ("MODEL", "CLIP") + FUNCTION = "load_lora" + + CATEGORY = "loaders" + + def load_lora(self, model, clip, lora_name, strength_model, strength_clip): + if strength_model == 0 and strength_clip == 0: + return (model, clip) + + lora_path = folder_paths.get_full_path("loras", lora_name) + lora = None + if self.loaded_lora is not None: + if self.loaded_lora[0] == lora_path: + lora = self.loaded_lora[1] + else: + temp = self.loaded_lora + self.loaded_lora = None + del temp + + if lora is None: + lora = comfy.utils.load_torch_file(lora_path, safe_load=True) + self.loaded_lora = (lora_path, lora) + + from pathlib import Path + with open(Path(__file__).parent.parent.parent / "sd_lora_keys.txt", "w") as lfile: + for key in lora: + lfile.write(f"{key}:\t{lora[key].size()}\n") + + #model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip) + #return (model_lora, clip_lora) + return (model, clip) + + \ No newline at end of file