Leveraging Latent Vector Prediction for Localized Control in Image Generation via Diffusion Models

TL;DR

The paper tackles the challenge of fine-grained local control in diffusion-based image generation by introducing a latent-space conditioning framework that localizes guidance to user-defined regions. It combines masking-based conditioning applied to T2I-Adapter features, prediction of the initial latent vector $\hat{z}_0$ at each step, and a Sobel-based similarity loss to enforce alignment inside the region of interest while preserving creative generation elsewhere. Empirical results on a large LAION-based dataset show improved fidelity within the ROI (lower DMSE_in) and controlled background generation (moderate DMSE_out) with competitive FID and CLIPScore compared to global-condition baselines like T2I-Adapter and ControlNet. The approach is efficient, compatible with existing diffusion models, and offers a practical pathway to precise, localized image editing and synthesis, with future work exploring color conditioning, attention-based localization, and hierarchical constraints.

Abstract

Diffusion models emerged as a leading approach in text-to-image generation, producing high-quality images from textual descriptions. However, attempting to achieve detailed control to get a desired image solely through text remains a laborious trial-and-error endeavor. Recent methods have introduced image-level controls alongside with text prompts, using prior images to extract conditional information such as edges, segmentation and depth maps. While effective, these methods apply conditions uniformly across the entire image, limiting localized control. In this paper, we propose a novel methodology to enable precise local control over user-defined regions of an image, while leaving to the diffusion model the task of autonomously generating the remaining areas according to the original prompt. Our approach introduces a new training framework that incorporates masking features and an additional loss term, which leverages the prediction of the initial latent vector at any diffusion step to enhance the correspondence between the current step and the final sample in the latent space. Extensive experiments demonstrate that our method effectively synthesizes high-quality images with controlled local conditions.

Figures (7)

• Figure 1: Comparison of generated images and their edges from current image-level models given the prompt “A living room with a sofa.” The first column shows the original image and its edges. The subsequent columns present the generated images and their corresponding edge maps. Edge images are divided into (green areas), which represent edges within the region of interest where a spatial condition is applied, and we aim to preserve edge similarity; and (red areas), which denote unconditioned regions where the model is expected to generate a coherent background with natural contours, without over or under generation.
• Figure 2: Overview of the T2I-Adapter integration. The input condition $c_f$ is processed by the adapter network $F_{\text{AD}}(c_f)$ to generate condition features $F^c = \{F^c_1, F^c_2, F^c_3, F^c_4\}$, which are designed to match the spatial dimensions of the corresponding encoder feature maps $F^{\text{enc}} = \{F^{\text{enc}}_1, F^{\text{enc}}_2, F^{\text{enc}}_3, F^{\text{enc}}_4\}$ produced by the UNet. At each encoder layer, the adapter features are fused with the UNet features via element-wise addition, enabling the integration of image-level conditioning information into the diffusion process.
• Figure 3: Overview of our training methodology at each step $t$ of the denoising process. Given the textual prompt $p$, the control condition $c_{f}$ and region of interest defined by the mask $M$, an image I is passed through $\mathcal{E}$ to obtain its latent representation $z_{0}$. At each training step, we apply Local Control (see Section \ref{['sec:local_control']}) to the T2I-Adapter features. A timestep $t$ is then selected, and diffusion is performed on $z_{0}$ to generate the noisy input $z_{t}$. Latent variable $z_{t}$ jointly with $y$ and the T2IAdapter features are passed into the denoising network $U$, to predict $\epsilon_{\theta}$. We subsequently Predict the initial latent vector$\widehat{z}_{0}$ (see Section \ref{['sec:predicting_z0']}) using $\epsilon_{\theta}$ and $z_{t}$ to Enforce Similiarity (see Section \ref{['sec:enforcing_similarity']}) by applying a Sobel filter $\mathcal{S}$ on the original and predicted latents $z_{0}$, $\widehat{z}_{0}$ within the region of interest.
• Figure 4: Comparison between direct masking of the edge map and our feature-level masking approach. When the mask is applied directly to the edge map, the model cannot distinguish between unstructured regions (intended smooth areas) and unconditioned regions (where no guidance should be applied), resulting in a loss of control outside the masked area. In contrast, our approach applies the mask to the T2I-Adapter features, preserving three distinct zones: structured regions with edge information, unstructured regions intended to be smooth, and unconditioned regions outside the mask. This separation enables precise, localized control while allowing the model to generate coherent and diverse content beyond the ROI.
• Figure 5: Comparison of the effect of applying the mask either to the edge condition (Masking Edges) or to the adapter features (Masking Features) for both T2I-Adapter and our method. When the mask is applied to the edge maps, both models interpret areas without contours as smooth, resulting in blurred or empty regions outside the masked area. In contrast, when the mask is applied to the adapter features, T2I-Adapter continues to treat unconditioned areas as contourless, while our method correctly distinguishes between the ROI and the background. This allows our approach to preserve structure inside the mask and generate more coherent, realistic backgrounds beyond it.