From 8d5d3be3a077b61b7580edd8fdbbabd05f05efe3 Mon Sep 17 00:00:00 2001
From: Codex <codex@local>
Date: Mon, 6 Jul 2026 16:52:40 +0800
Subject: [PATCH] fix: load xpred checkpoints from diffusion models

---
 README.md   |  4 ++--
 __init__.py | 17 ++---------------
 2 files changed, 4 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index 2035177..f681bb9 100644
--- a/README.md
+++ b/README.md
@@ -31,13 +31,13 @@ export RUM_ANIMA_XPRED_ROOT=/path/to/RUM-anima-xpred
 Put model files under ComfyUI's normal `models/` folders:
 
 ```text
-ComfyUI/models/anima_xpred/       xpred-adapter-checkpoint.safetensors
+ComfyUI/models/diffusion_models/  xpred-adapter-checkpoint.safetensors
 ComfyUI/models/text_encoders/     qwen_3_06b_base.safetensors
 ComfyUI/models/vae/               qwen_image_vae.safetensors
 ```
 
 The loader uses ComfyUI dropdowns for these files instead of absolute path text boxes.
-The x-pred checkpoint is loaded directly through the normal Anima DiT loader, so a separate base DiT file is not required.
+The x-pred checkpoint is an Anima DiT checkpoint and is loaded directly through the normal Anima DiT loader, so it belongs in the same `diffusion_models` folder as ordinary Anima DiT files. A separate base DiT file is not required.
 
 ## Nodes
 
diff --git a/__init__.py b/__init__.py
index 0395850..7669a81 100644
--- a/__init__.py
+++ b/__init__.py
@@ -12,7 +12,6 @@ import torch
 
 
 NODE_ROOT = Path(__file__).resolve().parent
-XPRED_FOLDER = "anima_xpred"
 
 
 def _backend_root() -> Path:
@@ -46,18 +45,6 @@ from rum_xpred.anima import anima_euler_step, make_shifted_sigma_schedule, xpred
 from rum_xpred.cache_batches import make_seeded_eps_batch
 
 
-def _register_model_folders() -> None:
-    try:
-        import folder_paths
-
-        folder_paths.add_model_folder_path(XPRED_FOLDER, str(Path(folder_paths.models_dir) / XPRED_FOLDER))
-    except Exception:
-        pass
-
-
-_register_model_folders()
-
-
 def _model_names(folder_name: str, placeholder: str) -> list[str]:
     try:
         import folder_paths
@@ -199,7 +186,7 @@ class AnimaXPredModelLoader:
     def INPUT_TYPES(cls):
         return {
             "required": {
-                "checkpoint": (_model_names(XPRED_FOLDER, "put_xpred_checkpoint_in_models_anima_xpred.safetensors"),),
+                "checkpoint": (_model_names("diffusion_models", "put_xpred_checkpoint_in_models_diffusion_models.safetensors"),),
                 "text_encoder": (_model_names("text_encoders", "put_qwen_text_encoder_in_models_text_encoders.safetensors"),),
                 "vae": (_model_names("vae", "put_qwen_image_vae_in_models_vae.safetensors"),),
                 "prediction_type": (["x", "v"], {"default": "x"}),
@@ -228,7 +215,7 @@ class AnimaXPredModelLoader:
         fp8: bool,
         fp8_scaled: bool,
     ):
-        checkpoint_path = _model_path(XPRED_FOLDER, checkpoint)
+        checkpoint_path = _model_path("diffusion_models", checkpoint)
         text_encoder_path = _model_path("text_encoders", text_encoder)
         vae_path = _model_path("vae", vae)
         for label, path in {