Post-Training Pruning for Diffusion Transformers #14

Is your feature request related to a problem? Please describe.

diffusers is an inference-first library and ships no mechanism for pruning Diffusion Transformers (DiTs). As DiTs like FLUX.1-dev grow large, there is no in-repo path to produce a sparser, smaller checkpoint that preserves image quality. The paper Post-Training Pruning for Diffusion Transformers (DiT-Pruning, arXiv:2607.00927) demonstrates this is achievable — 0.001 CLIP-score loss at 50% sparsity on FLUX.1-dev at 512×512 on MJHQ, with the advantage widening under high sparsity — but diffusers has no pruning module or host for such a method, and scripts/ is exclusively checkpoint-conversion utilities.

Describe the solution you'd like.

If pursued as a diffusers deliverable, target the smallest in-scope slice: a scripts/prune_flux.py utility that loads a FluxTransformer2DModel (FLUX.1-dev), applies an energy-based weight+activation saliency metric with clustering-aware granularity computed over a small MJHQ calibration set, and exports a sparse checkpoint via the existing save path. Evaluate at 50% sparsity against (a) the dense baseline and (b) an LLM-derived saliency baseline, reporting CLIP score and FID at 512×512. If the calibration/saliency infra is out of scope for core, scope this down to a checkpoint-conversion path that consumes a precomputed pruning mask.

The paper's concrete contributions, which a port would need to reproduce:

1. An energy-based saliency metric that jointly balances weight and activation contributions. The authors argue LLM-derived saliency metrics over-emphasize weight contribution, and that DiT weights have substantially larger magnitudes than LLM weights, so reusing LLM metrics causes large quality drops.
2. Clustering-aware pruning granularity, motivated by an observed distinct clustering pattern in the 2D weight space, which enables better sparse allocation than prior fixed granularities.

This is a method (an algorithm that computes which elements to prune and how to allocate sparsity), validated on multiple DiTs, not a drop-in checkpoint format diffusers could simply consume.

Describe alternatives you've considered.

• External / companion library. Novel pruning algorithms are arguably out of scope for diffusers core. Quantization is handled today by offloading to external libraries (bitsandbytes, quanto, torchao) through a thin diffusers/quantizers wrapper rather than hosting in-tree algorithms — pruning could follow the same optimum-style companion-library pattern.
• Pruned-checkpoint conversion path only. Rather than porting the saliency algorithm itself, consume a precomputed / externally-computed pruning mask applied to FLUX.1-dev and export the sparse checkpoint. This is the slice most likely to fit diffusers' existing scripts/ conversion conventions.

Additional context.

Why this is an Issue rather than a PR. No implementable subset or call site holds up against the repo layout: diffusers is inference-first with no pruning module, scripts/ is conversion-only, and the required calibration/saliency infrastructure (forward passes over MJHQ, energy-based saliency + clustering optimization) is training-adjacent with no existing scaffold. The most realistic implementation would land as a freestanding module that no pipeline or model code invokes — pipelines never call a pruner — so the scope decisions below should be settled before implementation.

Open questions to unblock this:

• Scope decision: is a novel pruning algorithm in scope for diffusers core, or should it live in an external / optimum-style companion library, consistent with how quantization is handled today?
• Intended host: if in scope, is this a scripts/prune_flux.py conversion-style utility, or a real diffusers/pruning module? Which existing code path calls it?
• Deliverable shape: should the contribution be a pruned-checkpoint conversion path (apply a precomputed mask and export) rather than porting the saliency algorithm itself?
• Assets & data: do we have access to the paper's released saliency code and the MJHQ calibration set, and is FLUX.1-dev the agreed target?

License & code availability: 🟡 No code repository surfaced — no LICENSE could be fetched to evaluate. Worth confirming the paper has an open release before investing in adoption (no repository/model URL surfaced; license (none detected), class no-code-link, compat 0.30).

Reopen this Issue if you want Outrider to revisit this paper later. While it stays closed, the orchestrator will not re-recommend the same paper.

Drafted by Outrider — paper: arXiv:2607.00927.