Flux2TTRTrainer

Phase-1: distill Flux TTR linear attention modules from native attention.

By ttulttul·Created 6 months ago·Updated 6 months ago· 1

Flux2TTRTrainer

model
latents
conditioning
training_config

MODEL
loss_value

◄steps512►

◄trainingtrue►

◄training_preview_ttrtrue►

◄checkpoint_path/tmp/ComfyUI/models/approximate_attention/flux2_ttr.pt►

◄feature_dim256►

◄query_chunk_size256►

◄key_chunk_size1024►

◄landmark_fraction0.080►

◄landmark_min64►

◄landmark_max0►

◄text_tokens_guess77►

◄alpha_init0.100►

◄training_query_token_cap128►

◄replay_buffer_size8►

◄replay_offload_cputrue►

◄replay_max_mb768►

◄train_steps_per_call1►

◄readiness_threshold0.1200►

◄readiness_min_updates24►

◄enable_memory_reservefalse►

◄layer_start-1►

◄layer_end-1►

◄cfg_scale1.000►

◄min_swap_layers1►

◄max_swap_layers-1►

◄inference_mixed_precisiontrue►

◄controller_checkpoint_path►

Categoryadvanced/attention

Inputs (31)

Name	Type	Default	Description
model	MODEL		—
latents	LATENT		—
conditioning	CONDITIONING		—
steps	INT	5120–200000	—
training	BOOLEAN	true	Train TTR layers by distillation when enabled.
training_preview_ttr	BOOLEAN	true	When training, output TTR student attention for visual preview instead of teacher passthrough.
checkpoint_path	STRING	/tmp/ComfyUI/models/approximate_attention/flux2_ttr.pt	Checkpoint file to load/save TTR layer weights (defaults to ComfyUI/models/approximate_attention/flux2_ttr.pt).
feature_dim	INT	256128–8192	Kernel feature dimension (must be a multiple of 256).
query_chunk_size	INT	2561–4096	Query chunk size for kernel attention evaluation.
key_chunk_size	INT	10241–8192	Key chunk size for kernel KV/Ksum accumulation.
landmark_fraction	FLOAT	0.0800.01–0.5	Fraction of image tokens used as landmarks for exact softmax residual.
landmark_min	INT	641–1024	Minimum landmark count used at low resolution.
landmark_max	INT	00–2048	Maximum landmark count used at high resolution (0 means unlimited).
text_tokens_guess	INT	770–1024	Assumed number of text tokens at the start of sequence for landmark selection.
alpha_init	FLOAT	0.1000–10	Initial residual gate for landmark softmax branch.
training_query_token_cap	INT	1281–4096	Max number of query tokens per replay sample; keys/values always stay full length.
replay_buffer_size	INT	81–4096	Replay buffer capacity per layer for distillation samples.
replay_offload_cpu	BOOLEAN	true	Store replay samples on CPU (reduced precision) to reduce VRAM pressure.
replay_max_mb	INT	76864–65536	Global replay memory budget in MB across all layers.
train_steps_per_call	INT	11–32	Number of replay optimization steps run per attention call.
readiness_threshold	FLOAT	0.12000–10	Enable student inference for a layer only when EMA loss is below this threshold.
readiness_min_updates	INT	240–100000	Minimum replay updates before a layer can be marked ready.
enable_memory_reserve	BOOLEAN	false	Call ComfyUI free_memory before HKR attention allocations (can offload aggressively).
layer_start	INT	-1-1–512	Only apply TTR to single blocks with index >= layer_start (-1 disables).
layer_end	INT	-1-1–512	Only apply TTR to single blocks with index <= layer_end (-1 disables).
cfg_scale	FLOAT	1.0000–100	Default CFG scale used when transformer_options does not provide guidance scale.
min_swap_layers	INT	10–512	Training mode: minimum number of eligible single layers to swap to TTR per diffusion step.
max_swap_layers	INT	-1-1–512	Training mode: maximum swapped layers per diffusion step (-1 means all eligible layers).
inference_mixed_precision	BOOLEAN	true	Use input dtype (bf16/fp16) for TTR inference on CUDA for speed.
controller_checkpoint_path	STRING		Optional Phase-2 controller checkpoint used for inference-time layer routing.
training_configopt	TTR_TRAINING_CONFIG		—

Outputs (2)

Name	Type	Description
MODEL	MODEL	—
loss_value	FLOAT	—