Rapid Diffusion: Building Domain-Specific Text-to-Image Synthesizers with Fast Inference Speed
Year: 2023

Abstract:
Text-to-Image Synthesis (TIS) aims to generate images based on textual inputs. Recently, several large pre-trained diffusion models have been released to create high-quality images with pre-trained text encoders and diffusionbased image synthesizers. However, popular diffusion-based models from the open-source community cannot support industrial domainspecific applications due to the lack of entity knowledge and low inference speed. In this paper, we propose Rapid Diffusion, a novel framework for training and deploying superresolution, text-to-image latent diffusion models with rich entity knowledge injected and optimized networks. Furthermore, we employ BladeDISC, an end-to-end Artificial Intelligence (AI) compiler, and FlashAttention techniques to optimize computational graphs of the generated models for online deployment. Experiments verify the effectiveness of our approach in terms of image quality and inference speed. In addition, we present industrial use cases and integrate Rapid Diffusion to an AI platform to show its practical values. 1

heading: 1 Introduction
text: Text-to-Image Synthesis (TIS) is a prevalent multimodal task that aims to generate realistic images based on textual inputs, which supports real-world applications such as product appearance design and art creation. Apart from Generative Adversarial Network (GAN)-based approaches (Agnese et al., 2020), recently, pre-trained diffusion models (Rombach et al., 2022; Ramesh et al., 2022) have been
∗B. Liu and W. Lin contributed equally to this work. †C. Wang and K. Jia are co-corresponding authors.
1The source code is publicly available in the EasyNLP framework (Wang et al., 2022). URL: https://github.com/ alibaba/EasyNLP.
proposed to create artistic images with qualities comparable to or better than those from humans.
Despite the exciting advancement, for industrial domain-specific applications, we suggest that popular latent diffusion models from the open-source community (such as the Stable Diffusion model series2) are incapable of supporting those applications. The reasons are are twofolds. i) For diffusion-based methods, a CLIP-based text encoder (or other similar models) is required to encode the input texts, providing conditional inputs for the U-Net model (Rombach et al., 2022). As entities (or objects) are usually the key elements for generated images, CLIP models pre-trained over text-image pairs collected from the Web may need more abilities of concept understanding and are challenging to capture the specific entity knowledge required for realistic image generation (Ma et al., 2022). ii) For industrial applications, the model inference speed and the computational cost are vital factors to be considered. The cumbersome computation of the iterative diffusion process is often the bottleneck of fast inference (Song et al., 2021). Therefore, obtaining knowledgeable diffusion models to generate high-resolution images with moderate parameter sizes and optimized implementations that support fast online inference is desirable.
To address the above issues, we propose Rapid Diffusion, a novel framework for the training and deploying text-to-image diffusion models with rich entity knowledge injected and networks optimized. In Rapid Diffusion, a knowledge-enhanced CLIP model is effectively trained for learning entity knowledge from knowledge graphs (KGs). To
2https://stability.ai/blog/ stable-diffusion-public-release
295
generate high-resolution images and avoid parameter explosion, we integrate an ESRGAN-based network (Wang et al., 2018) after the diffusion block for image super-resolution, instead of directly leveraging a large-scale hierarchical diffusion model. For online deployment, an efficient inference pipeline is designed with the neural architectures optimized based on FlashAttention (Dao et al., 2022). The Intermediate Representation (IR) of computational graphs built from the generated models are further processed by a recently released Artificial Intelligence (AI) compiler (Zhu et al., 2021).
In the experiments, we evaluate the effectiveness of Rapid Diffusion in terms of the qualities of generated images from multiple application domains and the model inference speed for online deployment. We also provide industrial use cases to show how our framework benefits real-world applications. In addition, we have integrated the proposed training and deployment workflows into an industrial, cloud-native AI platform to facilitate zero-code model training and elastic inference on distributed GPU clusters. In summary, the major contributions of this work are as follows:
• We propose the Rapid Diffusion framework for the training and deployment of domainspecific diffusion-based TIS models. Specifically, a new knowledge-enhanced model pipeline is designed for super-resolution TIS. An efficient inference pipeline is further designed to optimize the computational graphs of our model for faster model inference. • Experiments over multiple domains show the effectiveness of Rapid Diffusion in terms of both image quality and inference speed, achieving an average FID score of 21.90 and ×1.73 acceleration ratio compared to all the counterparties. • We demonstrate the industrial use case and the integration of Rapid Diffusion to a cloudnative AI platform to show its practical values for real-world applications.

heading: 2 Related Work

heading: 2.1 Text-to-Image Synthesis (TIS)
text: TIS is a multi-modal task of converting texts to images with the same semantic meanings. In the early years, traditional methods (Zhu et al., 2007) mainly focused on analyzing the correlations be-
tween sentences and images but could not generate new images on the pixel level. Generative Adversarial Network (GAN) (Goodfellow et al., 2014) was proposed in 2014 and became the mainstream approach in the image synthesis field (Agnese et al., 2020). GANs and their variants (Reed et al., 2016; Liu et al., 2022b) have proved their effectiveness in TIS but still lack the ability to generate high-resolution images. Diffusion models (Ho et al., 2020; Sohl-Dickstein et al., 2015) have attracted the attention of researchers in recent years. Leveraging large-scale text-image datasets, pre-trained diffusion models (Rombach et al., 2022; Ramesh et al., 2022) become competitive with human painters. However, these diffusion models need help with the efficiency problem and more knowledge for the generation process.

heading: 2.2 Efficient Methods for Diffusion Models
text: Diffusion models (Ho et al., 2020; Sohl-Dickstein et al., 2015) typically add noise to images (or latent tensors generated from images) and then learn to denoise step by step. The number of steps while training may be very large, making the sampling time-consuming. To improve the sampling efficiency, a recent study (Salimans and Ho, 2022) introduces knowledge distillation to diffusion models. This acceleration method can reduce the steps but requires additional training. Another family of methods tries to construct new samplers without further training. For example, DDIM (Song et al., 2021) uses a deterministic generative process to produce images much faster. Some numerical solvers, including forward Euler and linear multistep method (Butcher, 2000), are leveraged to reduce the steps (Karras et al., 2022). By using a pseudo numerical algorithm to solve differential equations on manifolds, PNDM (Liu et al., 2022a) further improves the generation quality within a few given steps. The better implementation of attention algorithms can speedup the process, which requires fewer IO accesses (Dao et al., 2022). Colossal-AI (Bian et al., 2021) accelerates the training speed of diffusion models and reduces the GPU memory usage for deployment. In addition, when diffusion models are deployed online, the amount of computation can be reduced with better-complied computational graphs of these models. For example, TensorRT3 provides an inference optimizer and runtime that achieves lower
3https://github.com/NVIDIA/TensorRT
latency of model inference. In our work, we integrate various techniques from both modeling and engineering aspects to deliver better training and inference experiences for diffusion-based, domainspecific TIS applications.

heading: 3 The Proposed Framework
text: In this section, we formally present the techniques of the proposed Rapid Diffusion framework in detail. A brief overview of Rapid Diffusion is presented in Figure 1.

heading: 3.1 Model Architecture
text: As seen in Figure 1, our model converts input texts to high-resolution images by modeling the transformation and interaction between three representation spaces: i) knowledge-enhanced text embedding space, ii) latent space, and iii) pixel space.

heading: 3.1.1 Knowledge-enhanced Text Embedding Space
text: In this stage, we aim to encode the semantics of input texts to text embeddings. A common practice is to leverage the text encoder of CLIP (Radford et al., 2021), which jointly learns textual and visual representations in a unified space. Yet, CLIP pre-trained over plain text-image pairs may have weak representation power of entities. In our work, we leverage the 100 million text-image pairs from Wukong (Gu et al., 2022) as our multimodal pre-training corpus, as our real-world applications mostly focus on the Chinese language. For entities, we leverage the largest Chinese KG available to us, i.e., OpenKG4 (containing over 16 million entities and 140 million relation triples). During the CLIP pre-training process, the input representation of an entity token e appearing in a sentence of the Wukong corpus is augmented by: e⃗ = e⃗txt + e⃗kg where e⃗txt is the vanilla token embedding of the entity e, and e⃗kg is the KG embedding derived by the TransE algorithm (Bordes et al., 2013) due to its effectiveness and simplicity. Note that although we focus on the pre-training of Chinese Knowledge-enhanced CLIP (CKCLIP) models here, our method is language-invariant and can be applied to other languages with minor modifications.
During the fine-tuning process of our domainspecific TIS models, the parameters of the text
4http://openkg.cn/
encoder of our CKCLIP model are set to be trainable to capture more domain-related semantics. We further add some text prompts to the input text according to the application scenario (e.g., “the photo of [object]”, “the picture of [object]”) and obtain its CLIP representation as the conditional input to the next stage.

heading: 3.1.2 Latent Space
text: According to the given knowledge-enhanced text embedding e⃗, we use a latent diffusion model to generate image encoding with similar semantic meaning in a latent space. The model architecture is U-Net (Ronneberger et al., 2015) with a cross-attention mechanism capturing the textual conditioning information. In the training stage, the image x is encoded into the latent space and then we add Gaussian noise ϵ ∼ N (0, 1) to obtain xt, where t = 1, . . . , T is the step in the diffusion process. The loss function of image reconstruction is formulated as follows: LLDM = Ex,t∼U(1,...,T ),ϵ∼N (0,1) ( ||ϵ− ϵθ(xt, t)||22 )
(1) The generation process is the reverse of the diffusion process. Starting from the random Gaussian noise xT , the latent diffusion model gradually denoises the latent tensor and successively calculates xT−1, xT−2, . . . , x1. To improve the correlation between the generated image and the input prompt, we use the classifier-free guidance (Ho and Salimans, 2021) to generate the corresponding images. Additionally, to avoid the efficiency problem caused by too large step T , we employ PNDM (Liu et al., 2022a) scheduler to reduce the steps. In our work, we also pre-train the latent diffusion model using the Wukong dataset (Gu et al., 2022) and fine-tune the model using every downstream dataset respectively.

heading: 3.1.3 Pixel Space
text: After generating the final latent x0, a KLregularized decoder D reconstructs the image from x0 in the pixel space. In our design, the generated images are not necessarily in high resolution. Instead, an ESRGAN-based network (Wang et al., 2018) is applied after the decoder such that after a single forward pass, a corresponding highresolution image can be generated. An alternative design choice is directly generating high-resolution images using the diffusion model. However, this setting can be sub-optimal to satisfy the requirements of moderate model size and fast inference
speed. Consider a base 256× 256 diffusion-based U-Net model that employed for the TIS task, where one with a 256× 256 −→ 1024× 1024 diffusionbased super-resolution U-Net model and the other with an ESRGAN-based model. The former contains more parameters and requires more steps for inference, resulting in its inference time being several times slower than the latter. Therefore, we adopt an ESRGAN-based model to generate highresolution images efficiently.

heading: 3.2 Inference Speedup Designs
text: The inference process of the proposed model in this paper consists of three main components. We profile the inference speed of the original PyTorch model in eager mode and observe that the bottleneck is primarily located in the loop of the U-Net model, where the cross-attention computation dominates the inference time. The profiling result can be seen in Figure 2. To resolve this issue, we incorporate automatic slicing and compilation optimization techniques to optimize the entire pipeline in an end-to-end manner and introduce an IO-aware attention implementation to enhance the inference performance further.

heading: 3.2.1 Compilation Optimization
text: Our algorithm generates various low-level runtime flows for models with dynamic shapes on specific devices. It is achieved by enhancing a set of IR to create a complete dynamic shape representation (Zhu et al., 2021). For the operations
with intensive memory access, we fully utilized shared memory to design larger-grained kernel fusion strategies, effectively reducing the CPU/GPU switches (Zheng et al., 2022). Optimal graph partitioning and kernel implementation selection are performed for optimal inference speed. The optimization has been applied throughout the computing module, resulting in a significant improvement in inference speed.

heading: 3.2.2 Effective IO-Aware Attention
text: Based on the automatic compilation optimization, we further utilize the FlashAttention technique (Dao et al., 2022) for the cross-attention operator of U-Net, which is the core of the network’s inference bottleneck. The technique is based on the attention IO characteristics and performs tiling operations on the attention calculation to reduce memory read-write computation. We introduce different FlashAttention kernel implementations for
various combinations of computing devices and hardware architectures and dynamic inputs. The technique mentioned in the previous section effectively assists us in automatically finding the optimal implementation. As a result, the cross-attention calculation can be accelerated without deviation, yielding a 1.9× speed-up for the U-Net module.

heading: 4 Experiments

heading: 4.1 Experimental Settings
text: We first pre-train the CKCLIP model in the following experiments using the text-image pairs from Wukong (Gu et al., 2022) and the OpenKG. After that, the text encoder of CKCLIP and our diffusion model are pre-trained using the same data source. We fine-tune and evaluate the model over three domain-specific datasets to show the values of Rapid Diffusion in real-world applications. Implementation details and parameter settings can be found in the appendix.

heading: 4.2 Results of Three Application Scenarios
text: We report the performance of Rapid Diffusion over three domain-specific scenarios (i.e., E-commerce5, Chinese Painting (Li et al., 2021) and Cuisine, which are closely related to our applications) in terms of Frechet Inception Distance (FID) (Heusel et al., 2017) score. Details of the three datasets, together with the training/validation/testing splits, are given in Table 4 in the appendix. We compare our model with three popular open-source diffusion models, namely Stable diffusion6, Stable diffusion 27, and Taiyi Diffusion8 (which is the largest Chinese diffusion model available so far). Note that Stable diffusion and Stable diffusion 2 mainly support English text inputs. Hence, we leverage the Chinese-English translation model (Wei et al., 2022) to translate our texts to English. The results are shown in Table 1. It can be seen that Rapid Diffusion outperforms all counterparties over the three datasets, achieving the average FID score at 21.90. The results indicate that our knowledge-enhanced models over domain-specific scenarios understand domain knowledge better and can generate more realistic and varied images.
5https://tianchi.aliyun.com/muge 6https://huggingface.co/CompVis/
stable-diffusion-v-1-4-original 7https://huggingface.co/stabilityai/ stable-diffusion-2 8https://huggingface.co/IDEA-CCNL/ Taiyi-Stable-Diffusion-1B-Chinese-v0.1

heading: 4.3 Effectiveness of Knowledge-enhanced Chinese CLIP
text: As CLIP models aim to learn cross-modal representations, we first intrinsically evaluate our model by text-image retrieval. We compare the vanilla Chinese CLIP model and our CKCLIP model using the same pre-training text-image corpus. Pretraining details can also be found in the appendix. For evaluation, we employ the standard split of Flickr30K-CN (Lan et al., 2017), and then finetune both models. Table 2 reports the text-to-image and text-to-image retrieval results over the testing set. Our CKCLIP model improves retrieval performance by significant margin (especially for R@1 metric), showing its ability to learn cross-modal representations. In addition, we provide some qualitative results from the Cuisine dataset to show how more entity knowledge can lead to better representation and generation of the key objects in images, as shown in Figure 3.

heading: 4.4 Results of Inference Speedup
text: For the implementation of compilation optimization, we employ BladeDISC (Zhu et al., 2021) as
our underlying AI compiler, which is an end-toend dynamic shape compiler for machine learning workloads. Results of the comparison between our implementation and Torch Native (eager mode) are displayed in Table 3. According to the results, U-Net inference takes the longest in the entire process (in 1900ms), while the text encoder and the image decoder take only 0.012ms and 0.04ms, respectively. However, with the optimization by BladeDISC, we speed up the inference time of the text encoder, U-Net and the image decoder by ×3, ×1.91, and ×1.9 times compared to the eager mode. Additionally, FlashAttention assists us in further optimizing U-Net, which decreases the inference time from 994ms to 759ms. Finally, we are able to generate images more quickly. Note that the underlying GPU used in the experiments is NVIDIA A100 (80GB), and the scheduler runs 50 in steps.

heading: 4.5 Results of Super-resolution
text: For image super-resolution, the ESRGAN-based network can be efficiently leveraged to achieve upscaling results. We can directly use the ESRGANbased network following the latent diffusion model because it has been pre-trained on common-used image datasets such as DIV2K (Agustsson and Timofte, 2017). However, considering the uniqueness and consistency of domain-specific images, we conjecture that fine-tuning enables the model to perform better. Experiments show that after finetuning, the model beats the pre-trained model according to our qualitative and quantitative results, which achieves 23.1 in terms of Peak Signal to Noise Ratio9, while the pre-trained model achieves only 22.7. Figure 4 further compares images with and without our pre-trained/fine-tuned models.

heading: 4.6 Case Studies
text: We provide more cases from each application domain to show how much our model outperforms previous ones. Refer to Figure 5 in the appendix.
9https://en.wikipedia.org/wiki/Peak_ signal-to-noise_ratio

heading: 5 Applications
text: In this section, we demonstrate the practical values of Rapid Diffusion by industrial use case and the integration to Alibaba Cloud PAI (Machine Learning Platform for AI).

heading: 5.1 Industrial Use Case
text: Here, we briefly discuss two real-world use cases. The first is a fashion design for e-commerce manufacturers. The inputs to our system consist of keywords for multiple elements, such as trend, fabric, color and style. An automatic prompt generation process is called to provide TIS models naturallanguage-like inputs. For a single request of fashion design, a handful of prompts can be generated, each associated with multiple generated images. The images are then regarded as materials for designers. The cuisine dataset described previously is from our online food delivery and local life service platform. Our diffusion model for cuisine generation provides the inspiration functionality to help service providers to create innovative menus where users can select or freely enter all kinds of foodrelated keywords to generate images. Note that the images will be marked as “AI-generated” before they are sent to our applications.

heading: 5.2 Integration to AI Platform
text: To allow users to create their models, we have integrated Rapid Diffusion into a cloud-native AI platform to facilitate zero-code model training and elastic inference. For model training, after uploading training/validation datasets and checking hyper-
parameters, a training job is automatically submitted to our deep learning container, where the training command and the docker image have already been prepared. After the job is completed, the resulting model is available for deployment. Based on Query Per Second (QPS) requirements, our prediction service can scale to an adjustable number of machines in the cloud. We can call the TIS service via a RESTful API by HTTP requests.

heading: 6 Conclusion and Future Work
text: We present the Rapid Diffusion framework for the training and deploying knowledge-enhanced, domain-specific, high-resolution, diffusion-based TIS models. Experimental results show the effectiveness of Rapid Diffusion in both image quality and inference speed, achieving an average FID score of 21.90 and ×1.73 acceleration ratio compared to all the counterparties. We further show its practical values through industrial use cases and the integration into an AI platform. In the future, we will extend the functionality of Rapid Diffusion and further increase the inference speed by advanced compilation optimization techniques.
Ethical Considerations
The techniques for inference speedup presented in this work are fully methodological. Hence, there are no direct negative social impacts. However, as the models automatically generate the images, they may have some negative impacts, such as the generation of toxic content and the existence of social biases. We suggest that the produced models should not be used to generate offensive or inappropriate images for people intentionally. Users should carefully deal with the potential risks by filtering out these images when the models are deployed online.

heading: Acknowledgements
text: This work was partially supported by Alibaba Innovative Research Foundation (No. D8200510), NSFC (Grant No. 61936003), the Program for Guangdong Introducing Innovative and Enterpreneurial Teams (No. 2017ZT07X183), Zhuhai Industry Core and Key Technology Research Project (No. 2220004002350), and by Alibaba Cloud Group through Research Talent Program with South China University of Technology.

heading: A Generated Results of Case Study
text: We show more generated images from our model and baselines, presented in Figure 5.

heading: B Data Statistics
text: Table 4 shows the data statistics of our experiments. We divide the E-commerce and Chinese Painting datasets into training, validation and testing sets according to the ratio of 80%, 10%, and 10%. For the Cuisine dataset, we divide 10% for validation, 10 thousand images for testing and the rest for training. Among these datasets, E-commerce and Chinese Painting are public datasets, while Cuisine is an inhouse dataset provided by our online food delivery and local life service platform.

heading: C Hyper-parameters Settings
text: For knowledge-enhanced CLIP pre-training, we follow the hyper-parameter settings in (Gu et al., 2022). For training the latent diffusion model, we set the learning rate as 5× 10−5, the batch size as 80, and the image size as 256 × 256. The latent dimension of the auto-encoder is 32 × 32. The hidden dimension of the text-encoder is 768.
For fine-tuning the super-resolution model, we obtain low-resolution images by down-sampling high-resolution images using the bi-cubic kernel function. Different from the original two-stage training process, we directly employ the pre-trained ESRGAN model as an initialization for the generator and the discriminator. We use Adam with β1 = 0.9, β2 = 0.999. The batch size is set to 16 and the learning rate is set to 1× 10−4 and halved at [50k, 100k, 200k, 300k] iterations.
During model training, all the experiments are conducted on a single server with 8 NVIDIA A100 GPUs (80G).
C.1 Hyper-parameters of Model Architectures
Table 5 shows the model sizes of all the experiment models, including Stable Diffusion, Stable Diffusion 2, Taiyi Diffusion and our Rapid Diffusion
model. Compared with the other three baselines, Rapid diffusion is the most compact model with better performance in our scenarios.
We further provide the detailed settings of the entire Rapid Diffusion model pipelines in Table 6.

heading: Future Work
text: In the future, we will extend the functionality of Rapid Diffusion and further increase the inference speed by advanced compilation optimization techniques. Ethical Considerations
The techniques for inference speedup presented in this work are fully methodological. Hence, there are no direct negative social impacts. However, as the models automatically generate the images, they may have some negative impacts, such as the generation of toxic content and the existence of social biases. We suggest that the produced models should not be used to generate offensive or inappropriate images for people intentionally. Users should carefully deal with the potential risks by filtering out these images when the models are deployed online