FastFLUX: Pruning FLUX with Block-wise Replacement and Sandwich Training
This work addresses deployability issues for users of large text-to-image models, though it is incremental as it builds on existing pruning and fine-tuning techniques.
The paper tackles the problem of slow inference and high memory usage in large text-to-image models like FLUX by proposing FastFLUX, a pruning framework that replaces complex residual branches with linear layers and uses localized fine-tuning, resulting in maintained image quality and significantly improved inference speed with 20% of the hierarchy pruned.
Recent advancements in text-to-image (T2I) generation have led to the emergence of highly expressive models such as diffusion transformers (DiTs), exemplified by FLUX. However, their massive parameter sizes lead to slow inference, high memory usage, and poor deployability. Existing acceleration methods (e.g., single-step distillation and attention pruning) often suffer from significant performance degradation and incur substantial training costs. To address these limitations, we propose FastFLUX, an architecture-level pruning framework designed to enhance the inference efficiency of FLUX. At its core is the Block-wise Replacement with Linear Layers (BRLL) method, which replaces structurally complex residual branches in ResBlocks with lightweight linear layers while preserving the original shortcut connections for stability. Furthermore, we introduce Sandwich Training (ST), a localized fine-tuning strategy that leverages LoRA to supervise neighboring blocks, mitigating performance drops caused by structural replacement. Experiments show that our FastFLUX maintains high image quality under both qualitative and quantitative evaluations, while significantly improving inference speed, even with 20\% of the hierarchy pruned. Our code will be available soon.