SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Dongxu Li†, Junnan Li†, Steven C.H. Hoi†  
Salesforce AI Research
†Corresponding authors: {li.d, junnan.li, shoi}@salesforce.com https://github.com/salesforce/LAVIS/tree/main/projects/blip-diffusion
Figure 1: Leveraging the pre-trained subject representation, BLIP-Diffusion enables subject-driven generation under efficient fine-tuning or zero-shot setups. The model can also serve as a foundation subject-driven text-to-image generation model and supports applications such as controlled generation and image editing, when combined with techniques such as ControlNet [1] and prompt-to-prompt [2].
Subject-driven text-to-image generation models create novel renditions of an input subject based on text prompts. Existing models suffer from lengthy fine-tuning and difficulties preserving the subject fidelity. To overcome these limitations, we introduce BLIP-Diffusion, a new subject-driven image generation model that supports multimodal control which consumes inputs of subject images and text prompts. Unlike other subject-driven generation models, BLIP-Diffusion introduces a new multimodal encoder which is pre-trained to provide subject representation. We first pre-train the multimodal encoder following BLIP-2 to produce visual representation aligned with the text. Then we design a subject representation learning task which enables a diffusion model to leverage such visual representation and generates new subject renditions. Compared with previous methods such as DreamBooth, our model enables zero-shot subject-driven generation, and efficient fine-tuning for customized subject with up to 20x speedup. We also show that BLIP-Diffusion can be flexibly combined with existing techniques such as ControlNet and prompt-to-prompt to enable novel subject-driven generation and editing applications.
Text-to-image generation models have developed significantly and enabled creation of high-quality images based on textual prompts [3-7]. One of their applications is subject-driven generation, which aims to render novel renditions of an input subject while preserving its appearance. The common approach to subject-driven generation [8-11] is through inverting subject visuals into text embedding space. Specifically, with a pretrained text-to-image generation model, a placeholder text embedding is optimized to reconstruct a set of subject images. The embedding is then composed into natural language prompts to create different subject renditions. One known inefficiency of this approach is that it requires reiterating hundreds [9, 10] or thousands [8] tedious fine-tuning steps for each new subject, which hinders it from efficiently scaling to a wide range of subjects.
We attribute such inefficiency to the fact that most pre-trained text-to-image models do not natively support multimodal control - using both images and texts as control input. As a result, it becomes challenging to learn subject representation that aligns with the text space while capturing the subject visuals with high fidelity. To overcome these limitations, we introduce BLIP-Diffusion, the first subject-driven text-to-image generation model with pre-trained generic subject representation, which enables subject-driven generation in zero-shot or with few-step fine-tuning. Our model builds upon a vision-language encoder (i.e. BLIP-2 [12]) and a latent diffusion model [6] (i.e. Stable Diffusion). The BLIP-2 encoder takes as input the subject image and its category text; it produces text-aligned subject representation as output. We then infix the subject representation in the prompt embedding to guide the latent diffusion model for subject-driven image generation and editing.
To enable controllable and high-fidelity generation, we propose a new two-stage pre-training strategy to learn generic subject representation. In the first pre-training stage, we perform multimodal representation learning, which enforces BLIP-2 to produce text-aligned visual features based on the input image. In the second pre-training stage, we design a subject representation learning task where the diffusion model learns to generate novel subject renditions based on the input visual features. To achieve this, we curate pairs of input-target images with the same subject appearing in different contexts. Specifically, we synthesize input images by composing the subject with a random background. During pre-training, we feed the synthetic input image and the subject class label through BLIP-2 to obtain the multimodal embeddings as subject representation. The subject representation is then combined with a text prompt to guide the generation of the target image.
Benefiting from the pre-trained subject representation, BLIP-Diffusion achieves promising zero-shot subject-driven generation results and superior fine-tuning efficiency. For example, BLIP-Diffusion takes 40-120 fine-tuning steps to specialize for a given subject, achieving up to  $20\mathrm{x}$  speedup compared to DreamBooth [9]. Furthermore, BLIP-Diffusion inherits behaviours of the constituent latent diffusion model and can be flexibly extended to support various subject-driven generative applications without further training. Following the prompt-to-prompt [2] approach, BLIP-Diffusion enables editing images with subject-specific visuals. When combined with ControlNet [1], it enables subject-driven generation with various additional structure control.
2.1 Diffusion Models for Text-to-Image Generation
Diffusion models [3, 4, 6, 7, 13-15] generate images by progressively denoising a random variable drawn from a Gaussian distribution. In this work, we are particularly interested in the pre-trained text-to-image latent diffusion models [6]. Given a latent variable  $z$  and its noisy version  $z_{t}$  obtained by gradually adding noises to  $z$  for  $t$  steps, latent diffusion models optimize the following objective:
$$
\mathbb {E} _ {z, c, \epsilon \sim \mathcal {N} (0, 1), t} \left[ \| \epsilon - \epsilon_ {\theta} \left(z _ {t}, t\right) \| _ {2} ^ {2} \right], \tag {1}
$$
which is the squared error between the added noise  $\epsilon$  and the predicted noise  $\epsilon_{\theta}(z_t, t)$  by a neural model  $\epsilon_{\theta}$  at time step  $t$ , given  $c$  a text prompt as condition. During training, the latent variable  $z$  is obtained by passing the image into a pre-trained encoder [16]. For inference, a decoder is employed to convert the denoised latent into an image. In addition to text prompts, our model also conditions on subject representation, rendering an image generation architecture with multimodal conditions.
(a) Multimodal Representation Learning
(b) Subject Representation Learning
Figure 2: Illustration of the two-staged pre-training for BLIP-Diffusion. Left: in the multimodal representation learning stage, we follow prior work [12] and pretrained BLIP-2 encoder to obtain text-aligned image representation. Right: in the subject representation learning stage, we synthesize input images by composing subjects with random background. The BLIP-2 is then tasked to produce subject prompt embedding, which is later used by the latent diffusion model to generate output subject image. The image encoder remains frozen during pre-training.
2.2 Subject-driven Text-to-Image Generation
Given a few images of a subject, the task of subject-driven text-to-image generation aims at generating the subject in novel context based on text prompts. In the era of diffusion models, Textual Inversion [8] proposes to represent visual concepts using a placeholder text embedding, and optimize the embedding to reconstruct the subject images. DreamBooth [9] shares a similar methodology while additionally fine-tunes the diffusion model, which leads to better expressiveness and subject fidelity. One known drawback for both methods is their lengthy fine-tuning time for each new subject, which prevents the approaches from easily scaling up. More recent efforts attempt to reduce the required time and effort for fine-tuning. Concurrent to our effort, the work [11, 17, 18] pre-train the diffusion model on domain-specific images, such as cat and human face images. These models provide class-specific prior for generation thus being more efficient for fine-tuning. However, they are also constrained to a narrow list of subject categories and are not able to easily generalize to generic subjects. The work SuTI [19] proposes a knowledge distillation approach, which learns zero-shot generation from millions of fine-tuned expert models. Their model shows less flexibility in subject poses and is likely to be distracted by the background of the input images. In contrast, the pre-trained representation in our model is generic to a wide range of subjects, while generalizing efficiently to different subjects.
We propose BLIP-Diffusion, the first image diffusion model that features multimodal control through built-in generic pre-trained subject representation. Specifically, we adapt BLIP-2 encoder to extract multimodal subject representation, which is later used together with text prompt to guide generation.
We aim to learn subject representation that captures the subject-specific visual appearances while in the meantime aligns well with the text prompt. To this end, we propose a two-staged pre-training strategy as shown in Figure 2. First, a multimodal representation learning stage produces text-aligned generic image representation. Second, a subject representation learning stage prompts the diffusion model with text and subject representation for subject-driven generation. In this section, we delineate the model design and pre-training strategies.
3.1 Multimodal Representation Learning with BLIP-2
We use Stable Diffusion [6] as the latent diffusion model, which relies on CLIP [20] text encoder to produce prompt embeddings. In order to guide the generation using both text and subject representation as prompt, it is important that the subject embedding and the text embedding are well-aligned to ensure they can cooperate with each other. Inspired by the recent vision-language pre-trained model BLIP-2 [12], which produces high-quality text-aligned visual representation, we decide to adapt it to extract text-aligned subject representation.
Specifically, as shown in Figure 2a, we employ two main modules from BLIP-2 to learn multimodal representation: a frozen pre-trained image encoder to extract generic image features, and a multimodal
Figure 3: Left: example training image pairs, target images (top) and input images (bottom) with random background. Images are from the OpenImage dataset [22]. Right: comparison between the subject-driven image generation of our model (top), and image variations from DALLE-2 [4] (bottom). We prompt our model with the text prompt "a dog" and the subject representation, while feeding the input image into DALLE-2 for variation generation. Our model learns subject representation with minimal background information encoded, which enables faithful and flexible control when used together with text prompts.
encoder (i.e. Q-Former) for image-text alignment. The multimodal encoder is a transformer that accepts a fixed number of learnable query tokens and an input text. The query tokens interact with text through self-attention layers, and interact with frozen image features through cross-attention layers, and produces text-aligned image features as output. The output is of the same dimension as the number of query tokens. Empirically, we find that the originally implemented 32 output features often overpower the CLIP text embeddings when used in combination for image generation. Therefore, we instead half the number of query tokens and output 16 features.
Following BLIP-2 pre-training, we jointly train three vision-language pre-training objectives, including an image-text contrastive learning (ITC) loss that aligns the text and image representation by maximizing their mutual information, an image-grounded text generation (ITG) loss that generates texts for input images, and an image-text matching (ITM) loss that captures fine-grained image-text alignment via a binary prediction. We conduct multimodal representation learning on generic image-text paired data, which allows the model to learn a diverse set of visual and textual concepts.
3.2 Subject Representation Learning with Stable Diffusion
As the result of the multimodal representation learning, we obtain text-aligned visual representation of the input image. These features capture the generic semantic information of the input image. However, they are not specifically tailored to serve as guidance for the diffusion model. To this end, the subject representation learning stage aims to enable the diffusion model to leverage such visual representation, and generate different renditions of the subjects when combining with text prompts. In particular, we consider two desired properties when injecting the subject representation into a diffusion model. First, we expect the subject representation to well coordinate with the text prompts for the purpose of text-guided subject-driven generation. In this regard, prior methods [9, 18, 19] do not address the text prompts during training. They are thus not directly suitable to be used for scalable pre-training. Second, the behavior of the underlying diffusion model should ideally be maintained. This allows the subject-driven generation model to take advantage of techniques built on top of the original model on the fly, such as image editing and structure-controlled generation.
Model Architecture. The proposed model architecture is shown in Figure 2b. We connect the output of the BLIP-2 multimodal encoder to the input of the diffusion model's text encoder. During pretraining, the multimodal encoder takes as input a subject image and a text of the subject category, and produces a category-aware subject visual representation. We then transform the subject representation using a feed-forward layer consisting of two linear layers with GELU [21] activation in-between. The projected features are appended to the text prompt token embeddings as a soft visual subject prompt. Specifically, when combining the text token and subject embeddings, we use the template "[text prompt], the [subject text] is [subject prompt]". We pass text tokens through the CLIP embedding layer to obtain text token embeddings. We then concatenate subject embeddings and text token embeddings before passing them to the subsequent CLIP model layers. The resultant CLIP embeddings serve as guidance for the diffusion model to generate the output image. The soft visual prompt makes minimal architectural change to the underlying diffusion model, rendering an effective solution to inject subject representation while in the meantime largely inherits the modeling capabilities of the underlying diffusion model.
Subject-generic Pre-training with Prompted Context Generation. We aim to pre-train the model such that it learns to represent generic subjects from the input image. To this end, a naive approach
Figure 4: Left: BLIP-Diffusion cooperates with ControlNet for structure and subject controllable generation. Our model maintains the modeling capabilities of the underlying diffusion model and requires no re-training of the ControlNet parameters. Right: combined with cross-attention control techniques present in prompt-to-prompt, our model can be used for subject-driven image editing. The figure illustrates one denoising step where we mix latent maps generated with and without subject representation.
is to use the same image as both input to the multimodal encoder and output to the diffusion model as in [8, 9]. However, our preliminary experiments suggest that this leads to trivial solutions where generations are significantly interfered by the background in the inputs, or even models copying the input image as output, rendering generations not respecting the text prompts. On the other hand, while it is possible to collect multiple images of the same subject in different context, and thereby using different images as input and target, such an approach is laborious to scale up to generic subjects.
To address these issues, we propose a new pre-training task for learning subject-generic representation, called prompted context generation, where we curate input-target training pairs by synthesizing images of the subject in random background. The model takes the synthesized subject image as input, and aims to generate the original subject image as output according to a text prompt. Specifically, given an image containing a subject, we first feed the image and a category text of the subject to a text-prompted segmentation model CLIPSeg [23] with confidence thresholding. We then build a trimap by taking the segmentation map with higher confidence as the known foreground, the lower confidence as the uncertain region and the rest as the known background. Given the trimap, we use closed-form matting [24, 25] to extract the foreground, namely the subject. Then we compose the extracted subject onto a random background image via alpha blending. Finally, we use the synthetic image as the input and the original subject image as the output to serve as one training image pair.
As shown in Figure 3, such synthetic pairs effectively separate the foreground subject and the background context, preventing the subject-irrelevant information from being encoded in the subject prompt. In this way, we encourage the diffusion model to consider jointly the subject prompt and the text prompt for generation, leading to a pre-trained model that can be faithfully and flexibly controlled by both the subject image and the text prompt.
During pre-training, we freeze the image encoder and jointly train the BLIP-2 multimodal encoder along with the text encoder and U-Net of the latent diffusion model. To better preserve the original text-to-image generation capability, we find it beneficial to randomly drop the subject prompt at a  $15\%$  probability while using only text prompts to guide the diffusion.
3.3 Fine-tuning and Controllable Inference
The pre-trained subject representation enables both zero-shot generation and efficient fine-tuning for specific custom subjects. In addition, our model provides high-level visual control while inheriting the modeling capabilities of the underlying diffusion model. This enables us to leverage established image generation and editing techniques on the fly with BLIP-Diffusion as the foundation generation model. Below we first describe the efficient few-step subject-specific fine-tuning for custom subject generation. Then we present the extension capabilities of BLIP-Diffusion by incorporating existing techniques including ControlNet [1] and prompt-to-prompt image editing [2].
Subject-specific Fine-tuning and Inference. The pre-trained generic subject representation enables efficient fine-tuning for highly personalized subjects. Given a few subject images and the subject category text, we first use the multimodal encoder to obtain the subject representation individually. We then initialize the subject prompt embedding using the mean subject representation of all the subject images. In this way, we cache the subject prompt embedding without needing a forward pass of the multimodal encoder during fine-tuning. The diffusion model is fine-tuned to generate subject images as target by considering the text prompt embedding and the mean subject embedding. We also freeze the text encoder of the diffusion model, which we find helpful to counteract overfitting to target images. We use batch size 3 and a constant learning rate of 5e-5 with AdamW [26] optimizer across all the subjects, and generally observe decent results after 40-120 training steps, which takes 20-40 seconds to complete on a single A100 GPU.
Structure-controlled Generation with ControlNet. Our model introduces a multimodal conditioning mechanism for subject-control. In the meanwhile, the architecture is also compatible to integrate with ControlNet [1] to achieve simultaneous structure-controlled and subject-controlled generation. Figure 4 illustrates such integration, where we attach the U-Net of the pre-trained ControlNet to that of BLIP-Diffusion via residuals. In this way, the model takes into account the input structure condition, such as edge maps and depth maps, in addition to the subject cues. Since our model inherits the architecture of the original latent diffusion model, we observe satisfying generations using off-the-shelf integration with pre-trained ControlNet without further training.
Subject-driven Editing with Attention Control. Our model combines subject prompt embeddings with text prompt embeddings for multimodal controlled generation. Inspired by prompt-to-prompt [2], our model enables subject-driven image editing by manipulating the cross-attention maps of prompt tokens. In Figure 4, we show such capability where the model edits the original image with subject-specific visuals. For this purpose, we assume the generation process of the original image is known, or can be derived via inversion [13, 27] for real images. To edit the image, we first specify text tokens to edit, for example the token "dog". Next, we use the cross-attention maps of the specified token to extract automatically a mask for regions to edit. In order to preserve the layout and semantics in unedited regions, we keep the attention maps from the original generation while generating new attention maps for the inserted subject embeddings. We mix the denoising latents at each step based on the extracted editing mask. Namely, latents of the unedited regions are from the original generation whereas latents of the edited regions are from the subject-driven generation. In this way, we obtain the edited image with subject-specific visuals while also preserving the unedited regions.
4.1 Pre-training Datasets and Details.
For multimodal representation learning, we follow BLIP-2 [12] and pretrain the model on 129M image-text pairs, including 115M image-text pairs from LAION [28] with CapFilt [29] captions, COCO [30], Visual Genome [31] and Conceptual Captions [32, 33]. We use  $\mathrm{ViT}_{\mathrm{Large}}$  from CLIP [20] as the image encoder, and initialize Q-Former with  $\mathrm{BERT}_{\mathrm{base}}$  [34]. As aforementioned, we use 16 queries to learn subject representation. Other training hyperparameters follow [29].
For subject representation learning, we use a subset of 292K images from OpenImage-V6 [22], each containing a salient subject. We also remove images with human-related subjects. We use BLIP-2  $\mathrm{OPT}_{6.7\mathrm{B}}$  to generate captions as text prompts. We obtain a set of 59K background images from the web to synthesize subject inputs. We use Stable Diffusion v1-5 as the foundation diffusion model. We use a total batch size 16 with a constant learning rate 2e-6 for 500K steps using AdamW [26] optimizer, taking 6 days to finish on 16 A100 40Gb GPUs. More details on hyperparameters and data filtering procedure are included in Appendix for reference.
Main Qualitative Results. In Figure 5, we show qualitative generation results of BLIP-Diffusion. Thanks to the pre-trained subject representation, our model facilitates zero-shot subject-driven generation (row #1), producing meaningful results even for highly customized subjects. The model also enables efficient fine-tuning (row #3-6), demonstrating high-fidelity generation for re-contextulization, artistic stylization, textual modification, property modification and accessorization. Compared with
row #1-2
cat
on a skateboard
as plushie
car
painting by Van Gogh painted green
h painted green.
in a forest
with sunset
wearing top hat
on a red rug
Palace of Versailles
Few-step Fine-tuned Subject-driven Generation
seen from back
made of lego
at grand canyon.
decorated on cup
wearing glasses
Re-contextualization
at the Acropolis
in a jungle
in a comic
pencil sketch
printed on a shirt
Artistic Stylization
as book cover
oil painting
Texture Modification
gold sculpture
paper sculpture
as perfume atomizer
Property Modification
bottle as lamp
as teapot
in superman suit
Accessories
in police suit
in nurse suit
in iron man suit
in batman suit
row #7-8
teapot
depth map
a sofa
Structure Controlled Subject-driven Generation
canny edges
skateboarding at times square
HED edges
pouring milk in to a cup
row #9-10
#
attention map