From ef719b21733db21f0d06c88ea3050978384e66cf Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 8 Jul 2026 11:07:58 +0800 Subject: [PATCH] feat: add anima xpred lora loader --- README.md | 11 ++++++ __init__.py | 44 ++++++++++++++++++++++ tests/test_lora_loader.py | 79 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 134 insertions(+) create mode 100644 tests/test_lora_loader.py diff --git a/README.md b/README.md index 2ddeddf..d000eb9 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ Put model files under ComfyUI's normal `models/` folders: ComfyUI/models/diffusion_models/ xpred-adapter-checkpoint.safetensors ComfyUI/models/text_encoders/ qwen_3_06b_base.safetensors ComfyUI/models/vae/ qwen_image_vae.safetensors +ComfyUI/models/loras/ anima-lora.safetensors ``` The loader uses ComfyUI dropdowns for these files instead of absolute path text boxes. @@ -38,6 +39,10 @@ The x-pred checkpoint is an Anima DiT checkpoint and is loaded directly through - Loads the Anima DiT x-pred checkpoint, text encoder, and VAE. - Use `xpred-adapter-checkpoint.safetensors`, not `xpred-train-state.pt`. +- `Load Anima XPred LoRA` + - Loads an Anima DiT LoRA from ComfyUI's `models/loras` folder. + - Takes and returns `RUM_ANIMA_XPRED`, so it belongs between the model loader and sampler. + - `Sample Anima XPred` - Runs the dedicated x-pred sampler with `heun` or `euler`: `v = (z - x_pred) / sigma` @@ -53,4 +58,10 @@ Connect: Load Anima XPred Model -> Sample Anima XPred -> Preview Image / Save Image ``` +With a LoRA: + +```text +Load Anima XPred Model -> Load Anima XPred LoRA -> Sample Anima XPred -> Preview Image / Save Image +``` + Do not connect an x-pred checkpoint to a regular Anima/FM velocity sampler. The checkpoint predicts clean latent `x`, not velocity. diff --git a/__init__.py b/__init__.py index 30dfd97..0b766f2 100644 --- a/__init__.py +++ b/__init__.py @@ -2,6 +2,7 @@ from __future__ import annotations import gc import sys +from copy import copy from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace @@ -229,6 +230,47 @@ class AnimaXPredModelLoader: return (LoadedAnimaXPred(args=args, adapter=adapter, student=student, device=device, dtype=dtype, prediction_type=prediction_type),) +class AnimaXPredLoraLoader: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("RUM_ANIMA_XPRED",), + "lora": (_model_names("loras", "put_anima_lora_in_models_loras.safetensors"),), + "strength": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.05}), + } + } + + RETURN_TYPES = ("RUM_ANIMA_XPRED",) + RETURN_NAMES = ("model",) + FUNCTION = "load_lora" + CATEGORY = "RUM/Anima XPred" + + def load_lora(self, model: LoadedAnimaXPred, lora: str, strength: float): + lora_path = _model_path("loras", lora) + if not lora_path.is_file(): + raise FileNotFoundError(f"lora not found: {lora_path}") + + args = copy(model.args) + args.teacher_lora = str(lora_path) + args.teacher_lora_weight = float(strength) + + adapter = create_adapter(args, device=model.device, dtype=model.dtype) + checkpoint_path = getattr(args, "student_init", None) or getattr(args, "dit", None) + student = adapter.load_student_xpred(init_checkpoint=checkpoint_path) + student.to(device=model.device, dtype=model.dtype).eval().requires_grad_(False) + return ( + LoadedAnimaXPred( + args=args, + adapter=adapter, + student=student, + device=model.device, + dtype=model.dtype, + prediction_type=model.prediction_type, + ), + ) + + class AnimaXPredSampler: @classmethod def INPUT_TYPES(cls): @@ -346,10 +388,12 @@ class AnimaXPredSampler: NODE_CLASS_MAPPINGS = { "AnimaXPredModelLoader": AnimaXPredModelLoader, + "AnimaXPredLoraLoader": AnimaXPredLoraLoader, "AnimaXPredSampler": AnimaXPredSampler, } NODE_DISPLAY_NAME_MAPPINGS = { "AnimaXPredModelLoader": "Load Anima XPred Model", + "AnimaXPredLoraLoader": "Load Anima XPred LoRA", "AnimaXPredSampler": "Sample Anima XPred", } diff --git a/tests/test_lora_loader.py b/tests/test_lora_loader.py new file mode 100644 index 0000000..1e31a59 --- /dev/null +++ b/tests/test_lora_loader.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path +from types import SimpleNamespace + +import torch + + +ROOT = Path(__file__).resolve().parents[1] + + +def load_node_module(): + spec = importlib.util.spec_from_file_location("rum_anima_xpred_nodes", ROOT / "__init__.py") + assert spec is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +class DummyAdapter: + def __init__(self, args, device, dtype): + self.args = args + self.device = device + self.dtype = dtype + + def load_student_xpred(self, init_checkpoint): + self.init_checkpoint = init_checkpoint + return torch.nn.Linear(1, 1) + + +def test_lora_loader_returns_new_rum_model(monkeypatch, tmp_path): + module = load_node_module() + lora_path = tmp_path / "style.safetensors" + lora_path.write_bytes(b"placeholder") + + created = {} + + def fake_model_path(folder_name, filename): + assert folder_name == "loras" + assert filename == "style.safetensors" + return lora_path + + def fake_create_adapter(args, device, dtype): + created["args"] = args + return DummyAdapter(args, device, dtype) + + monkeypatch.setattr(module, "_model_path", fake_model_path) + monkeypatch.setattr(module, "create_adapter", fake_create_adapter) + + original_args = SimpleNamespace( + dit="/models/diffusion_models/base.safetensors", + student_init="/models/diffusion_models/base.safetensors", + text_encoder="/models/text_encoders/qwen.safetensors", + vae="/models/vae/qwen_vae.safetensors", + teacher_lora=None, + teacher_lora_weight=1.0, + ) + original = module.LoadedAnimaXPred( + args=original_args, + adapter=object(), + student=torch.nn.Linear(1, 1), + device=torch.device("cpu"), + dtype=torch.float32, + prediction_type="x", + ) + + (loaded,) = module.AnimaXPredLoraLoader().load_lora(original, "style.safetensors", 0.75) + + assert loaded is not original + assert loaded.args is not original_args + assert loaded.args.teacher_lora == str(lora_path) + assert loaded.args.teacher_lora_weight == 0.75 + assert loaded.args.student_init == original_args.student_init + assert loaded.prediction_type == "x" + assert created["args"] is loaded.args