Flux2TTRTrainer

Flux2TTRTrainer is Phase 1 of the ComfyUI-Taylor-Attention project: it distills the fast "TTR" attention layers from native Flux attention, and it does it online, in the middle of normal sampling. You don't launch a separate training script. You build a normal-looking sampling graph with this node in the chain, and the training happens in the background of the images you're already generating. The catch: it's a research node, it's slow, and a useful run means leaving it on for a while.

Why this exists

Why the whole project exists: Flux's attention is quadratic in token count, which is why big resolutions choke consumer GPUs. This node trains a linear-time kernel-regression replacement for those layers - not a new checkpoint, not a LoRA, just swapped attention internals - and the pack's paper reports a few hours on a single RTX 4090 is enough to get roughly half of Flux's layers distillable.

How the distillation works

The mechanism is worth understanding because it's genuinely careful. The teacher is native Flux attention computed on real prompts and sampling states; the student is the TTR layer, which fuses a linear-time kernel branch with a small exact-softmax branch over "landmark" tokens, blended by a bounded weight. Samples pile into a replay buffer that keeps full keys and values but subsamples queries, so the student learns the same regression problem the teacher solved without the storage bill (CPU offload is on by default, budgeted to 768 MB). Per-layer readiness gating is fail-closed: a layer doesn't serve the student until it has accumulated enough updates and its EMA cosine distance to the teacher is low, with hysteresis so it doesn't flap at the boundary. Untrained layers silently fall back to teacher attention instead of wrecking your images. The loss is a Kendall-style uncertainty-weighted blend of huber and cosine, and the checkpoint persists per-layer optimizer state so you can resume across runs.

Inputs worth setting

Inputs worth setting on a first run:

• model, latents, conditioning - the standard sampling inputs.
• steps - default 512. This is your training budget per run.
• checkpoint_path - defaults to ComfyUI/models/approximate_attention/flux2_ttr.pt.
• feature_dim - default 256, must be a multiple of 256.

Everything else - chunk sizes, landmark fraction/min/max, replay knobs, readiness threshold, layer_start/layer_end - has sane defaults. training toggles distillation on; training_preview_ttr makes it output the student for visual comparison. Outputs are a patched MODEL (feed it to a KSampler) and loss_value.

Common issues

The honest gotchas: this is real training cost, not a slider trick. Keep an eye on VRAM - replay offloads to CPU by default, and enable_memory_reserve is off precisely because flipping it on can make ComfyUI evict memory aggressively. The pack requires a recent ComfyUI with the v3 node API and heavy Python deps. ## Installing

Install via ComfyUI Manager (search "Taylor-Attention") or:

cd ComfyUI/custom_nodes
git clone https://github.com/ttulttul/ComfyUI-Taylor-Attention
uv pip install -e custom_nodes/ComfyUI-Taylor-Attention

(README's one-liner says ComfyUI-Approximate-Attention - the old folder name; use the one you cloned.) No model files are downloaded; checkpoints land in ComfyUI/models/approximate_attention/. And don't expect tutorials - this pack has almost no community footprint. It's research code, well-engineered for what it is, and if you're not prepared to babysit a multi-hour distillation run, start with the controller article's opposite number and come back when you've got time.