IFControlNet: Spatially Controlled Diffusion

• IFControlNet is a conditional generative diffusion extension that enforces spatial fidelity and reconstructs missing image details using intermediate feature alignment.
• It integrates lightweight auxiliary control branches and convolutional probes into pretrained latent diffusion models, ensuring precise spatial control and artifact suppression.
• Applied to multi-focus image fusion, IFControlNet restores fine details and enhances overall image quality, demonstrating superior performance in key metrics.

IFControlNet is a conditional generative diffusion extension designed to enforce spatial fidelity and reconstruct missing image content by leveraging intermediate feature alignment during the denoising process. It augments pretrained latent diffusion models (e.g., Stable Diffusion) with lightweight auxiliary control branches and convolutional probes, ensuring alignment between generated outputs and external spatial conditions or intermediate restoration targets. IFControlNet demonstrates substantial improvements in tasks requiring precise spatial control, notably within multi-focus image fusion, where it refines all-in-focus images by restoring lost details and suppressing artifacts.

1. Core Architectural Principles

IFControlNet builds upon the latent diffusion backbone (e.g., Stable Diffusion 2.1-base), integrating auxiliary mechanisms for conditional guidance:

• VAE Encoder/Decoder: A frozen variational autoencoder maps input images $x$ to latent codes $z=E(x)$ (dimension $64 \times 64 \times 4$ for $512 \times 512$ images) and reconstructs outputs via $D$.
• ControlNet Branch ($f_\phi$): A lightweight U-Net operating in parallel, accepting noisy latents $z_t$, conditional latents $c_\text{IF}$ derived from the initial fused image, and a time embedding for each step.
• Latent Diffusion U-Net ($\epsilon_\theta$): The original (frozen) backbone, except at injection points where predicted residuals $\Delta z_t$ from $f_\phi$ are added element-wise at each denoising block.
• Sampler: A DDIM/DDPM sampling process produces progressively denoised latents.

At each denoising stage, IFControlNet introduces a residual correction $\Delta z_t$ via $f_\phi(z_t, c_\text{IF}, t)$, augmenting the backbone without disrupting its generative prior. This injects structural priors from the initial fused image directly into the sampling trajectory, steering generation toward desired content and spatial alignment (Xie et al., 25 Dec 2025).

Additionally, lightweight timestep-conditioned convolutional probes extract intermediate decoder features within the UNet, reconstructing external controls (e.g., edges, depth maps) from noisy latents at every denoising step. This enables efficient alignment feedback throughout the diffusion process (Konovalova et al., 3 Jul 2025).

2. Diffusion and Conditioning Formulation

The latent diffusion process employs standard DDPM-style forward noising and reverse denoising, mathematically defined as:

• Forward (noising):

$$q(z_t | z_{t-1}) = \mathcal{N}(z_t; \sqrt{1-\beta_t}z_{t-1}, \beta_t I)$$

with cumulative $\bar\alpha_t = \prod_{s=1}^t (1-\beta_s)$. Sampling: $$z_t = \sqrt{\bar\alpha_t} z_0 + \sqrt{1-\bar\alpha_t} \epsilon$$, $\epsilon \sim \mathcal{N}(0, I)$.

• Reverse (denoising):

$$p_\theta(z_{t-1}|z_t) = \mathcal{N}\left(z_{t-1};\, \mu_\theta(z_t, t), \Sigma_\theta(t)\right)$$

$$\mu_\theta(z_t, t) = (z_t - \sqrt{1-\alpha_t}\, \epsilon_\theta(z_t, t)) / \sqrt{\alpha_t}$$.

• IFControlNet Conditional Injection:

$\tilde{z}_t = z_t + f_\phi(z_t, c_\text{IF}, t)$

$$z_{t-1} = \frac{1}{\sqrt{\alpha_t}}\left( \tilde{z}_t - \frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}} \epsilon_\theta(\tilde{z}_t, t) \right) + \sigma_t \epsilon'$$

($\sigma_t = 0$ for DDIM sampling).

• Intermediate Probe Alignment (InnerControl):

For decoder feature $h_i(t)$ at timestep $t$,

$\hat{c}_i(t) = f_i(h_i(t); \phi_i, t)$

Alignment enforced across all layers and steps:

$$L_\text{align} = \sum_{t=0}^T \sum_{i \in \text{Layers}} \lambda_i \cdot \|f_i(h_i(t)) - c_\text{spatial}\|_2^2$$

(Konovalova et al., 3 Jul 2025).

3. Training Strategies and Loss Functions

IFControlNet utilizes multiple training objectives to optimize spatial control and image quality:

• Conditional Denoising Loss ($L_\text{simple}$):

$$L_\text{simple}(\phi) = \mathbb{E}_{z_0, t, \epsilon} [\, \| \epsilon - \epsilon_\theta(z_t + f_\phi(z_t, c_\text{IF}, t), t) \|_2^2 \,]$$

• InnerControl Alignment Loss: Enforces signal reconstruction via probes from UNet features at every diffusion step.
• Cycle-Consistency Reward Loss (optional, ControlNet++ style):

For single-step reconstruction $x_0'$ and reward model $D$,

$$L_\text{reward} = \| D(x_0') - c_\text{spatial} \|^2$$

Applied only for $t \leq t_\text{reward threshold}$ (e.g., $t \leq 200$ for edge, $t \leq 400$ for depth).

• Combined Objective:

$$L_\text{total} = L_\text{diffusion} + \alpha L_\text{reward} + \beta L_\text{align}$$

• Optimization Details:
  • AdamW (learning rates $1{\times}10^{-5}$ / $3{\times}10^{-6}$), batch sizes 8–256.
  • ControlNet and probe weights updated; backbone and VAE weights frozen.
  • Probe architectures use conv-bottleneck layers and timestep embeddings; self-attention for depth guidance (Konovalova et al., 3 Jul 2025, Xie et al., 25 Dec 2025).

4. Application to Multi-Focus Image Fusion

Within the GMFF (Generative Multi-Focus Fusion Framework) pipeline, IFControlNet refines initial outputs from deterministic fusion models (e.g., StackMFF V4):

• Stage 1 (Deterministic Fusion): StackMFF V4 combines available focal plane images into an initial all-in-focus image $I_\text{fused}$.
• Stage 2 (Generative Restoration via IFControlNet):
  • $I_\text{fused}$ encoded to latent $c_\text{IF}$, used as the conditional input for IFControlNet.
  • The generative branch restores fine details, reconstructs missing regions (e.g., where no input image is truly in focus), and suppresses edge artifacts from hard-selection and uncertain estimation inherent in deterministic fusion.
  • Cross-attention layers align spatial features between $c_\text{IF}$ and noisy latents, steering denoising toward realistic completions (Xie et al., 25 Dec 2025).

Experiments use synthetic stacks generated from datasets including DUTS, NYU Depth V2, DIODE, Cityscapes, ADE20K. Variable levels (0–50%) of missing focal planes simulate incomplete data scenarios.

5. Experimental Evaluation and Results

Evaluation of IFControlNet covers both spatial controllability (edge, depth, fusion) and restoration fidelity:

Metric	ControlNet v1.1	ControlNet++	Ctrl-U	IFControlNet
Depth RMSE↓	35.90	28.32	29.06	26.09
HED SSIM↑ / FID↓	– / –	0.8097/15.01	–	0.8207/13.27
LineArt SSIM↑/FID↓	– / –	0.8399/13.88	–	0.8258/12.08

• Control fidelity (SSIM for edges, RMSE for depth) shows IFControlNet’s consistent superiority over prior methods, especially in scenarios with missing or noisy inputs.
• Image quality (FID) is maintained or slightly improved, with no adverse trade-off from increased control regularization.
• Prompt relevance (CLIP-score $\approx32$) remains unchanged.
• Fusion perceptual quality (BRISQUE, PIQE): GMFF (StackMFF V4 + IFControlNet) achieves substantial reductions on Mobile Depth (BRISQUE 14.98$\rightarrow$9.20, PIQE 28.00$\rightarrow$27.25), Middlebury (BRISQUE 25.87$\rightarrow$13.67, PIQE 44.28$\rightarrow$29.35), outperforming previous deblurring and fusion models.

Qualitative inspections reveal:

• Edge-artifact suppression at focus boundaries
• Hallucination of missing details in regions with no sharp source input
• Refinement of textural and micro-structural features (e.g., serrations, background tiles)
• Robust performance when applied to initial outputs from other fusion models, confirming decoupled applicability.

6. Implementation, Efficiency, and Practical Considerations

• Model Size: IFControlNet totals $\sim6.35$ billion parameters (diffusion + ControlNet).
• Computational Cost: $\sim493$ G FLOPs per $512\times512$ image; inference $17.6$ seconds/image (A6000 GPU).
• Training Budget: Typical runs require 8 H100 GPUs ($\sim$6 hours), or 2 A6000s ($\sim$16 hours).
• Initialization: Control branch $\phi$ is often initialized from IRControlNet (DiffBIR) checkpoints, aiding stability.
• Applicability: Can be integrated into multi-stage restoration pipelines and retrofitted onto outputs of non-diffusion fusion models.

7. Contextual Significance and Implications

IFControlNet advances conditional generative modeling and image restoration by:

• Enforcing spatial consistency and control fidelity across all diffusion steps, not solely at final outputs.
• Providing lightweight, stepwise feedback via convolutional probes, enabling fine-grained alignment even at high noise levels.
• Delivering state-of-the-art performance in multi-focus image fusion, with additive gains in image and perceptual quality over both deterministic algorithms and previous ControlNet variants.

A plausible implication is that this feature feedback approach generalizes to broader conditional generative tasks—potentially benefiting workflows requiring strict geometric or semantic structure in generated images. The methodology demonstrates that intermediate-feature alignment during the entire denoising trajectory materially affects generative outcomes, providing a new axis for conditional control in diffusion architectures (Konovalova et al., 3 Jul 2025, Xie et al., 25 Dec 2025).