BLIP-Diffusion: Pre-trained Subject Representation for Controllable Text-to-Image Generation and Editing

In plain words

Suppose you want a picture of your own dog on a beach. Most ways of teaching a text-to-image model about one specific subject retrain part of the model for a while on that subject, and the result is a separate model copy per subject.

BLIP-Diffusion instead learns a general way to turn a photo of any subject into a short run of hidden "subject tokens". The image generator reads those tokens next to your text prompt, so a single photo can steer the result, often with no extra training. A short fine-tune on a few photos makes the result closer to one particular subject when that matters.

The model is trained in two stages. First, its image encoder learns to describe pictures in a form that lines up with text. Second, it is trained on made-up pairs: a subject is cut out of its photo and pasted onto a random background, and the model has to recreate the original from the pasted version and a prompt. That pushes it to focus on the subject rather than on what surrounds it. The same design also works with structure controls and with text-driven editing of existing images.

TL;DR

• Subject-driven text-to-image model built on Stable Diffusion v1-5 and BLIP-2 §3 §4.1.
• A pre-trained subject encoder gives zero-shot generation from one photo, and fine-tuning takes 40-120 steps; the authors report up to 20x more efficiency than DreamBooth §3.3 app. A.2.
• Author-reported DreamBench DINO: 0.594 zero-shot and 0.670 fine-tuned, against 0.668 for DreamBooth tab. 1.
• Pick it if you need subject fidelity with fast personalization. The stage-2 pre-training is not in the released code, so you cannot reproduce the base model from this repository.

Problem

• Inversion-based personalization (Textual Inversion, DreamBooth) needs hundreds to thousands of fine-tuning steps for each new subject §1.
• Most pre-trained text-to-image models do not accept images as a control input, so a learned subject embedding is hard to align with the text space while keeping the subject's look §1.
• Prior fine-tuning methods do not use the text prompt during training, so they are not directly suited to large-scale pre-training §3.2.
• Training with the same image as input and target leads to trivial solutions: the background leaks into the subject, or the model copies its input §3.2.

Method

Overview. A frozen BLIP-2 image encoder and a trainable BLIP-2 multimodal encoder (Q-Former) turn a subject image and its category word into a fixed set of subject tokens. A small projection maps these into the CLIP text-embedding space, where they are added to the prompt tokens. The Stable Diffusion U-Net then denoises as usual §3.2 fig. 2.

Step 1: multimodal representation (BLIP-2). Use a frozen ViT-L/14 image encoder from CLIP and a Q-Former initialised from BERT-base §4.1. The Q-Former has 16 learnable queries, half of the 32 used in BLIP-2, because 32 outputs overpowered the CLIP text embeddings during generation §3.1. Pre-train with the three BLIP-2 objectives: image-text contrastive, image-grounded text generation and image-text matching §3.1.

Step 2: project the subject tokens. Feed the 16 Q-Former outputs through a feed-forward block with two linear layers and a GELU between them §3.2. The projected tokens are appended to the prompt token embeddings as a soft visual subject prompt §3.2. The authors describe the prompt as a template of the form "[text prompt], the [subject text] is [subject prompt]" §3.2. The released code places the tokens differently; see Paper vs code.

Step 3: condition the U-Net. Pass the combined sequence through the CLIP text encoder of Stable Diffusion; its output conditions the U-Net §3.2. Only the Q-Former, the projection, the text encoder and the U-Net are involved, so the change to the base model stays small §3.2.

Step 4: build subject-generic training pairs (prompted context generation). This is paper-only; the repository has no code for it §3.2 fig. 3.

1. Take a subset of 292K OpenImages V6 images, each with a salient subject; filters remove groups, interiors, extreme aspect ratios and sizes, person- related classes and cluttered classes §4.1 app. A.4.
2. Run a text-prompted segmentation model (CLIPSeg) on the image with the category word, and keep only confident regions §3.2.
3. From the confident foreground, the uncertain band and the background, build a trimap §3.2.
4. Extract the subject with closed-form matting §3.2.
5. Alpha-blend the subject onto a random background from a set of 59K web images §3.2 §4.1.
6. Use the composite as model input and the original image as target, with a caption from BLIP-2 OPT-6.7B as the prompt §3.2 §4.1.

Step 5: subject representation pre-training (stage 2). Freeze the image encoder; train the Q-Former, the text encoder and the U-Net jointly §3.2. Drop the subject prompt with 15% probability so the model keeps plain text-to-image behaviour §3.2. Pre-training uses batch size 16, a constant learning rate of 2e-6, AdamW and 500K steps; it took 6 days on 16 A100 40GB GPUs §4.1.

Step 6: fine-tune one subject. Encode each of the few subject images with the multimodal encoder, then average the resulting embeddings into one cached subject embedding §3.3 app. A.5. Freeze the text encoder and fine-tune the U-Net on the subject images with the text prompt and the cached embedding §3.3. Use batch size 3 and AdamW §3.3 app. A.5. The fine-tune loss is the noise-prediction MSE of a standard latent-diffusion step; the paper gives no formula for it, so see the pseudocode below blip_diffusion.py.

Step 7: inference and extensions. Generation uses the PNDM scheduler and classifier-free guidance with scale 7.5 app. A.5. The model can attach a pre-trained ControlNet to its U-Net through residual inputs, without further training §3.3 fig. 4. For editing, the cross-attention maps of a chosen token give a mask; the unedited regions keep the original generation's latents, and the edited regions take the subject-driven result §3.3 fig. 4.

Pseudocode for the conditioning path (steps 1 to 3), as implemented:

# Subject tokens from the BLIP-2 multimodal path  [§3.1] [§3.2]
# Input:  subject_image (B, 3, 224, 224); subject_text: one word per image
# Output: subject_tokens (B, 16, 768)
def encode_subject(subject_image, subject_text):
    # Frozen ViT features, then 16 text-aligned Q-Former outputs  [§3.1]
    # [code: lavis/models/blip_diffusion_models/blip_diffusion.py#L86-L90]
    query_outputs = qformer(subject_image, subject_text)  # (B, 16, 768)
    # Projection into the text embedding space  [§3.2]
    # Code also adds a LayerNorm and a residual; see Paper vs code
    subject_tokens = project(query_outputs)               # (B, 16, 768)
    return subject_tokens

# Splice subject tokens into the prompt
# [code: lavis/models/blip_diffusion_models/modeling_ctx_clip.py#L223-L233]
# Input:  token_embeds (S, 768); subject_tokens (16, 768)
# Output: conditioning (S + 16, 768)
def splice_subject(token_embeds, subject_tokens, begin=2):
    # Insert after the first two tokens; nothing is removed
    head = token_embeds[:begin]
    tail = token_embeds[begin:]
    return concatenate(head, subject_tokens, tail)

# One fine-tuning step, noise-prediction MSE
# [code: lavis/models/blip_diffusion_models/blip_diffusion.py#L224-L268]
# Input:  target_image (B, 3, 512, 512); input_image (B, 3, 224, 224)
# Output: scalar loss
def training_step(target_image, input_image, caption, subject_text):
    # Encode the target with the frozen VAE and apply its scale
    # [code: lavis/models/blip_diffusion_models/blip_diffusion.py#L225-L226]
    latents = vae_encode(target_image) * 0.18215
    # Add Gaussian noise at a random diffusion timestep  [§3.3]
    noise = random_normal_like(latents)
    timestep = random_timestep()
    noisy_latents = add_noise(latents, noise, timestep)
    # Subject tokens and the caption go through the text encoder  [§3.2]
    subject_tokens = encode_subject(input_image, subject_text)
    text_embeds = clip_text(caption, subject_tokens)
    # The U-Net predicts the noise; MSE against the sampled noise
    noise_pred = unet(noisy_latents, timestep, text_embeds)
    return mse(noise_pred, noise)

# Cached subject embedding for fine-tuning  [§3.3] [app. A.5]
# [code: blip_diffusion.py#L193-L215]
# Input:  subject_images (N, 3, 224, 224), N few-shot photos
# Output: cached_embedding (1, 16, 768)
def cache_subject(subject_images, subject_text):
    # Encode each photo, then average over the photos
    tokens = encode_subject(subject_images, subject_text)  # (N, 16, 768)
    return tokens.mean(dim=0, keepdim=True)                 # (1, 16, 768)

Hyperparameters

Per-subject fine-tuning steps on DreamBench tab. 3. These are the counts the authors used; mean 76 app. A.5.

Results

Author-reported. Setup: DreamBench, 30 subjects with 4-7 images each §4.2 app. A.3. Four images per prompt, 3,000 images in total, averaged over 10 fixed sets of random seeds §4.2. Base model Stable Diffusion v1-5 §4.1. Timings are on a single A100 §3.3 app. A.5.

DreamBench quantitative comparison tab. 1. The ± values are as printed in the table; its caption describes them as differences across 10 runs.

Ablation, zero-shot, after 250K subject pre-training steps tab. 2 §4.2:

The paper also reports qualitative results, zero-shot style transfer, subject interpolation and failure cases fig. 5 fig. 8 fig. 9 fig. 10. Failure modes include the model's own weak handling of fine-grained composition relations and of text prompts §5. Both image-text and subject alignment rise with more subject pre-training steps, shown in the right panel of tab. 2 §4.2.

Paper vs code

Official code: the LAVIS repository, the blip-diffusion project. The paper links it, and the authors are at Salesforce README.md. Released at commit 05cb670f47296522e24e7727e366b4a997d28582, BSD-3-Clause.

• Stage-2 pre-training is absent. The repository has fine-tuning configs only (lavis/projects/blip_diffusion/) plus a download of the released checkpoint blip_diffusion_base.yaml. The prompted-context-generation pipeline, the 15% prompt dropout, the 500K-step schedule and the CLIPSeg/matting stage are not in the code. Treat these numbers as unverified against code.
• Projection block. Paper: two linear layers with GELU §3.2. Code: LayerNorm, Linear 768→3072, QuickGELU, Linear 3072→768, dropout 0.1, and the un-normalised input added back as a residual [code: lavis/models/blip_diffusion_models/blip_diffusion.py#L38-L56].
• Prompt construction. Paper: "[text prompt], the [subject text] is [subject prompt]" §3.2. Code: prompt "a {subject} {prompt}", repeated 20 times joined by commas blip_diffusion.py blip_diffusion.py. The 16 subject tokens are placed after the BOS token and "a" [code: lavis/models/blip_diffusion_models/modeling_ctx_clip.py#L223-L233]. The paper never states the insertion position.
• Inference steps. Paper: PNDM, 100 steps app. A.5. Code: generate defaults to 50 steps blip_diffusion.py; the README examples use 25 and 50 README.md README.md.
• Subject embedding during fine-tuning. Paper: initialise from the mean and fine-tune the diffusion model §3.3. Code: the cached embedding becomes a trainable nn.Parameter with requires_grad=True [code: lavis/models/blip_diffusion_models/blip_diffusion.py#L215]. Whether the optimiser updates it was not checked in the runner code.
• Fine-tune learning rate. Paper says 5e-5 in §3.3 and 5e-6 in app. A.5. Code uses 5e-6 finetune-db-dog.yaml.
• Fine-tune example. The dog example config trains on the DreamBench dog8 folder for 40 iterations finetune-db-dog.yaml finetune-db-dog.yaml. This matches the table entry for dog8 tab. 3.
• Editing hyperparameters. The paper gives no numeric editing thresholds. Code defaults are in the Hyperparameters table above.
• Metric code. The DINO, CLIP-I and CLIP-T scoring code is not in the repository. A search of the Python files for DINO found no scoring code.
• Subject-prompt dropout. The forward pass in code has no dropout of the subject prompt blip_diffusion.py. Only the paper describes it §3.2.

The guidance scale (7.5) and the frozen text encoder during fine-tuning match between paper and code blip_diffusion.py.

Pitfalls

• Learning rate conflict. §3.3 says 5e-5; the appendix and the code say 5e-6. Use 5e-6 unless you have a reason to test otherwise app. A.5.
• Steps range. The text says 40-120 steps §1 §3.3. The per-subject table lists dog5 at 20 steps tab. 3. Plan for 20-120.
• Zero-shot versus fine-tuned Textual Inversion. §4.2 says zero-shot results beat fine-tuned Textual Inversion. The table agrees on DINO (0.594 vs 0.569) and CLIP-T (0.300 vs 0.255). On CLIP-I it does not: 0.779 against 0.780 tab. 1.
• Ablation step count. The ablation runs use 250K steps, while the main model uses 500K §4.2 §4.1.
• Prompt splicing. The subject tokens are inserted, not substituted, and the prompt is lengthened by 16 positions. In the code, text is truncated to 61 tokens when the query tokens are included, because the total length is capped at 77 blip_diffusion.py. The code comment says a special token is removed, but the slice keeps it [code: lavis/models/blip_diffusion_models/modeling_ctx_clip.py#L223-L233].
• Prompt amplification. The code repeats the prompt 20 times. This is a deliberate trick, not in the paper, and it changes the text conditioning [code: lavis/models/blip_diffusion_models/blip_diffusion.py#L270-L276].
• Hidden fine-tune cost. Checkpoints are selected by hand against validation prompts app. A.5. The 40-120 step count does not include this selection.
• Timings. §3.3 says 20-40 seconds per subject. app. A.5 says about 40 seconds on average. Both are single-A100 figures.
• Data restrictions. Person-related subjects are removed from the pre-training data on purpose. The authors say this makes person-as-subject control weaker app. A.1 app. A.4.
• Base model inheritance. The model keeps the weaknesses of Stable Diffusion, including bias and NSFW output. The authors advise content moderation app. A.1.
• Figure and table numbers. Numbers were taken from source order, not from a compiled PDF. Check them against the arXiv PDF before citing.
• Attribution. Metrics are author-reported and follow the DreamBooth metric set app. A.3. The repository does not include the scoring code.

Lineage

• Builds on: BLIP-2 (arXiv:2301.12597) for the multimodal encoder and its three objectives §3.1; Stable Diffusion, the latent diffusion model (arXiv:2112.10752) §3.2.
• Compared against: DreamBooth (arXiv:2208.12242), Textual Inversion (arXiv:2208.01618), Re-Imagen (arXiv:2209.14491), InstructPix2Pix (arXiv:2211.09800) §4.2 app. A.2.
• Related fine-tuning methods: Custom Diffusion (/papers/2212.04488), which trains cross-attention weights for personalization.
• Combined with: ControlNet (arXiv:2302.05543) and Prompt-to-Prompt (arXiv:2208.01626) for structure-controlled generation and editing §3.3.
• Superseded by: not known from content in this wiki.