feat: add anima xpred lora loader

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.

__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",
 }

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