Make It Count: Text-to-Image Generation with an Accurate Number of Objects

TL;DR

The approach, CountGen, does not depend on external source to determine object layout, but rather uses the prior from the diffusion model itself, creating prompt-dependent and seed-dependent layouts, which strongly outperforms the count-accuracy of existing baselines.

Abstract

Despite the unprecedented success of text-to-image diffusion models, controlling the number of depicted objects using text is surprisingly hard. This is important for various applications from technical documents, to children's books to illustrating cooking recipes. Generating object-correct counts is fundamentally challenging because the generative model needs to keep a sense of separate identity for every instance of the object, even if several objects look identical or overlap, and then carry out a global computation implicitly during generation. It is still unknown if such representations exist. To address count-correct generation, we first identify features within the diffusion model that can carry the object identity information. We then use them to separate and count instances of objects during the denoising process and detect over-generation and under-generation. We fix the latter by training a model that predicts both the shape and location of a missing object, based on the layout of existing ones, and show how it can be used to guide denoising with correct object count. Our approach, CountGen, does not depend on external source to determine object layout, but rather uses the prior from the diffusion model itself, creating prompt-dependent and seed-dependent layouts. Evaluated on two benchmark datasets, we find that CountGen strongly outperforms the count-accuracy of existing baselines.

Figures (18)

• Figure 1: CountGen generates the correct number of objects specified in the input prompt while maintaining a natural layout that aligns with the prompt.
• Figure 2: Architecture outline. Given a prompt that includes a quantity, we begin generating a corresponding image using pretrained SDXL until timestep $t=500$. We then perform Instance Localization, where we combine cross-attention maps corresponding with the object, and self-attention features extracted at timestep $t$ to generate object clusters for each generated object. Then we apply ReLayout, which generates an object layout with the correct number of instances, while preserving the composition of the extracted layout. Finally, we perform Layout Guided generation, which applies an inference time optimization based on the layout through cross-attention loss $L_\text{cross}$ and self-attention masking.
• Figure 3: PCA Visualization. to explore the notion of objectness inside SDXL latent space, we visualize dimension-reduced self-attention feature maps from various layers across the network at timestep $t=500$. We notice that although most layers do not exhibit a clear separation between object instances, layer $l^{up}_{52}$ displays a robust separation indicated by different ones having distinct colors. Visualization across different timesteps is shown in the appendix (\ref{['fig:pca_timestamp']}).
• Figure 4: Correcting under-generation. we show examples for the ReLayout correction of cases where SDXL generates less objects than specified in the prompt. It is evident that the generated layouts are natural and obey the same composition of the original generation, with the correct number of objects.
• Figure 5: A training set for a ReLayout. We created pairs of images using SDXL, using the same seed and prompts that only differ by object count. We filtered out images that did not conform to the prompt, using the techniques described in Section \ref{['sec:layout_design']}. The resulting image pairs preserve the scene and layout except adding one object.