From 41aac18ccdb07bc463c7522bb461f7d3c9f341a7 Mon Sep 17 00:00:00 2001 From: Alisson Anjos Date: Tue, 12 May 2026 20:20:55 +0100 Subject: [PATCH] Add Edit Anything module sidecar workflow --- ltxv_editanything.py | 186 +++++++++++++++++++++++++++++++++++-- pyproject.toml | 3 +- split_editanything_lora.py | 82 ++++++++++++++++ 3 files changed, 262 insertions(+), 9 deletions(-) create mode 100644 split_editanything_lora.py diff --git a/ltxv_editanything.py b/ltxv_editanything.py index bfcbedf..86b86f5 100644 --- a/ltxv_editanything.py +++ b/ltxv_editanything.py @@ -13,11 +13,12 @@ Usage in ComfyUI: import types import logging +import os import torch import torch.nn.functional as F import folder_paths -from safetensors.torch import load_file +from safetensors.torch import load_file, save_file logger = logging.getLogger("ltxv.editanything") @@ -27,6 +28,172 @@ _KW_APP_LATENTS = "ea_app_latents" _KW_ROLE_STRENGTH = "ea_role_strength" _KW_ADALN_COND = "ea_adaln_cond" +_EA_LOADER_SKIP_PREFIXES = ( + "role_embedding.", + "_role_embedding.", + "ref_adaln_proj.", + "diffusion_model.role_embedding.", + "diffusion_model.adaln_single.linear.", + "diffusion_model.audio_adaln_single.linear.", +) + + +def _split_editanything_lora_state_dict(sd): + standard = {} + module = {} + for key, value in sd.items(): + if key.startswith(_EA_LOADER_SKIP_PREFIXES): + module[key] = value + else: + standard[key] = value + return standard, module + + +def _filter_editanything_lora_for_standard_loader(sd): + """Remove Edit Anything side-channel weights before ComfyUI's LoRA merger. + + The Apply node reads these weights from the original checkpoint. Passing them + through the standard LoRA loader can make ComfyUI try to reshape AdaLN LoRA + matrices against a different LTXV variant and print shape errors. + """ + filtered, module = _split_editanything_lora_state_dict(sd) + return filtered, list(module.keys()) + + +class LTXVEditAnythingSplitLora: + """Split an Edit Anything checkpoint into standard LoRA and sidecar module.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "lora_name": (folder_paths.get_filename_list("loras"), { + "tooltip": "Original Edit Anything LoRA checkpoint to split.", + }), + "overwrite": ("BOOLEAN", {"default": False}), + }, + } + + RETURN_TYPES = ("STRING", "STRING") + RETURN_NAMES = ("standard_lora_name", "editanything_module_name") + FUNCTION = "split" + CATEGORY = "LTXV/EditAnything" + + def split(self, lora_name, overwrite): + lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) + sd = load_file(lora_path) + standard, module = _split_editanything_lora_state_dict(sd) + + root, ext = os.path.splitext(lora_path) + standard_path = f"{root}.standard{ext}" + module_path = f"{root}.editanything_module{ext}" + + if not overwrite: + for path in (standard_path, module_path): + if os.path.exists(path): + raise FileExistsError( + f"{path} already exists. Enable overwrite to regenerate it." + ) + + save_file(standard, standard_path) + save_file(module, module_path) + + lora_root = folder_paths.get_folder_paths("loras")[0] + standard_name = os.path.relpath(standard_path, lora_root) + module_name = os.path.relpath(module_path, lora_root) + logger.info( + "Split Edit Anything LoRA into %s (%d keys) and %s (%d keys).", + standard_name, + len(standard), + module_name, + len(module), + ) + return (standard_name, module_name) + + +class LTXVEditAnythingModuleLoader: + """Load the sidecar module generated by LTXV Edit Anything Split LoRA.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "module_name": (folder_paths.get_filename_list("loras"), { + "tooltip": "The .editanything_module.safetensors sidecar file.", + }), + }, + } + + RETURN_TYPES = ("LTXV_EA_MODULE",) + RETURN_NAMES = ("editanything_module",) + FUNCTION = "load" + CATEGORY = "LTXV/EditAnything" + + def load(self, module_name): + module_path = folder_paths.get_full_path_or_raise("loras", module_name) + return ({"path": module_path, "state_dict": load_file(module_path)},) + + +class LTXVEditAnythingLoraLoader: + """LoRA loader for Edit Anything checkpoints. + + Use this instead of the standard ComfyUI LoRA loader for LoRAs trained with + the video_to_video_ref_adaln/Edit Anything strategy. + """ + + def __init__(self): + self.loaded_lora = None + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "lora_name": (folder_paths.get_filename_list("loras"), { + "tooltip": "Edit Anything LoRA checkpoint.", + }), + "strength_model": ("FLOAT", { + "default": 1.0, + "min": -100.0, + "max": 100.0, + "step": 0.01, + }), + }, + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "load_lora" + CATEGORY = "LTXV/EditAnything" + + def load_lora(self, model, lora_name, strength_model): + if strength_model == 0: + return (model,) + + import comfy.sd + import comfy.utils + + lora_path = folder_paths.get_full_path_or_raise("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: + self.loaded_lora = None + + if lora is None: + raw_lora = comfy.utils.load_torch_file(lora_path, safe_load=True) + lora, skipped = _filter_editanything_lora_for_standard_loader(raw_lora) + self.loaded_lora = (lora_path, lora) + if skipped: + logger.info( + "LTXVEditAnythingLoraLoader skipped %d Edit Anything/AdaLN keys before standard LoRA load.", + len(skipped), + ) + + model_lora, _ = comfy.sd.load_lora_for_models(model, None, lora, strength_model, 0) + return (model_lora,) + # --------------------------------------------------------------------------- # Model finder @@ -296,13 +463,13 @@ class LTXVEditAnythingApply: "ref_image": ("IMAGE", { "tooltip": "Reference image for appearance. Needs to be a clean image, aligned and facing the camera.", }), + "editanything_module": ("LTXV_EA_MODULE", { + "tooltip": "Sidecar module loaded by LTXV Edit Anything Module Loader.", + }), "resize_mode": (["pad_to_fit", "center_crop", "stretch"], { "default": "pad_to_fit", "tooltip": "How to resize the reference image to match the video's resolution. 'pad_to_fit' avoids distortion.", }), - "lora_name": (folder_paths.get_filename_list("loras"), { - "tooltip": "LoRA trained with the video_to_video_ref_adaln strategy.", - }), }, "optional": { "guide_frames": ("IMAGE", { @@ -350,15 +517,14 @@ class LTXVEditAnythingApply: CATEGORY = "LTXV/EditAnything" def apply( - self, model, positive, negative, vae, latent, ref_image, lora_name, resize_mode="pad_to_fit", + self, model, positive, negative, vae, latent, ref_image, editanything_module, resize_mode="pad_to_fit", guide_frames=None, guide_strength=1.0, ref_strength=1.0, role_strength=1.0, adaln_scale=1.0, enable_adaln=True, enable_role_embedding=True, debug=False, ): import comfy.utils import comfy_extras.nodes_lt as nodes_lt - lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) - sd = load_file(lora_path) + sd = editanything_module["state_dict"] # ── ① Load weights from LoRA ───────────────────────────────────────── # Role embedding: new format [N_slots, 128] @@ -723,12 +889,18 @@ class LTXVResizeReferenceByMask: NODE_CLASS_MAPPINGS = { + "LTXVEditAnythingSplitLora": LTXVEditAnythingSplitLora, + "LTXVEditAnythingModuleLoader": LTXVEditAnythingModuleLoader, + "LTXVEditAnythingLoraLoader": LTXVEditAnythingLoraLoader, "LTXVEditAnythingApply": LTXVEditAnythingApply, "LTXVApplyNeutralMask": LTXVApplyNeutralMask, "LTXVResizeReferenceByMask": LTXVResizeReferenceByMask, } NODE_DISPLAY_NAME_MAPPINGS = { + "LTXVEditAnythingSplitLora": "LTXV Edit Anything Split LoRA", + "LTXVEditAnythingModuleLoader": "LTXV Edit Anything Module Loader", + "LTXVEditAnythingLoraLoader": "LTXV Edit Anything LoRA Loader", "LTXVEditAnythingApply": "LTXV Edit Anything (Apply)", "LTXVApplyNeutralMask": "LTXV Apply Neutral Mask", "LTXVResizeReferenceByMask": "LTXV Resize Reference By Mask", diff --git a/pyproject.toml b/pyproject.toml index 27ee826..26357dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "bfsnodes" description = "Custom ComfyUI nodes including LTXV Edit Anything conditioning helpers." -version = "1.2.0" +version = "1.2.1" license = {file = "LICENSE"} dependencies = [ "torch", @@ -39,4 +39,3 @@ DisplayName = "ComfyUI-BFSNodes" Icon = "" includes = [] # "requires-comfyui" = ">=1.0.0" # ComfyUI version compatibility - diff --git a/split_editanything_lora.py b/split_editanything_lora.py new file mode 100644 index 0000000..cd8ed84 --- /dev/null +++ b/split_editanything_lora.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Split an LTXV Edit Anything LoRA into standard LoRA and sidecar module. + +Usage: + python split_editanything_lora.py path/to/edit_anything.safetensors +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from safetensors.torch import load_file, save_file + + +SIDECAR_PREFIXES = ( + "role_embedding.", + "_role_embedding.", + "ref_adaln_proj.", + "diffusion_model.role_embedding.", + "diffusion_model.adaln_single.linear.", + "diffusion_model.audio_adaln_single.linear.", +) + + +def split_lora(input_path: Path, overwrite: bool = False) -> tuple[Path, Path]: + if not input_path.exists(): + raise FileNotFoundError(input_path) + if input_path.suffix != ".safetensors": + raise ValueError(f"Expected a .safetensors file, got: {input_path}") + + standard_path = input_path.with_name(f"{input_path.stem}.standard.safetensors") + module_path = input_path.with_name( + f"{input_path.stem}.editanything_module.safetensors" + ) + + if not overwrite: + for path in (standard_path, module_path): + if path.exists(): + raise FileExistsError( + f"{path} already exists. Re-run with --overwrite to replace it." + ) + + sd = load_file(str(input_path)) + standard = {} + module = {} + + for key, value in sd.items(): + if key.startswith(SIDECAR_PREFIXES): + module[key] = value + else: + standard[key] = value + + save_file(standard, str(standard_path)) + save_file(module, str(module_path)) + + print(f"Input: {input_path}") + print(f"Standard: {standard_path} ({len(standard)} keys)") + print(f"Module: {module_path} ({len(module)} keys)") + + if not module: + print("Warning: module file is empty; no Edit Anything sidecar keys matched.") + + return standard_path, module_path + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Split an LTXV Edit Anything LoRA into a clean standard LoRA and an Edit Anything module." + ) + parser.add_argument("input", type=Path, help="Path to the original .safetensors LoRA.") + parser.add_argument( + "--overwrite", + action="store_true", + help="Replace existing .standard/.editanything_module outputs.", + ) + args = parser.parse_args() + split_lora(args.input, overwrite=args.overwrite) + + +if __name__ == "__main__": + main()