Initial AnimaFastTrain nodes

README.md

@@ -0,0 +1,72 @@
+# ComfyUI AnimaFastTrain
+
+Experimental custom nodes for training Anima reference context in ComfyUI with
+in-memory reference context tokens.
+
+## Installation
+
+Clone this repository into your ComfyUI `custom_nodes` directory:
+
+```bash
+cd ComfyUI/custom_nodes
+git clone https://github.com/quinteroac/ComfyUI-AnimaFastTrain.git
+```
+
+Then restart ComfyUI. The nodes will appear under:
+
+```text
+Anima/AnimaFastTrain
+```
+
+To update an existing installation:
+
+```bash
+cd ComfyUI/custom_nodes/ComfyUI-AnimaFastTrain
+git pull
+```
+
+Restart ComfyUI after updating.
+
+## Nodes
+
+- **AnimaFastTrain - Train Context Tokens**
+  - Inputs: `MODEL`, `CLIP`, `VAE`, and a reference `IMAGE`.
+  - Trains per-block context tokens in memory.
+  - Does not save safetensors or any weight cache to disk.
+  - Training uses scale `1.0`, matching the experimental Python pipeline.
+    Runtime strength is controlled by the patch node.
+  - Default parameters:
+    - `training_steps=80`
+    - `learning_rate=0.02`
+    - internal Anima shape: `num_blocks=28`, `text_dim=1024`
+
+- **AnimaFastTrain - Patch Model**
+  - Inputs: `MODEL` and the in-memory context object.
+  - Outputs a patched `MODEL`.
+  - Connect the patched model to a normal ComfyUI sampler.
+  - The patch is applied during the model call and restored immediately after
+    each call.
+  - Default runtime strength:
+    - `reference_context_scale=0.45`
+
+## Important graph order
+
+Connect the final model into **Patch Model** after any LoRA/model patch nodes.
+For example:
+
+`Checkpoint Loader -> LoRA Loader -> Patch Model -> KSampler`
+
+If another model patch node runs after **Patch Model**, it may clone the model
+and drop this experimental wrapper.
+
+## Suggested workflow
+
+1. Load the Anima model, CLIP, and VAE.
+2. Load or create a reference image. A face-dominant crop often improves
+   facial identity consistency.
+3. Run **Train Context Tokens**.
+4. Run **Patch Model** with the returned context.
+5. Use the patched model with normal conditioning, latent, sampler, and VAE
+   decode nodes.
+
+Everything is intentionally memory-only for experimentation.

__init__.py

@@ -0,0 +1,7 @@
+"""Experimental AnimaFastTrain nodes for ComfyUI."""
+
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+WEB_DIRECTORY = None
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]

nodes.py

@@ -0,0 +1,618 @@
+"""Experimental AnimaFastTrain custom nodes.
+
+These nodes keep identity context tokens in memory only:
+
+- Train Reference Context Tokens: optimizes per-block Anima context tokens from
+  one reference image and returns an in-memory token object.
+- Patch Model With Context Tokens: returns a MODEL whose UNet call temporarily
+  injects those tokens into Anima block cross-attention.
+
+No safetensors cache is written by these nodes.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+from dataclasses import dataclass
+from types import MethodType
+from typing import Any, Callable
+
+import torch
+
+LOGGER = logging.getLogger("ComfyUI-AnimaFastTrain")
+
+ANIMA_NUM_BLOCKS = 28
+ANIMA_TEXT_DIM = 1024
+
+
+@dataclass
+class AnimaReferenceContext:
+    tokens: Any
+    losses: list[float]
+    training_steps: int
+    learning_rate: float
+    training_seed: int
+    reference_count: int
+    num_blocks: int
+    num_tokens: int
+    text_dim: int
+    application_count: int = 0
+
+
+def _normalize_prompt_text(text: str) -> str:
+    return " " if text == "" else text
+
+
+def _encode_prompt(clip: Any, text: str) -> Any:
+    tokens = clip.tokenize(_normalize_prompt_text(text))
+    return clip.encode_from_tokens_scheduled(tokens)
+
+
+def _reference_image_seed(reference_image: Any) -> int:
+    images = reference_image.detach().cpu().clamp(0, 1)
+    image_bytes = (images * 255).to(torch.uint8).numpy().tobytes()
+    digest = hashlib.sha256(image_bytes).digest()
+    return int.from_bytes(digest[:8], byteorder="little", signed=False)
+
+
+def _resize_reference_image(reference_image: Any, size: int) -> Any:
+    if size <= 0:
+        return reference_image[:1]
+    image = reference_image[:1].movedim(-1, 1)
+    image = torch.nn.functional.interpolate(
+        image,
+        size=(size, size),
+        mode="bicubic",
+        align_corners=False,
+    )
+    return image.movedim(1, -1).clamp(0, 1)
+
+
+def _prepare_reference_images(
+    reference_images: list[Any],
+    *,
+    size: int,
+) -> Any:
+    resized_images = []
+    for image in reference_images:
+        if image is None:
+            continue
+        resized_images.append(_resize_reference_image(image, size))
+    if not resized_images:
+        raise ValueError("at least one reference image is required")
+    return torch.cat(resized_images, dim=0)
+
+
+def _create_reference_context_tokens(
+    reference_image: Any,
+    *,
+    num_blocks: int,
+    num_tokens: int,
+    text_dim: int,
+    dtype: str,
+    std: float,
+) -> Any:
+    if num_blocks < 1:
+        raise ValueError("num_blocks must be >= 1")
+    if num_tokens < 1:
+        raise ValueError("num_tokens must be >= 1")
+    if text_dim < 1:
+        raise ValueError("text_dim must be >= 1")
+    if std < 0:
+        raise ValueError("init_std must be >= 0")
+
+    torch_dtype = getattr(torch, dtype)
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(_reference_image_seed(reference_image))
+    return torch.randn(
+        (num_blocks, num_tokens, text_dim),
+        generator=generator,
+        dtype=torch_dtype,
+    ) * std
+
+
+def _get_anima_diffusion_model(model: Any) -> Any:
+    if hasattr(model, "model") and hasattr(model.model, "diffusion_model"):
+        return model.model.diffusion_model
+    if hasattr(model, "diffusion_model"):
+        return model.diffusion_model
+    if hasattr(model, "inner_model"):
+        return _get_anima_diffusion_model(model.inner_model)
+    if hasattr(model, "inner_model_wrapper"):
+        return _get_anima_diffusion_model(model.inner_model_wrapper)
+    if hasattr(model, "blocks"):
+        return model
+    raise TypeError("model must be an Anima diffusion model or ComfyUI ModelPatcher")
+
+
+def _prepare_block_context_tokens(tokens: Any, block_count: int) -> Any:
+    if len(tokens.shape) == 2:
+        tokens = tokens.unsqueeze(0).unsqueeze(0)
+    elif len(tokens.shape) == 3:
+        if tokens.shape[0] == block_count:
+            tokens = tokens.unsqueeze(1)
+        else:
+            tokens = tokens.unsqueeze(0)
+    elif len(tokens.shape) != 4:
+        raise ValueError(
+            "context tokens must be shaped [tokens, text_dim], "
+            "[batch, tokens, text_dim], [blocks, tokens, text_dim], "
+            "or [blocks, batch, tokens, text_dim]"
+        )
+
+    if tokens.shape[0] == 1 and block_count != 1:
+        tokens = tokens.repeat(block_count, 1, 1, 1)
+    if tokens.shape[0] != block_count:
+        raise ValueError(
+            f"context token block count must be 1 or match model blocks ({tokens.shape[0]} != {block_count})"
+        )
+    if not torch.is_floating_point(tokens):
+        raise ValueError("context tokens must be floating point")
+    return tokens
+
+
+def _apply_anima_context_injection(
+    model: Any,
+    tokens: Any,
+    *,
+    scale: float,
+    mode: str,
+    debug_state: dict[str, int] | None = None,
+) -> Callable[[], None]:
+    if mode not in {"positive", "all"}:
+        raise ValueError("mode must be one of: positive, all")
+
+    diffusion_model = _get_anima_diffusion_model(model)
+    blocks = getattr(diffusion_model, "blocks", None)
+    if blocks is None:
+        raise TypeError("Anima diffusion model must expose a blocks attribute")
+    if debug_state is not None and not debug_state.get("patch_logged"):
+        debug_state["patch_logged"] = 1
+        LOGGER.info(
+            "Resolved Anima diffusion model for context injection: active_type=%s diffusion_type=%s blocks=%s",
+            type(model).__name__,
+            type(diffusion_model).__name__,
+            len(blocks),
+        )
+
+    block_tokens = _prepare_block_context_tokens(tokens, len(blocks))
+    originals = []
+
+    for block_index, block in enumerate(blocks):
+        original_forward = block.forward
+        originals.append((block, original_forward))
+
+        def _forward_with_context_injection(
+            self: Any,
+            x_B_T_H_W_D: Any,
+            emb_B_T_D: Any,
+            crossattn_emb: Any,
+            *args: Any,
+            _block_index: int = block_index,
+            _original_forward: Any = original_forward,
+            **kwargs: Any,
+        ) -> Any:
+            injected = block_tokens[_block_index].to(
+                device=crossattn_emb.device,
+                dtype=crossattn_emb.dtype,
+            )
+            if scale != 1.0:
+                injected = injected * scale
+            if injected.shape[0] == 1 and crossattn_emb.shape[0] != 1:
+                injected = injected.repeat(crossattn_emb.shape[0], 1, 1)
+            if injected.shape[0] != crossattn_emb.shape[0]:
+                raise ValueError(
+                    "context token batch size must be 1 or match cross-attention batch size"
+                )
+            transformer_options = kwargs.get("transformer_options", {})
+            cond_or_uncond = transformer_options.get("cond_or_uncond")
+            if mode == "positive" and cond_or_uncond:
+                flags = torch.tensor(
+                    cond_or_uncond,
+                    device=crossattn_emb.device,
+                    dtype=injected.dtype,
+                )
+                if crossattn_emb.shape[0] % flags.shape[0] == 0:
+                    flags = flags.repeat_interleave(crossattn_emb.shape[0] // flags.shape[0])
+                if flags.shape[0] == crossattn_emb.shape[0]:
+                    positive_mask = (flags == 0).to(injected.dtype).view(-1, 1, 1)
+                    injected = injected * positive_mask
+            if injected.shape[2] != crossattn_emb.shape[2]:
+                raise ValueError(
+                    "context token text_dim must match cross-attention context dim "
+                    f"({injected.shape[2]} != {crossattn_emb.shape[2]})"
+                )
+            if debug_state is not None and _block_index == 0:
+                debug_state["block_calls"] = debug_state.get("block_calls", 0) + 1
+                block_calls = debug_state["block_calls"]
+                if block_calls <= 3:
+                    LOGGER.info(
+                        "Injected Anima context inside block: block=%s call=%s crossattn_before=%s injected=%s injected_norm=%.6f",
+                        _block_index,
+                        block_calls,
+                        tuple(crossattn_emb.shape),
+                        tuple(injected.shape),
+                        float(injected.detach().float().norm().cpu()),
+                    )
+            patched_context = torch.cat([injected, crossattn_emb], dim=1)
+            return _original_forward(
+                x_B_T_H_W_D,
+                emb_B_T_D,
+                patched_context,
+                *args,
+                **kwargs,
+            )
+
+        block.forward = MethodType(_forward_with_context_injection, block)
+
+    def restore() -> None:
+        for block, original_forward in originals:
+            block.forward = original_forward
+
+    return restore
+
+
+def _optimize_reference_context_tokens(
+    *,
+    model: Any,
+    vae: Any,
+    positive: Any,
+    reference_images: Any,
+    initial_tokens: Any,
+    steps: int,
+    learning_rate: float,
+    seed: int,
+    training_dtype: str,
+) -> tuple[Any, list[float]]:
+    if steps < 1:
+        raise ValueError("steps must be >= 1")
+    if learning_rate <= 0:
+        raise ValueError("learning_rate must be > 0")
+
+    import comfy.model_management
+    from comfy_extras.nodes_train import TrainGuider
+
+    dtype_map = {
+        "bfloat16": torch.bfloat16,
+        "bf16": torch.bfloat16,
+        "float16": torch.float16,
+        "fp16": torch.float16,
+        "float32": torch.float32,
+        "fp32": torch.float32,
+    }
+    if training_dtype not in dtype_map:
+        raise ValueError(f"unsupported training_dtype: {training_dtype}")
+    autocast_dtype = dtype_map[training_dtype]
+
+    encoded = vae.encode(reference_images)
+    latent_bank = encoded["samples"] if isinstance(encoded, dict) else encoded
+    if latent_bank.shape[0] < 1:
+        raise ValueError("encoded reference latents must not be empty")
+
+    losses: list[float] = []
+
+    with torch.inference_mode(False), torch.enable_grad():
+        context_tokens = torch.nn.Parameter(initial_tokens.detach().clone().float())
+        optimizer = torch.optim.AdamW([context_tokens], lr=learning_rate)
+        loss_fn = torch.nn.MSELoss()
+
+        class ContextTokenTrainSampler:
+            def sample(
+                self,
+                model_wrap: Any,
+                sigmas: Any,
+                extra_args: dict[str, Any],
+                callback: Any,
+                noise: Any,
+                latent_image: Any = None,
+                denoise_mask: Any = None,
+                disable_pbar: bool = False,
+            ) -> Any:
+                device = latent_image.device
+                latent_count = latent_image.shape[0]
+
+                for step in range(steps):
+                    latent_index = step % latent_count
+                    latent_batch = latent_image[latent_index : latent_index + 1]
+                    generator = torch.Generator(device=device)
+                    generator.manual_seed(seed + step * 1009)
+                    batch_noise = torch.randn(
+                        latent_batch.shape,
+                        generator=generator,
+                        device=device,
+                        dtype=latent_batch.dtype,
+                    )
+                    sigma = model_wrap.inner_model.model_sampling.percent_to_sigma(
+                        torch.rand((1,), device=device).item()
+                    )
+                    batch_sigmas = torch.tensor([sigma], device=device)
+
+                    xt = model_wrap.inner_model.model_sampling.noise_scaling(
+                        batch_sigmas,
+                        batch_noise,
+                        latent_batch,
+                        False,
+                    )
+                    x0 = model_wrap.inner_model.model_sampling.noise_scaling(
+                        torch.zeros_like(batch_sigmas),
+                        torch.zeros_like(batch_noise),
+                        latent_batch,
+                        False,
+                    )
+
+                    optimizer.zero_grad(set_to_none=True)
+                    with torch.autocast(xt.device.type, dtype=autocast_dtype):
+                        pred = model_wrap(
+                            xt.requires_grad_(True),
+                            batch_sigmas.requires_grad_(True),
+                            **extra_args,
+                        )
+                        loss = loss_fn(pred.float(), x0.float())
+                    loss.backward()
+                    optimizer.step()
+                    losses.append(float(loss.detach().cpu()))
+                    if not disable_pbar:
+                        print(f"context step {step + 1}/{steps}: loss={losses[-1]:.6f}")
+
+                return latent_image
+
+        model_patcher = model.clone()
+        model_patcher.model.requires_grad_(False)
+        model_patcher.set_model_compute_dtype(autocast_dtype)
+        token_norm_before = float(context_tokens.detach().float().norm().cpu())
+        LOGGER.info(
+            "Starting Anima context token training: steps=%s lr=%s dtype=%s tokens_shape=%s token_norm_before=%.6f reference_shape=%s grad_enabled=%s inference_mode=%s",
+            steps,
+            learning_rate,
+            training_dtype,
+            tuple(initial_tokens.shape),
+            token_norm_before,
+            tuple(reference_images.shape),
+            torch.is_grad_enabled(),
+            torch.is_inference_mode_enabled(),
+        )
+        restore_context_injection = _apply_anima_context_injection(
+            model_patcher,
+            context_tokens,
+            scale=1.0,
+            mode="positive",
+        )
+        try:
+            comfy.model_management.load_models_gpu(
+                [model_patcher],
+                memory_required=1e20,
+                force_full_load=True,
+            )
+            guider = TrainGuider(model_patcher, offloading=False)
+            guider.set_conds(positive)
+            sigmas = torch.tensor([0])
+            dummy_noise = torch.zeros_like(latent_bank)
+            guider.sample(
+                dummy_noise,
+                latent_bank,
+                ContextTokenTrainSampler(),
+                sigmas,
+                disable_pbar=True,
+                seed=seed,
+            )
+        finally:
+            restore_context_injection()
+
+        final_loss = losses[-1] if losses else 0.0
+        token_norm_after = float(context_tokens.detach().float().norm().cpu())
+        LOGGER.info(
+            "Finished Anima context token training: steps=%s final_loss=%.6f token_norm_before=%.6f token_norm_after=%.6f token_norm_delta=%.6f tokens_device=%s",
+            steps,
+            final_loss,
+            token_norm_before,
+            token_norm_after,
+            token_norm_after - token_norm_before,
+            context_tokens.device,
+        )
+        return context_tokens.detach().cpu(), losses
+
+
+class AnimaFastTrainReferenceContext:
+    @classmethod
+    def INPUT_TYPES(cls) -> dict[str, Any]:
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "clip": ("CLIP",),
+                "vae": ("VAE",),
+                "reference_image": ("IMAGE",),
+                "training_prompt": (
+                    "STRING",
+                    {
+                        "default": (
+                            "masterpiece, best quality, score_7, safe, same character, "
+                            "solo, clean white background, clean anime illustration"
+                        ),
+                        "multiline": True,
+                    },
+                ),
+                "seed": (
+                    "INT",
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 0xffffffffffffffff,
+                        "control_after_generate": "fixed",
+                    },
+                ),
+                "training_steps": ("INT", {"default": 80, "min": 1, "max": 10000}),
+                "learning_rate": ("FLOAT", {"default": 0.02, "min": 0.000001, "max": 1.0, "step": 0.001}),
+                "training_image_size": ("INT", {"default": 512, "min": 0, "max": 2048, "step": 64}),
+                "num_tokens": ("INT", {"default": 16, "min": 1, "max": 256}),
+                "init_std": ("FLOAT", {"default": 0.02, "min": 0.0, "max": 1.0, "step": 0.001}),
+                "training_dtype": (["bf16", "fp16", "fp32"], {"default": "bf16"}),
+            },
+            "optional": {
+                "reference_image_2": ("IMAGE",),
+                "reference_image_3": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ("ANIMA_REFERENCE_CONTEXT", "STRING")
+    RETURN_NAMES = ("context", "info")
+    FUNCTION = "train"
+    CATEGORY = "Anima/AnimaFastTrain"
+
+    def train(
+        self,
+        model: Any,
+        clip: Any,
+        vae: Any,
+        reference_image: Any,
+        training_prompt: str,
+        seed: int,
+        training_steps: int,
+        learning_rate: float,
+        training_image_size: int,
+        num_tokens: int,
+        init_std: float,
+        training_dtype: str,
+        reference_image_2: Any = None,
+        reference_image_3: Any = None,
+    ) -> tuple[AnimaReferenceContext, str]:
+        LOGGER.info(
+            "Train node received reference image: shape=%s dtype=%s seed=%s",
+            tuple(reference_image.shape),
+            getattr(reference_image, "dtype", None),
+            seed,
+        )
+        reference_images = _prepare_reference_images(
+            [reference_image, reference_image_2, reference_image_3],
+            size=training_image_size,
+        )
+        reference_count = reference_images.shape[0]
+        LOGGER.info(
+            "Prepared Anima reference image batch: count=%s shape=%s",
+            reference_count,
+            tuple(reference_images.shape),
+        )
+        initial_tokens = _create_reference_context_tokens(
+            reference_images,
+            num_blocks=ANIMA_NUM_BLOCKS,
+            num_tokens=num_tokens,
+            text_dim=ANIMA_TEXT_DIM,
+            dtype="float32",
+            std=init_std,
+        )
+        positive = _encode_prompt(clip, training_prompt)
+        optimized_tokens, losses = _optimize_reference_context_tokens(
+            model=model,
+            vae=vae,
+            positive=positive,
+            reference_images=reference_images,
+            initial_tokens=initial_tokens,
+            steps=training_steps,
+            learning_rate=learning_rate,
+            seed=seed,
+            training_dtype=training_dtype,
+        )
+        context = AnimaReferenceContext(
+            tokens=optimized_tokens,
+            losses=losses,
+            training_steps=training_steps,
+            learning_rate=learning_rate,
+            training_seed=seed,
+            reference_count=reference_count,
+            num_blocks=ANIMA_NUM_BLOCKS,
+            num_tokens=num_tokens,
+            text_dim=ANIMA_TEXT_DIM,
+        )
+        final_loss = losses[-1] if losses else 0.0
+        initial_norm = float(initial_tokens.detach().float().norm().cpu())
+        optimized_norm = float(optimized_tokens.detach().float().norm().cpu())
+        info = (
+            f"Anima context tokens trained in memory: steps={training_steps}, "
+            f"lr={learning_rate}, references={reference_count}, "
+            f"tokens={tuple(optimized_tokens.shape)}, "
+            f"final_loss={final_loss:.6f}, "
+            f"token_norm_before={initial_norm:.6f}, "
+            f"token_norm_after={optimized_norm:.6f}, "
+            f"token_norm_delta={optimized_norm - initial_norm:.6f}"
+        )
+        LOGGER.info(info)
+        return context, info
+
+
+class AnimaFastTrainPatchModel:
+    @classmethod
+    def INPUT_TYPES(cls) -> dict[str, Any]:
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "context": ("ANIMA_REFERENCE_CONTEXT",),
+                "reference_context_scale": ("FLOAT", {"default": 0.45, "min": 0.0, "max": 5.0, "step": 0.01}),
+                "reference_context_mode": (["positive", "all"], {"default": "positive"}),
+            }
+        }
+
+    RETURN_TYPES = ("MODEL",)
+    RETURN_NAMES = ("model",)
+    FUNCTION = "patch"
+    CATEGORY = "Anima/AnimaFastTrain"
+
+    def patch(
+        self,
+        model: Any,
+        context: AnimaReferenceContext,
+        reference_context_scale: float,
+        reference_context_mode: str,
+    ) -> tuple[Any]:
+        patched_model = model.clone()
+        previous_wrapper = patched_model.model_options.get("model_function_wrapper")
+        debug_state = {"block_calls": 0, "wrapper_calls": 0}
+        LOGGER.info(
+            "Patch node created AnimaFastTrain model wrapper: scale=%s mode=%s tokens_shape=%s",
+            reference_context_scale,
+            reference_context_mode,
+            tuple(context.tokens.shape),
+        )
+
+        def _context_wrapper(apply_model: Any, args: dict[str, Any]) -> Any:
+            # ComfyUI passes the real BaseModel.apply_model bound method here.
+            # Patch that bound object so the injection affects the active model
+            # used by the sampler, even if the ModelPatcher was cloned upstream.
+            active_model = getattr(apply_model, "__self__", None) or patched_model
+            context.application_count += 1
+            debug_state["wrapper_calls"] = debug_state.get("wrapper_calls", 0) + 1
+            if context.application_count <= 3 or context.application_count % 10 == 0:
+                LOGGER.info(
+                    "Applying AnimaFastTrain context: call=%s active_type=%s scale=%s mode=%s input_shape=%s",
+                    context.application_count,
+                    type(active_model).__name__,
+                    reference_context_scale,
+                    reference_context_mode,
+                    tuple(args["input"].shape),
+                )
+            restore = _apply_anima_context_injection(
+                active_model,
+                context.tokens,
+                scale=reference_context_scale,
+                mode=reference_context_mode,
+                debug_state=debug_state,
+            )
+            try:
+                if previous_wrapper is not None:
+                    return previous_wrapper(apply_model, args)
+                return apply_model(args["input"], args["timestep"], **args["c"])
+            finally:
+                restore()
+
+        patched_model.set_model_unet_function_wrapper(_context_wrapper)
+        return (patched_model,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "AnimaFastTrainReferenceContext": AnimaFastTrainReferenceContext,
+    "AnimaFastTrainPatchModel": AnimaFastTrainPatchModel,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "AnimaFastTrainReferenceContext": "AnimaFastTrain - Train Context Tokens",
+    "AnimaFastTrainPatchModel": "AnimaFastTrain - Patch Model",
+}