Source-linked AI summary
OminiControl2: Efficient Conditioning for Diffusion Transformers
Zhenxiong Tan, Qiaochu Xue, Xingyi Yang, Songhua Liu, Xinchao Wang
TL;DR
Fine-grained DiT control is hindered by the computational cost of long and repeatedly processed conditional inputs. OminiControl2 compresses condition tokens and reuses conditional features across denoising steps, reporting over 90% lower conditional-processing overhead and a 5.9× speedup in multi-conditional generation.
Problem
Long or multiple condition-token sequences make DiT control computationally expensive, while condition tokens are repeatedly processed across denoising steps.
Method
OminiControl2 combines compact condition-token representations with conditional feature reuse across denoising steps.
Results
Over 90% lower conditional-processing overhead and up to 5.9× speedup are reported across multi-conditional generation scenarios.
Takeaways & Limitations
The reported efficiency gains support practical complex control scenarios with substantially lower computational demands.
Takeaways & Limitations
Naively keeping condition features static degrades output quality because subtle condition-feature changes during denoising create a training-inference discrepancy.
Abstract
from arXiv · showhide
Fine-grained control of text-to-image diffusion transformer models (DiT) remains a critical challenge for practical deployment. While recent advances such as OminiControl and others have enabled a controllable generation of diverse control signals, these methods face significant computational inefficiency when handling long conditional inputs. We present OminiControl2, an efficient framework that achieves efficient image-conditional image generation. OminiControl2 introduces two key innovations: (1) a dynamic compression strategy that streamlines conditional inputs by preserving only the most semantically relevant tokens during generation, and (2) a conditional feature reuse mechanism that computes condition token features only once and reuses them across denoising steps. These architectural improvements preserve the original framework's parameter efficiency and multi-modal versatility while dramatically reducing computational costs. Our experiments demonstrate that OminiControl2 reduces conditional processing overhead by over 90% compared to its predecessor, achieving an overall 5.9$\times$ speedup in multi-conditional generation scenarios. This efficiency enables the practical implementation of complex, multi-modal control for high-quality image synthesis with DiT models.
1. Introduction
OminiControl2 targets the prohibitive cost of long, multi-conditional DiT inputs with compact token representations and conditional feature reuse. It preserves OminiControl’s parameter efficiency and versatile control while reducing overhead and accelerating generation.
- OminiControl’s concatenation of control and noisy-image tokens enables multimodal attention and avoids auxiliary modules by reusing the pretrained VAE.This design reduces total parameter count while supporting diverse control signals.
- OminiControl’s long unified sequences become prohibitively expensive for extensive or multiple conditions because attention scales quadratically with sequence length.This computational bottleneck limits practical complex, multi-conditional use.
- OminiControl2 selectively retains informative condition tokens, using localized regions for partial edits and downsampled tokens for full-image generation.These strategies reduce the number of tokens processed during DiT inference while retaining essential conditioning information.
- Conditional token features are computed once initially and reused across later denoising steps, while generated-image token features continue updating.This removes redundant conditional computation across inference steps.
- Over 90% lower conditional-processing overhead and 5.9× faster multi-conditional generation are reported while maintaining generation quality.The framework also retains OminiControl’s parameter efficiency and universal control capabilities.
2. Related Works
Prior controllable diffusion methods largely target UNet architectures, while newer DiT approaches leverage multimodal attention but face growing costs as condition tokens increase. Separate DiT acceleration work explores pruning and caching to reduce iterative computation.
- UNet-based image-control methods cannot be directly applied to DiTs because the architectures and operational paradigms differ fundamentally.
- DiT control methods such as OminiControl and DSD use existing multimodal attention to incorporate image conditions without complex architectural modifications.Their limitation is associated with increasing condition-token counts, especially in multi-conditional settings.
- Diffusion-transformer acceleration research includes model pruning, which dynamically skips or removes computational components during inference.Examples cited include DiP-GO and TinyFusion.
- The related acceleration literature also considers computation caching to reduce repeated intermediate processing in diffusion transformers.
3. Methods
OminiControl2 makes controllable DiT generation more efficient by compressing condition tokens and reusing their features across denoising steps. It preserves spatial correspondence and addresses the computational redundancy of long, multi-condition sequences.
- 3.1. Preliminaries: OminiControl extends DiT sequences with VAE-encoded image condition tokens, enabling flexible interactions for diverse image-generation tasks.The unified sequence combines noisy, text, and image condition tokens for multimodal attention.
- 3.1. Preliminaries: Multiple image conditions increase token counts and computational cost, while attention becomes expensive for long unified sequences.The cost grows particularly in multi-conditional settings because self-attention operates over the combined sequence.
- 3.2. Compact token representation: Compact representation reduces condition-token processing through spatial compression, position correction, pruning, and inpainting-specific token integration.Compression reduces spatial dimensions before encoding; position correction preserves correspondence, pruning removes minimally informative tokens, and inpainting combines noisy and condition tokens based on the mask.
- 3.2. Compact token representation: Position correction restores spatial alignment after compression, preserving conditioning fidelity for spatially aligned generation.Without this mapping, compressed tokens no longer correspond directly to target regions and can cause structural inconsistencies.
- 3.3. Feature reuse in DiT: Condition features remain highly similar across denoising steps, motivating reuse, whereas noisy-image features change substantially.The similarity analysis identifies redundant recomputation of condition features during sampling.
- 3.3. Feature reuse in DiT: Naive feature reuse degrades quality because fully static condition features omit subtle timestep-dependent changes.The resulting training-inference discrepancy shows that similarity alone does not justify unrestricted reuse.
- 3.4. Computational complexity analysis: 112 theoretical speedup is predicted for condition-token processing when n = 28 and r = 0.25, while practical overhead is approximated as O(n · |C|).The analysis separates token-count reduction from repeated condition-feature computation; practical costs are dominated by token-independent operations for 1–4 conditions.
4. Experiments
Experiments evaluate OminiControl2 across conditional generation tasks, multi-condition settings, efficiency configurations, and quality-preservation comparisons. The combined optimizations reduce computational overhead while maintaining competitive generation quality, with benefits increasing for more conditions and higher resolutions.
- Experimental setup: The evaluation covers Canny-to-image, depth-to-image, mask inpainting, deblurring, and simultaneous application of all four condition types.Experiments use FLUX.1 with LoRA fine-tuning and assess latency, condition overhead, FID, CLIP, NIQE, and MUSIQ.
- Efficiency: 3.8×–5.9× speedups over original OminiControl are achieved across configurations, with larger gains at higher resolutions and more conditions.The combined approach provides maximum efficiency, while compact token representation and feature reuse contribute differently depending on sampling-step settings.
- Qualitative comparison: Consistent image quality and up to 3.85× speedup are maintained as conditions increase from 1 to 4, unlike naive caching and token merging.Baseline methods show quality degradation and visible artifacts with four conditions, while token merging also incurs grouping overhead and positional-embedding constraints.
- Quantitative evaluation: 96.2% conditional-overhead reduction, from 24.38s to 0.93s, is achieved for multi-condition tasks while retaining competitive FID scores.Across conditioning tasks, OminiControl2 better preserves generation quality than the baseline while reducing computational overhead.
- Position correcting for compressed tokens: Position correcting produces significantly better outputs by preserving spatial correspondence between compressed condition tokens and generated images.Without position correction, generation exhibits structural inconsistencies.
- Ablation studies: Feature reuse alone reduces overhead by 84.7% with minimal quality degradation, while compact token representation alone reduces overhead by approximately 65%.For inpainting, the full method reduces FID from 12.07 to 11.71 while cutting overhead by 90.1%.
5. Conclusion
OminiControl2 combines compact token representation with conditional feature reuse for controllable image generation with Diffusion Transformers. Across control tasks, it achieves speedups of up to 5.9×, addressing computational bottlenecks in multi-conditional generation and supporting more practical deployment.
- 5. Conclusion: OminiControl2 achieves speedups of up to 5.9× across control tasks by combining compact token representation with conditional feature reuse.The optimizations eliminate redundancies in the diffusion process and address computational bottlenecks in multi-conditional generation.