UniMC: Taming Diffusion Transformer for Unified Keypoint-Guided Multi-Class Image Generation

Dear Reviewer u8PW,

Thank you for your review and constructive comments. During the rebuttal period, we have made every effort to address your concerns. The detailed responses are below:

Q1: Why human shows mild improvement compared with ControlNet?

While our method shows only mild improvement on the human class compared to ControlNet, it brings substantial gains across other categories, especially animals, as shown in Table 3 of the main paper. This highlights the strength of our approach in handling diverse object classes.

Our core objective is to solve the task of keypoint-conditioned multi-class image generation, where human is just one of many categories. In contrast, ControlNet is specifically designed for human-centric keypoint-conditioned generation, and thus naturally excels in that narrow domain.

This difference is analogous to a specialist vs. generalist comparison. A specialist may outperform others on a single task, but a well-designed generalist can achieve strong performance across a broader spectrum. UNIMC is purposefully built as a generalist model that scales across multiple categories, not just humans.

Q2: Explanation of dropout ratio choices.

In our initial experiments, we applied a 15% dropout ratio to both bounding boxes and keypoints. However, we observed that the model quickly overfit to the bounding box information—likely because coarse object localization is much easier to learn than fine-grained keypoint patterns. As a result, the model tended to ignore the keypoint inputs, and the pose accuracy plateaued below 20%, regardless of the number of training iterations.

To address this issue, we increased the dropout ratio for bounding boxes, forcing the model to rely more on the keypoints. The table below shows how different dropout ratios for bounding boxes affect pose accuracy:

These results suggest that a moderate dropout (around 50%) strikes a good balance, encouraging the model to make better use of keypoint information without completely discarding spatial cues from bounding boxes.

Dear Reviewer Rf97,

Thank you for your review and constructive comments. During the rebuttal period, we have made every effort to address your concerns. The detailed responses are below:

Q1: The need for a joint human-animal keypoint-annotated dataset.

As shown in Table 4 and Figure 6 of the main paper, training on separate human and animal keypoint datasets makes it difficult to generate images with multiple classes, limiting the development of more general keypoint-conditioned generation.

Q2: Results on multi-instance subset.

We present results on the multi-instance subset below. Our method still performs well, while baseline models decline, demonstrating that UNIMC performs better in multi-class scenarios.

Q3: Results on COCO+APT testing datasets.

We have added results on the COCO+APT testing set, and similar to HAIG, our method demonstrates strong performance and surpasses all baseline methods.

Q4: Performance with and without multiple instances in the dataset.

The training performance of HAIG with and without multiple instances shows that adding multi-instance images improves all metrics.

Q5: Comparison with [1]MIGC and [2]InstanceDiffusion.

The table below presents the quantitative comparison with MIGC and InstanceDiffusion. MIGC controls via bounding boxes and classes, while InstanceDiffusion can control through bounding boxes, points, and classes. Compared to previous baseline methods, both of these models show significant improvements in class control accuracy. InstanceDiffusion, which adds point control, also improves keypoint control accuracy.

The quantitative comparison can be seen in Fig.5.

Q6: What is the number of images used for FID comparison?

As mentioned in the Implementation Details section, we used the HAIG-2.9M testing set for evaluation, with 1,224 images.

Q7: Why not annotate with the ensemble of models?

Relevant works, such as Pixart-sigma, have shown that using a better annotation model improves the model's performance. Additionally, animal pose estimation is not as mature as human pose estimation, and only selecting the best model ensures that we can achieve usable results.

Q8: Bias in Evaluation.

Thank you for your insightful question. Please refer to Q3 — our method also achieves strong performance on human-annotated datasets, which supports the validity of our evaluation results and the reliability of our evaluation methodology.

Q9: Difference from coordinate regression tasks.

Our task is to extend keypoint-conditioned image generation from a single category (human) to a multi-class setting. To achieve this, we use keypoint coordinates as conditional inputs to guide image generation. In contrast, prior works on coordinate regression use coordinates as outputs, not as conditioning inputs. While these works demonstrate that DiT can model coordinate distributions, they do not show that DiT can generate images conditioned on coordinates, which is the core challenge of our work. Therefore, the prior use of DiT for coordinate regression is only loosely related to our task and does not diminish the novelty of our proposed approach.

Dear Reviewer 2vGV,

Thank you for your review and constructive comments. During the rebuttal period, we have made every effort to address your concerns. The detailed responses are provided below:

Q1: Applicability to other models.

Our method is designed to be compatible with almost all DiT-based architectures, due to the following reasons:

1. The Unified Keypoint Encoder is model-agnostic and does not rely on any specific architectural design.
2. The Timestep-aware Keypoint Modulator is implemented using only self-attention and does not depend on any model-specific components.

To demonstrate the general applicability of our method, we conducted experiments on Hunyuan-DiT. The specific experimental results are shown in the table below. Our method successfully equips Hunyuan-DiT with the ability to perform keypoint-conditioned multi-class image generation.

Q2: Analysis of training and inference speed.

Please refer to R1-Q1 and R1-Q4. Our method is highly efficient in both training and inference.

Specifically, it requires only 576 A800 GPU-hours to complete training. At inference time, it is more efficient than GLIGEN and comparable to ControlNet. Moreover, our method can quickly achieve strong performance even when fine-tuned on small-scale datasets.

Q3: Comparison with various text-based and Visual-Language Model (VLM)-based image generation models.

Our method is orthogonal to general text-to-image (T2I) or Visual-Language Model (VLM)-based image generation approaches. These models lack the ability to generate images conditioned on keypoints from different object categories. In contrast, we introduce a framework that enables keypoint-conditioned multi-class image generation.

The table below provides a comparison between UNIMC and recent state-of-the-art T2I models. Notably, our method achieves substantial improvements in Human/Animal class accuracy and keypoint-level AP, demonstrating its superiority in controllability and structural alignment.

Fig.4 presents a qualitative comparison of UNIMC with Flux and Kolors. As shown, the latter two models, relying solely on the prompt, fail to generate images that align with the provided keypoints.

Q4: Why our method outperforms prior work in keypoint-conditioned multi-class image generation?

The core objective of our method is to enable keypoint-conditioned multi-class image generation.
To this end, we introduce a Unified Keypoint Encoder that explicitly models the relationship between object categories and their corresponding keypoints, and a Timestep-aware Keypoint Modulator that enables fine-grained, keypoint-level feature modulation.

In contrast to prior methods, our approach effectively addresses the challenges of class binding fusion and keypoint binding fusion, allowing the model to learn precise keypoint control across multiple classes simultaneously.
As a result, UNIMC achieves significant improvements over baseline methods on this task.

Q5: Additional contributions beyond HAIG-2.9M and UNIMC.

Our primary contribution is extending keypoint-conditioned image generation from a single class (human) to a general multi-class setting. To achieve this goal, we propose both the HAIG-2.9M dataset and the UNIMC model, which together enable keypoint-conditioned multi-class image generation across diverse object categories.

Dear Reviewer kiWJ,

Thank you for your review and constructive comments. During the rebuttal period, we have made every effort to address your concerns. The detailed responses are provided below:

Q1: Details about the computational requirements of UNIMC compared to baseline methods.

Each method is trained using different datasets and GPU configurations, and their reported training costs are measured in different formats. For example, ControlNet reports GPU-hours, while GLIGEN uses steps * batch size. Therefore, directly comparing training costs across methods is not meaningful. In the table below, we report the officially reported training costs of baseline methods, along with our method’s training cost in both formats. Despite achieving significantly better performance than the baselines, our method maintains a relatively low training cost. For inference, all methods are evaluated on a single RTX 3080 Ti with FP16 precision, using a 50-step DDIM sampler and Flash-Attention-v2. The generation resolution is 512×512.

Our method is more efficient than ControlNet, and although the inference time and memory usage are slightly higher than those of GLIGEN, the performance improvements over GLIGEN are substantial.

Q2: How well does the model generalize to unseen animal species not present in the HAIG-2.9M dataset?

Although our model is trained on only 31 classes, it generalizes well to unseen categories due to two key factors:

1. the structural similarity across different species, and
2. our use of a text encoder for class representation, which allows for extension to arbitrary categories.

In addition, the underlying pretrained T2I model is inherently capable of class-to-image generation for a wide range of categories. Together, these components enable our method to theoretically scale to any class.

Taking kangaroo as an example, we observe that its keypoint structure is similar to that of humans. Thus, we reused human keypoints to generate a kangaroo. As shown in Fig.1, the generated pose and spatial structure roughly align with the keypoints, suggesting that our model possesses a certain degree of generalization to unseen species.

Therefore, we believe that as long as the training set covers a sufficient diversity of structurally representative animal classes, our model can generalize to a broad range of animal species.

Q3: Failure cases or limitations of UNIMC.

We observed that for some categories, a few keypoints are not perfectly controllable. For example, in the case of cat shown in Fig.2, the tail fails to be accurately generated. Beyond such rare cases, we did not observe significant failure cases in the test set.

Q4: What would be the minimum computational setup required to fine-tune UNIMC on a smaller dataset for a specific application?

Our use of 8 A800 GPUs was intended to accelerate training on a large-scale dataset. However, the UNIMC model itself is lightweight, with only 0.93B parameters (and 0.35B trainable). Under FP16 precision and batch size = 1, it requires only 12GB of GPU memory.

We fine-tuned UNIMC on a dataset of 1,000 cat images using a single RTX 3090, with batch size = 4 and 2K training steps, taking about 4 hours. The resulting AP for the cat category reached 28.05, and qualitative results are shown in Fig.3.

Q5: How did you ensure the quality of keypoint annotations in HAIG-2.9M?

As described in Section 4.2, we first randomly sampled 5K images and annotated them using several expert models. We then conducted a user study to select the best-performing model. As a result, our keypoint annotations reach the upper bound achievable by current pretrained models.

Moreover, the animal species in our dataset are already included in the training data of the selected pretrained keypoint estimator. During annotation, we manually sample and discard images with low-quality annotations, and we also filter out images where keypoint predictions have low confidence.

To further reduce ambiguous keypoint placements, we feed the bounding box detection results into the keypoint estimator and restrict keypoint detection to the corresponding bounding boxes. This ensures that the estimator focuses only on the relevant regions, thereby improving annotation accuracy and reliability.