CountDiffusion: Text-to-Image Synthesis with Training-Free Counting-Guidance Diffusion

TL;DR

This work tackles the challenge of generating images with accurate object quantities in text-to-image diffusion models. It introduces CountDiffusion, a training-free framework that first evaluates an intermediate denoising result to count objects and then applies a correction stage that modifies attention maps via universal guidance to adjust counts, enabling plug-and-play compatibility with diffusion models. Key contributions include a two-stage counting framework, a multi-class correction strategy, and a dataset built with LLMs for evaluation, with results showing substantial gains in accuracy, MAE, CLIP similarity, and human preferences. The approach offers a practical path to more quantitatively faithful T2I generation, though it notes limitations with large object sets and points to future improvements in scalability and counting robustness.

Abstract

Stable Diffusion has advanced text-to-image synthesis, but training models to generate images with accurate object quantity is still difficult due to the high computational cost and the challenge of teaching models the abstract concept of quantity. In this paper, we propose CountDiffusion, a training-free framework aiming at generating images with correct object quantity from textual descriptions. CountDiffusion consists of two stages. In the first stage, an intermediate denoising result is generated by the diffusion model to predict the final synthesized image with one-step denoising, and a counting model is used to count the number of objects in this image. In the second stage, a correction module is used to correct the object quantity by changing the attention map of the object with universal guidance. The proposed CountDiffusion can be plugged into any diffusion-based text-to-image (T2I) generation models without further training. Experiment results demonstrate the superiority of our proposed CountDiffusion, which improves the accurate object quantity generation ability of T2I models by a large margin.

Figures (9)

• Figure 1: Samples generated by SDXL, Pixart-$\Sigma$, Ranni, InstanceDiffusion, and CountDiffusion (Ours). Existing models struggle with generating images with correct objects counts. Even when provided with bounding boxes, the models may still generate target objects in the background. Furthermore, Ranni, which combines a LLM for T2I generation, is limited by the language model and is still highly likely to generate incorrect objects counts. CountDiffusion can correct the object quantity based on the generation results of the baseline model, $i.e.$, Pixart-$\Sigma$.
• Figure 2: Pipeline of the proposed CountDiffusion. CountDiffusion consists of a detection stage and a correction stage. In the detection stage, an intermediate denoising result generated by the diffusion-based T2I model is used to predict the final synthesized image with a one-step denoising. Then, a counting model, e.g., Grounded SAM, is used to get the quantity and segmentation information of the objects in this image. In the correction stage, a correction module is used to rectify the object quantity by modifying the latent feature of the synthesized image with universal guidance. Different correction strategy for single-class objects and multi-class objects.
• Figure 3: The synthesized images from intermediate denoising results of different denoising steps with one-step denoising. Total denoising step $T=40$ in this case. After approximately 15 denoising steps (i.e., $t_{mid}=25$), the image generated with one-step denoising is already able to provide sufficient object quantity and position information.
• Figure 4: Comparison between the proposed CountDiffusion and state of the arts across different object quantities on GPTSingleCount dataset.
• Figure 5: Qualitative comparisons of all models. Our method successfully generates the correct number of objects, while other methods struggle in some or all of the examples. The red smiling face in the bottom right corner of the image indicates that the correct number of objects were generated, while the green crying face indicates an error in generation.