From 08e82a3eb6f0fe0e65a55075d2b807bab9d76de6 Mon Sep 17 00:00:00 2001
From: Codex <codex@local>
Date: Mon, 6 Jul 2026 16:46:49 +0800
Subject: [PATCH] fix: load xpred checkpoint directly

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

diff --git a/README.md b/README.md
index 5ad0344..2035177 100644
--- a/README.md
+++ b/README.md
@@ -32,17 +32,17 @@ Put model files under ComfyUI's normal `models/` folders:
 
 ```text
 ComfyUI/models/anima_xpred/       xpred-adapter-checkpoint.safetensors
-ComfyUI/models/diffusion_models/  anima-base-v1.0.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.
 
 ## Nodes
 
 - `Load Anima XPred Model`
-  - Loads the Anima DiT x-pred checkpoint, text encoder, and VAE paths.
+  - Loads the Anima DiT x-pred checkpoint, text encoder, and VAE.
   - Use `xpred-adapter-checkpoint.safetensors`, not `xpred-train-state.pt`.
 
 - `Sample Anima XPred`
diff --git a/__init__.py b/__init__.py
index 5eae2fd..0395850 100644
--- a/__init__.py
+++ b/__init__.py
@@ -202,7 +202,6 @@ class AnimaXPredModelLoader:
                 "checkpoint": (_model_names(XPRED_FOLDER, "put_xpred_checkpoint_in_models_anima_xpred.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"),),
-                "base_dit": (_model_names("diffusion_models", "put_anima_base_dit_in_models_diffusion_models.safetensors"),),
                 "prediction_type": (["x", "v"], {"default": "x"}),
                 "precision": (["bf16", "fp16", "fp32"], {"default": "bf16"}),
                 "attn_mode": (["torch", "flash", "sageattn", "xformers"], {"default": "flash"}),
@@ -222,7 +221,6 @@ class AnimaXPredModelLoader:
         checkpoint: str,
         text_encoder: str,
         vae: str,
-        base_dit: str,
         prediction_type: str,
         precision: str,
         attn_mode: str,
@@ -233,12 +231,10 @@ class AnimaXPredModelLoader:
         checkpoint_path = _model_path(XPRED_FOLDER, checkpoint)
         text_encoder_path = _model_path("text_encoders", text_encoder)
         vae_path = _model_path("vae", vae)
-        base_dit_path = _model_path("diffusion_models", base_dit)
         for label, path in {
             "checkpoint": checkpoint_path,
             "text_encoder": text_encoder_path,
             "vae": vae_path,
-            "base_dit": base_dit_path,
         }.items():
             if not path.is_file():
                 raise FileNotFoundError(f"{label} not found: {path}")
@@ -246,7 +242,7 @@ class AnimaXPredModelLoader:
         device = _comfy_device()
         dtype = _dtype_from_name(precision)
         args = SimpleNamespace(
-            dit=str(base_dit_path),
+            dit=str(checkpoint_path),
             student_init=str(checkpoint_path),
             text_encoder=str(text_encoder_path),
             vae=str(vae_path),