FlexControl: Computation-Aware ControlNet with Differentiable Router for Text-to-Image Generation

TL;DR

FlexControl introduces a data-driven, computation-aware routing mechanism that dynamically activates control blocks in diffusion-based text-to-image generation. By copying all backbone blocks and gating them with a differentiable router, it eliminates manual block selection and reduces redundant computation while maintaining or improving image fidelity and alignment with prompts. The method optimizes a joint objective that combines standard diffusion loss with a FLOPs-aware cost to encourage efficient use of control blocks, achieving favorable Pareto frontiers versus static ControlNet variants on UNet and DiT backbones. Empirical results across depth maps, edges, and segmentation cues show consistent gains in FID, CLIP alignment, and controllability, with qualitative evidence of sharper details and better structure preservation. This work enables scalable, task-adaptive controllable diffusion with minimal architectural changes, facilitating practical deployment in diverse generative pipelines.

Abstract

ControlNet offers a powerful way to guide diffusion-based generative models, yet most implementations rely on ad-hoc heuristics to choose which network blocks to control-an approach that varies unpredictably with different tasks. To address this gap, we propose FlexControl, a novel framework that copies all diffusion blocks during training and employs a trainable gating mechanism to dynamically select which blocks to activate at each denoising step. With introducing a computation-aware loss, we can encourage control blocks only to activate when it benefit the generation quality. By eliminating manual block selection, FlexControl enhances adaptability across diverse tasks and streamlines the design pipeline, with computation-aware training loss in an end-to-end training manner. Through comprehensive experiments on both UNet (e.g., SD1.5) and DiT (e.g., SD3.0), we show that our method outperforms existing ControlNet variants in certain key aspects of interest. As evidenced by both quantitative and qualitative evaluations, FlexControl preserves or enhances image fidelity while also reducing computational overhead by selectively activating the most relevant blocks. These results underscore the potential of a flexible, data-driven approach for controlled diffusion and open new avenues for efficient generative model design. The code will soon be available at https://github.com/Anonymousuuser/FlexControl.

Figures (10)

• Figure 1: Overview of dynamic routing guided by the router unit. (a) In the training stage, Gumbel noise is added to the discrete mask to assist the gradient backpropagation. (b) In the inference stage, the router unit controls whether to activate the control block and whether to inject conditional control into the frozen block of the backbone according to the input latent variable. Once output the instruction of inactive, the corresponding control block and zero module will be skipped.
• Figure 2: Qualitative comparison of controllable generation methods. FlexControl achieves higher fidelity and structure preservation across Depth Map, Canny Edge, and Segmentation Mask conditions, reducing distortions (boxes) seen in other methods. It better aligns with input conditions while maintaining visual quality.
• Figure 3: Comparison of FlexControl and existing methods on SD1.5 for semantic consistency. FlexControl achieves better semantic alignment and structure preservation with varying sparsity levels, while ControlNet-based methods show inconsistencies in segmentation accuracy (highlighted in yellow boxes). Captions: A stone building surrounded by a stone wall and a grassy lawn.
• Figure 4: Comparison of FlexControl and existing methods on SD3.0 for edge preservation. FlexControl maintains better spatial consistency and object integrity across different sparsity levels, while ControlNet-based methods introduce distortions and inconsistencies (highlighted in red boxes). Captions: A room with large windows, a gray sofa, a table, and a TV stand.
• Figure A1: The distribution of activated control blocks under different timesteps. The hyperparameter $\gamma$ is set to 0.5 to approximate 50% sparsity, and timestep is set to 10, 20 and 50 repectively. The first line shows the results of the model based on SD1.5, and the second line shows the results of the model based on SD3.0. $\large \IfNoValueTF{HTML}{ }{ }$ and $\large \IfNoValueTF{HTML}{ }{ }$ denotes activated and inactivated blocks, respectively.