# ComfyUI AnimaFastTrain

Experimental custom nodes for training Anima reference context in ComfyUI with
in-memory reference context tokens.

## Installation

### ComfyUI Manager

1. Open ComfyUI Manager.
2. Select **Custom Nodes Manager**.
3. Search for **ComfyUI-AnimaFastTrain**.
4. Click **Install**.
5. Restart ComfyUI.

The nodes will appear under:

```text
Anima/AnimaFastTrain
```

### Manual Install

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.

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 - MaskTrain Context Tokens**
  - Inputs: `MODEL`, `CLIP`, `VAE`, a reference `IMAGE`, and a `MASK`.
  - Trains the same in-memory per-block context tokens as the regular train
    node, but weights the reconstruction loss by the mask in latent space.
  - White/1.0 mask areas contribute to training; black/0.0 areas are ignored.
  - Optional second and third reference images can use matching optional masks.
    If only one mask is connected, it is reused for every reference image.
  - This is mask-weighted training, not a runtime cross-attention mask.

- **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**, or **MaskTrain Context Tokens** if you want the
   training loss to focus on a masked region.
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.