HRDiT-HAP

Derivation of the per-head attention scopes used by HAP (Head-adaptive Attention Pruning) in HRDiT (ECCV 2026).

HAP gives every attention head its own attention scope — the range of image tokens that head is allowed to look at — and skips the attention outside it. Which scope each head gets is decided once, offline, and stored in a scope plan: a JSON file of per-head alphas and betas.

HRDiT generates images from a ready-made attention plan; this repository produces one. Run it for a different attention budget — tighter skips more attention and generates faster, looser skips less.

Producing a plan takes two stages:

1. Profiling — run the model on a small calibration set and record, for every head and every candidate scope, how much the attention outside that scope contributes to the output.
2. Plan search — solve a multiple-choice knapsack problem that assigns one scope per head, minimising the importance lost subject to an attention budget.

🔧 Installation

conda create -n hrdit-hap python=3.10 -y
conda activate hrdit-hap

pip install torch==2.7.0 torchvision --index-url https://download.pytorch.org/whl/cu126
pip install -r requirements.txt

Download the off-the-shelf base model:

huggingface-cli download black-forest-labs/FLUX.1-dev --local-dir ./pretrained_models/FLUX.1-dev

📐 Deriving a scope plan

1. Prepare a calibration set

A JSON list of image/caption pairs. Two dozen pairs is enough; the images should be at least 1024x1024.

[
  {"image": "/path/to/000.jpg", "caption": "a description of that image"},
  {"image": "/path/to/001.jpg", "caption": "a description of that image"}
]

2. Profile

Each run writes one shard; split the calibration set across the GPUs you have. With 24 pairs on 4 GPUs:

for g in 0 1 2 3; do
    CUDA_VISIBLE_DEVICES=$g python scripts/flux_pipeline/flux_profile_local.py \
        --grad_dir ./profile/shard$g \
        --model_path ./pretrained_models/FLUX.1-dev \
        --calib_json ./calibration.json \
        --max_length 2048 \
        --data_offset $((g * 6)) --num_data 6 &
done
wait

Profiling backpropagates through the attention matrices, so it needs more memory than inference. Lower --num_data if a shard runs out of memory; shards are averaged by sample count, so splitting them differently does not change the result.

3. Search the plan

python scripts/flux_pipeline/flux_elastic_generate_local.py \
    --shard_root ./profile \
    --output_dir ./plans \
    --elastic_res 2048 --extend_res 4096 \
    --density_bounds 0.25 0.25 \
    --min_width 2048 \
    --plan_name scope_plan_flux.json

The solver uses gurobipy when it is installed and licensed for the problem size, and otherwise falls back to a Lagrangian subgradient method. It prints the realized attention density at both lengths.

4. Use the plan

cp ./plans/scope_plan_flux.json /path/to/HRDiT/configs/
cd /path/to/HRDiT
python inference.py --prompt "..." --scope_plan configs/scope_plan_flux.json

Key arguments

flux_profile_local.py

argument	default	meaning
--grad_dir	required	where this shard's importance tensor is written
--calib_json	—	calibration pairs; alternatively --ref_dir with --meta_json
--max_length	2048	resolution the profiling runs at
--data_offset / --num_data	0 / 6	slice of the calibration set for this shard
--aggregating_block_size	64	block size the importance is aggregated over

flux_elastic_generate_local.py

argument	default	meaning
--shard_root	required	directory holding the profiling shards
--elastic_res / --extend_res	2048 / 4096	the two resolutions the budget is enforced at
--density_bounds	0.1 0.1	attention budget at those two resolutions
--min_width	512	per-head band-width floor in tokens
--num_alphas / --num_betas	8 / 9	size of the candidate scope grid

Acknowledgements

The code is built upon MoA, thanks for their work!

Citation

@misc{xue2026hrdittrainingfreehighresolutionimage,
      title={HRDiT: Training-Free High-Resolution Image Generation with Off-the-Shelf Diffusion Transformer Models}, 
      author={Yu Xue and Haoxuan Qu and Zhuoling Li and Hongbin Xu and Jianxiong Yin and Simon See and Hossein Rahmani and Jun Liu},
      year={2026},
      eprint={2608.07003},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2608.07003}, 
}