Flux2TTRControllerTrainer

Flux2TTRControllerTrainer is Phase 2 of the ComfyUI-Taylor-Attention pipeline. Phase 1 (Flux2TTRTrainer) distills fast TTR attention layers; this node trains the controller that decides, at every denoising step, which layers run the fast path and which keep full native attention. Swap every layer and you lose quality. Swap none and you lose the whole point. The controller learns the middle ground.

How the controller learns

The mechanism is policy-gradient training with quality-driven rewards, which sounds scary but maps cleanly to what's on screen. The node runs dual-path sampling: a teacher path samples with the original Flux model, while a student path samples with the TTR-patched model under controller routing. The controller gets (sigma, CFG scale, latent width, height), embeds them, and outputs one routing logit per layer. The reward is −quality_loss − λ_eff × efficiency_penalty + λ_entropy × entropy_bonus. Quality comes from comparing teacher and student latents or images - RMSE and cosine distance by default, with LPIPS, DreamSim, HPS, BIQA, and Gemini scoring available if you raise their weights in the shared config node. The efficiency term penalizes drifting from your target TTR ratio (target_ttr_ratio in Flux2TTRTrainingParameters). Updates are REINFORCE over eligible (ready) layers only, with a quality-floor-clamped baseline so bad runs don't poison the estimate. sigma_aware_training (on by default) recomputes log-probs under a grad-enabled context to match how routing is actually used during denoising.

The inputs

The inputs are a full sampling rig, so wire it like a KSampler with a few extras:

• model_original - the base Flux model.
• model_ttr - your Phase-1-patched TTR model.
• positive, negative, latent, vae - the usual conditioning and latent inputs.
• training_config - required, and it must come from Flux2TTRTrainingParameters; this is where all the weights live.
• steps, seed, cfg, sampler_name, scheduler, denoise - the sampling settings for the teacher/student paths.
• training_iterations - default 100, your policy-update budget.
• checkpoint_path - defaults to ComfyUI/models/approximate_attention/flux2_ttr_controller.pt.

Outputs are the useful part: IMAGE_ORIGINAL and IMAGE_PATCHED let you eyeball teacher vs. student side by side, and CONTROLLER is the trained router you'll load at inference with Flux2TTRController.

Why this beats a fixed schedule

Why this beats hand-rolling a schedule: the pack's paper found routing isn't uniform across the denoising trajectory. Early high-sigma steps get TTR (attention is diffuse, cheap approximation is fine); the detail-critical late steps fall back to full attention. That pattern emerges from the reward shaping rather than being coded in - which is the interesting result, and also why you should let it train rather than trying to force it.

Common issues

Gotchas: this node needs both a base model and a Phase-1-patched model as separate MODEL inputs, and training_config is mandatory here (it's only optional in Phase-1). Don't turn on every quality scorer at once - each is a heavy dependency or, in Gemini's case, a paid API call per image pair. The installer handles pyiqa/hpsv2/dreamsim when you need them, but image-reward is deliberately excluded because it pins an old timm that breaks pyiqa. Install the pack via ComfyUI Manager (search "Taylor-Attention") or clone it into custom_nodes, then run:

uv pip install -e custom_nodes/ComfyUI-Taylor-Attention

The README's one-liner names the old folder ComfyUI-Approximate-Attention - use the folder you actually cloned. No model downloads; checkpoints go to ComfyUI/models/approximate_attention/. Everything here is experimental and community-unknown, so budget for iteration, not a one-shot win.