Flux2TTRController

This is the node that finally does the thing the whole pack is about: makes Flux faster at high resolution without you babysitting which layers get replaced. By the time you're here, you've either trained the pack's Phase-1 and Phase-2 checkpoints yourself or found them somewhere - and this node is where they actually get used.

Quick context, because the pack name is misleading. Flux's softmax attention scales quadratically with token count, which is the wall you hit at 2048×2048 and beyond. The ComfyUI-Taylor-Attention project (note: the "Taylor" part is history - the v1 Taylor-expansion approach failed and v2 is hybrid kernel-regression attention) distills fast linear-time replacements for roughly half of Flux's attention layers, then trains a controller to decide which layers swap and when. This node is the inference end: it loads the two checkpoints, patches your model, and hands you a MODEL to feed into a KSampler.

How the routing works

The mechanism is a per-step routing decision. The controller reads sigma, CFG scale, and latent width/height, then predicts one routing logit per layer: high logit means full (native) attention, low logit means the TTR fast path. Because it re-decides every denoising step, it can run aggressive TTR during the early noisy steps and fall back to full attention when fine detail matters at the end - the pattern the pack's paper found actually holds up.

The inputs that matter

The inputs that matter:

• model - your Flux model, any of the usual quantized variants.
• ttr_checkpoint_path and controller_checkpoint_path - both required, both validated against the filesystem. Missing or empty path raises. The defaults point at ComfyUI/models/approximate_attention/flux2_ttr.pt and flux2_ttr_controller.pt.
• quality_speed (0–1, default 0.5) - mapped internally to a controller threshold of 0.1 + 0.8 × quality_speed. Lower keeps more layers on full attention (better quality, slower); higher routes more to TTR (faster, riskier). The default is the honest middle.
• policy_mode - stochastic (default) samples one routing mask per diffusion step, matching how the controller was trained; threshold uses a deterministic cutoff.
• policy_temperature - only matters in stochastic mode; lower makes decisions harder, higher makes them noisier.

Output is a single patched MODEL, which goes straight into your KSampler.

Common issues

Where people get burned: this node is useless without both checkpoints. A Phase-1-only setup gives you distilled layers but no controller, and the node refuses to run - it validates the controller file exists up front. Also note feature_dim is read from the TTR checkpoint's own metadata, not from a UI field, so a mismatched config can't silently break the load. And the quality/speed knob is a real trade, not a free lunch - crank quality_speed to 1 and you'll see the kernel-attention softening (less fine texture) the paper admits to.

Installing

Install via ComfyUI Manager (search "Taylor-Attention") or clone into custom_nodes, restart, then install dependencies - this is a research pack with genuine requirements and it needs a recent ComfyUI with the v3 node API:

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

The README's one-liner still says ComfyUI-Approximate-Attention; use the folder you cloned. No model files to download - the checkpoints are trained by you and land in ComfyUI/models/approximate_attention/ (auto-created).