DiffAgent: Fast and Accurate Text-to-Image API Selection with Large Language Model

Lirui Zhao,Yue Yang,Kaipeng Zhang,Wenqi Shao,Yuxin Zhang,Yu Qiao,Ping Luo,Rongrong Ji

2024-03-31

Abstract:Text-to-image (T2I) generative models have attracted significant attention and found extensive applications within and beyond academic research. For example, the Civitai community, a platform for T2I innovation, currently hosts an impressive array of 74,492 distinct models. However, this diversity presents a formidable challenge in selecting the most appropriate model and parameters, a process that typically requires numerous trials. Drawing inspiration from the tool usage research of large language models (LLMs), we introduce DiffAgent, an LLM agent designed to screen the accurate selection in seconds via API calls. DiffAgent leverages a novel two-stage training framework, SFTA, enabling it to accurately align T2I API responses with user input in accordance with human preferences. To train and evaluate DiffAgent's capabilities, we present DABench, a comprehensive dataset encompassing an extensive range of T2I APIs from the community. Our evaluations reveal that DiffAgent not only excels in identifying the appropriate T2I API but also underscores the effectiveness of the SFTA training framework. Codes are available at

Computation and Language,Artificial Intelligence

What problem does this paper attempt to address?

The paper mainly addresses the following issues:

1. **Quickly screening suitable Text-to-Image (T2I) generation models and parameters**: With a large number of personalized T2I models (for example, there are approximately 74,492 different models in the Civitai community), it is difficult for users to quickly find the model and parameter combination that best suits their specific needs. It usually requires multiple attempts to achieve satisfactory results.

2. **Improving the quality and style matching of generated images**: Existing Stable Diffusion (SD) models have limitations when dealing with certain specific styles (such as pixel art), and the generated images may not meet user expectations.

To address these issues, the paper proposes **DiffAgent**, an agent model based on large language models (LLM) designed to quickly select the most suitable T2I generation API based on user input prompts. DiffAgent achieves this goal through a two-stage training framework:

- **First Stage**: Supervised Fine-Tuning (SFT) to build the initial model DiffAgent-SFT. It is trained using a large-scale instruction-following dataset DABench collected from the Civitai website, enabling the model to generate corresponding T2I APIs based on user prompts.

- **Second Stage**: Further optimization using the Ranking Response with Human Feedback (RRHF) algorithm to reduce hallucination errors in the generated APIs and better match human preferences, resulting in the final model DiffAgent-RRHF.

Additionally, the paper introduces a carefully curated high-quality dataset **DABench**, specifically designed for the T2I API selection task, containing a large number of callable API instances. Experimental results show that DiffAgent not only significantly outperforms baseline models on multiple benchmarks but also effectively reduces hallucination errors and improves the quality and style matching of generated images.