LayoutDiffusion: Controllable Diffusion Model for Layout-to-image Generation

In plain words

Many image generators work from text alone. Text is vague about where things go, so a prompt like "a zebra to the left of a bus" often gives the wrong layout or drops an object. A layout is a more direct control: a list of boxes, each with a label such as "zebra" or "bus". The hard part is combining that box list with the picture the generator is building.

LayoutDiffusion does this with a diffusion model that is trained from scratch to take the box list as an input. The boxes are turned into tokens. A small transformer first lets the objects in the list look at each other, so the model sees the scene as a whole. Then, inside the image network, each region of the picture is given its own box position, and the picture regions attend to the objects. This lets each region find the object whose box it sits in.

The authors report that this gives sharper, more faithful objects than earlier GAN and diffusion methods on two datasets, and they describe a classifier-free setup that lets the strength of the layout be turned up or down at sampling time. The method needs labelled boxes for training, and the paper says that limits how well it can use large text-image pretraining.

TL;DR

LayoutDiffusion trains an image-space diffusion UNet from scratch on COCO-Stuff and Visual Genome layouts §3 §4.1. Each object is a bounding box plus a class; a six-layer transformer fuses the objects into a layout embedding §3.2. Object-aware cross-attention (OaCA) mixes that embedding with image patches that carry their own box positions §3.3. It is trained with 20% classifier-free dropout and sampled with DPM-Solver in 25 steps §3.4 app. C.1. Author-reported: COCO-Stuff 256 FID 15.61 and VG 256 FID 15.63 tab. 1. Relevant if you need box-level control and can train from scratch on box-labelled data. The code differs from the paper in several places, listed under Paper vs code.

Problem

Text-only diffusion models give little precise control over the position and size of several objects, and missing or misplaced objects are common §1 fig. 1. Layouts add spatial and category information, but earlier layout-to-image methods were mostly GANs, which can be unstable and collapse in mode §1 §2. The one prior diffusion method that accepts layouts, LDM, fuses them by cross-attention or concatenation and conditions only in its second stage §1 §2.

The paper's stated difficulty is that the image and the layout are different modalities: the layout constrains position, size and category more tightly than text does, and losing one layout token can remove an object entirely §1 §3.3.

Method

Overview §3. Four parts: (a) layout embedding of the input boxes and classes; (b) a layout fusion module (LFM) that lets objects interact; (c) an image-layout fusion module with a structural image patch and object-aware cross-attention (OaCA), plus a direct addition for global conditioning; (d) a layout-conditional diffusion model with classifier-free training and sampling fig. 2 §3.

Layout input §3.1. A layout is a set of n objects o_i = {b_i, c_i}. The box b_i = (x0, y0, x1, y1) lies in [0, 1]^4 and the class c_i is an integer id §3.1. The layout is padded to a fixed length k. The first slot is a whole-layout object o_l with box (0, 0, 1, 1) and class 0. Padding slots use box (0, 0, 0, 0) and class C+1 §3.1. Here C is the number of object classes, which the code sets to 185 for COCO-Stuff and 180 for VG app. C.1 tab. 4.

Layout embedding §3.1. The layout embedding is L = B_L + C_L, where B_L = b W_B and C_L = c W_C. W_B is a 4 by d_L matrix shared with the image patches. W_C is the class embedding. The paper uses d_L = 256 for the main models §3.1 tab. 4.

Layout fusion module §3.2. LFM is a transformer encoder over the k layout tokens: L' = LFM(L). It uses several self-attention layers §3.2. Token O'_1, at index 0 of L', is treated as the global summary of the whole layout. Tokens O'_i, i ≥ 2, are treated as objects in the context of the others §3.3.

Structural image patch §3.3. The image feature map I (height h, width w, channels d_I) is split into patches. Patch I_{u,v} is given the box b_{I_{u,v}} = (u/h, v/w, (u+1)/h, (v+1)/w), with u as the row index and v as the column index §3.3. The patch boxes are then embedded with the same W_B as the object boxes. Positional embeddings for image patches and layout objects are P_I = (b_I W_B) W_P and P_L = (B_L) W_P, with W_P a projection to d_I §3.3.

Global conditioning §3.3. The paper adds the global token to the image features at several resolutions, I' = I + O'_1 W, with W a projection §3.3. The shipped code does something else; see Paper vs code.

Object-aware cross-attention (OaCA) §3.3. For the image, queries, keys and values are Q_I, K_I, V_I = Conv(Norm(I)) §3.3. For the layout, keys and values are K_L, V_L = Conv(½ (Norm(C_L) + L')), where Conv is a 1x1 convolution and the class embedding is averaged with the fused layout §3.3. The paper then writes the attention inputs as concatenations over channels (Ψ1) and over length (Ψ2):

• Q = Ψ1(Q_I, P_L); the query has shape hw by 2 d_I §3.3.
• K = Ψ1(Ψ2(K_I, K_L), Ψ2(P_I, P_L)); the key has shape (hw + k) by 2 d_I §3.3.
• V = Ψ2(V_I, V_L); the value has shape (hw + k) by d_I §3.3.
• Attention is softmax(QK^T / sqrt(d_k)) V §3.3.

The paper's dimension statement (Q has hw rows) fits the image positional embedding P_I in Q, not P_L. The code uses P_I there; see Paper vs code.

Diffusion model §3.4 app. B.1. Forward noising is Gaussian, with a linear noise schedule over T = 1000 steps §3.4 app. B.1 LayoutDiffusion_large.yaml. The network predicts the noise epsilon. The paper's training objective is the DDPM simple loss, a mean-squared error between the true and predicted noise §3.4 app. B.1. The paper's experiments text says the variance schedule is fixed and linear §4.2; the code learns the variance. See Paper vs code.

Classifier-free training and sampling §3.4 app. B.2. During training, the layout is replaced by the null layout with a fixed probability, 0.2 in all runs app. C.1 tab. 4. The null layout is the whole-layout token plus padding §3.4. At sampling time the two predictions are combined. The main text writes eps_hat = (1 − s) · eps(x_t, t | l_phi) + s · eps(x_t, t | l) §3.4. Appendix B.2 writes eps_tilde = eps(c) + s (eps(c) − eps(∅)) app. B.2. The shipped code uses the appendix form; see Paper vs code.

Sampling speed §3.4 §4.4. The authors adapt DPM-Solver, a high-order solver for the diffusion ODE, to the conditional classifier-free sampler §3.4 app. C.2. The shipped sample script uses dpm_solver with 25 steps sample.bash.

Pseudocode, layout embedding and fusion (from §3.1 and §3.2). The paper gives no algorithm box; the sketch below follows the text, and the code's shapes are noted where they differ.

# Layout embedding and fusion  [§3.1] [§3.2]
# boxes (B, k, 4) as (x0, y0, x1, y1) in [0, 1]; classes (B, k) integer ids
# Index 0 is the whole-layout token; padding slots use the padding class
def embed_and_fuse_layout(boxes, classes):
    # Class embedding C_L, shape (B, k, d_L)  [§3.1]
    class_emb = class_table[classes]
    # Box embedding B_L = b W_B, shape (B, k, d_L)  [§3.1]
    box_emb = boxes @ W_B
    # Layout embedding L = B_L + C_L  [§3.1]
    layout = box_emb + class_emb
    # Fused layout L' from several self-attention layers  [§3.2]
    fused = layout_transformer(layout)
    # Global token O'_1 is fused[:, 0]  [§3.3]
    global_token = fused[:, 0]
    return class_emb, fused, global_token

Pseudocode, OaCA for one block (from §3.3). Written from the equations in §3.3. The paper's Q and K use the image positional embedding; the code is consistent with that and not with the paper's P_L in Q.

# Object-aware cross-attention, one block  [§3.3]
# image (B, C, h, w); class_emb, fused (B, k, C); boxes (B, k, 4)
# patch_boxes (hw, 4): one box per image cell, (x0, y0, x1, y1)
def object_aware_cross_attention(image, class_emb, fused, boxes):
    # Q_I, K_I, V_I from normalised image features  [§3.3]
    image_norm = group_norm(image)
    q_image, k_image, v_image = conv_qkv(image_norm)
    # Positional embedding of each image cell from its box  [§3.3]
    patch_pos = project_pos(box_proj(patch_boxes))       # P_I
    # Positional embedding of each object from its box  [§3.3]
    layout_pos = project_pos(box_proj(boxes))            # P_L
    # K_L, V_L from the average of normalised class and fused  [§3.3]
    mixed = (group_norm(class_emb) + fused) / 2
    k_layout, v_layout = conv_kv(mixed)
    # Queries: image content with image position  [§3.3]
    queries = concat_channels(q_image, patch_pos)
    # Keys: image and layout, each with its own position  [§3.3]
    keys = concat_length(
        concat_channels(k_image, patch_pos),
        concat_channels(k_layout, layout_pos),
    )
    # Values: image and layout content  [§3.3]
    values = concat_length(v_image, v_layout)
    # Scaled attention; code scale is 1 / sqrt(channels per head)  [§3.3]
    weights = softmax(queries^T @ keys / sqrt(d_k), dim=-1)
    attended = weights @ values^T
    # Residual connection, zero-initialised output projection
    # [code: layout_diffusion/layout_diffusion_unet.py#L544]
    return image + zero_init_proj(attended)

Training step (from §3.4 and app. B.1).

# One training example  [§3.4] [app. B.1] [code: train_util.py#L220-L236]
# image (B, 3, H, W); boxes, classes: padded layout
def training_step(image, boxes, classes):
    # With probability 0.2, replace the layout by the null layout  [app. C.1]
    if random() < 0.2:
        boxes, classes = null_layout(boxes, classes)
    # Sample a timestep and Gaussian noise  [app. B.1]
    t = uniform_int(1, 1000)
    noise = normal_like(image)
    # Forward-noise the image in closed form  [app. B.1]
    noisy = q_sample(image, t, noise)
    # Predict the noise and the variance from the network  [§3.4]
    # [code: layout_diffusion/gaussian_diffusion.py#L784]
    eps_hat, var_values = split(unet(noisy, t, boxes, classes), 2, dim=1)
    # Simple MSE on the noise  [§3.4]
    mse = mean((noise - eps_hat) ** 2)
    # Variance term, as in the code's RESCALED_MSE loss
    # [code: layout_diffusion/gaussian_diffusion.py#L799-L802]
    vb = variational_bound_term(eps_hat, var_values, noisy, t)
    # Scale by num_timesteps / 1000, which is 1 for the shipped T = 1000
    return mse + vb * 1000 / 1000

Sampling with classifier-free guidance (as shipped in the code).

# Guided noise prediction
# [code: scripts/classifier-free_sample_for_layout.py#L118-L125]
# s is sample.classifier_free_scale; shipped default 1.0
# [code: configs/COCO-stuff_256x256/LayoutDiffusion_large.yaml#L142]
def guided_mean(x_t, t, boxes, classes, s):
    # Conditional prediction with the real layout  [§3.4]
    cond = unet_mean(x_t, t, boxes, classes)
    # Unconditional prediction with the null layout  [app. B.2]
    uncond = unet_mean(x_t, t, null_boxes, null_classes)
    # Shipped form: cond + s (cond - uncond)
    # [code: scripts/classifier-free_sample_for_layout.py#L125]
    return cond + s * (cond - uncond)

Hyperparameters

Values are from the appendix table (Table 4, tab. 4) and the shipped configs. Columns: COCO-Stuff 256 (the main model), COCO-Stuff 256 small, COCO-Stuff 128, VG 256 and VG 128. The VG 128 config was not read. "All" means the same in every column.

Other settings, all from the paper's text or code:

Results

Author-reported. Setup: COCO-Stuff (25,210 training and 3,097 validation images) and Visual Genome (62,565 training, 5,062 validation, 5,096 test images) §4.1 app. D.1. Baselines are Grid2Im, LostGAN-v2 and PLGAN (GANs) §4.4. Table 1 values, with the reported ± where given tab. 1. COCO evaluation uses 3,097 layouts with 5 samples each; VG uses 5,096 layouts with 1 sample each tab. 1.

Source: tab. 1. "Not reported" means the paper leaves the cell blank. Grid2Im was only evaluated on COCO-Stuff tab. 1.

Ablation, COCO-Stuff 128, 300K iterations tab. 2. Each row removes or swaps a component; the last row is the full model.

Source: tab. 2. The paper's DS statement is in the ablation text §4.5.

Comparison with LDM, COCO-Stuff 256 tab. 3. Both LDM rows use a latent space. The authors generated the same 2,048 LDM images for a fair comparison tab. 3. The protocol differs from Table 1: the FID here is over 2,048 images, not 3,097 × 5.

Source: tab. 3. The appendix repeats this with a total column app. C.2 tab. 6. The hardware for the images/s column is not stated.

Headline claims [abstract] §1. The abstract's relative gains are FID and CAS on COCO-Stuff, and FID and CAS on VG, against the best baseline. They are not all from one row of Table 1. By our arithmetic from Table 1, the COCO-Stuff 256 FID gain against PLGAN is about 46.4% (29.10 to 15.61), which matches the abstract's 46.35%. The COCO-Stuff 256 CAS gain against PLGAN is about 26.8% (37.65 to 47.74), not the abstract's 26.70%. The VG CAS gain of 41.82% matches VG 256 against LostGAN-v2 (34.48 to 48.90). The IS gain of 29.29% in the text matches COCO-Stuff 128 against PLGAN (15.60 to 20.17). Each of these numbers is the authors' figure. We checked the arithmetic only.

Paper vs code

Code: github.com/ZGCTroy/LayoutDiffusion at 6dcc329a139025a58455679ddeabedf09ff7c3dd. The README says the code is "heavily based on openai/guided-diffusion" with three additions: distributed training, OmegaConf configs, and the layout encoder with OaCA README.md. Checked: the COCO-Stuff 256 main config, the COCO-Stuff 256 small config, the COCO-Stuff 128 main config, the VG 256 main config, the layout encoder, the OaCA block, the UNet forward, the training loop, the loss, and the sampling script.

• Variance. Paper: "We adopt the fixed linear variance schedule" §4.2. Code: model_var_type: "LEARNED_RANGE" and learn_sigma: True LayoutDiffusion_large.yaml. The training loss asserts LEARNED_RANGE gaussian_diffusion.py.
• Loss. Paper: the simple MSE on noise §3.4 app. B.1. Code: RESCALED_MSE, which is the MSE plus a variational term weighted by num_timesteps / 1000 gaussian_diffusion.py.
• Classifier-free formula. Main text: eps_hat = (1 − s) eps_uncond + s eps_cond §3.4. Appendix: eps_cond + s (eps_cond − eps_uncond) app. B.2. Code: cond + s * (cond - uncond) classifier-free_sample_for_layout.py. The code matches the appendix. With the code's form and scale s_code, the main-text form needs s_paper = 1 + s_code (derived from the two formulas). The shipped default s_code = 1.0 therefore equals s_paper = 2.0 in the main-text form.
• Global conditioning. Paper: add O'_1 W to the image features at several resolutions §3.3. Code: the projected global token (transformer_proj of token 0) is added to the timestep embedding, which feeds every block layout_encoder.py layout_diffusion_unet.py. Token O'_1 still reaches the image through OaCA.
• Query positional embedding. Paper: Q = Ψ1(Q_I, P_L) §3.3. Code: the queries carry P_I layout_diffusion_unet.py. The paper's own shape (Q has hw rows) matches P_I, not P_L.
• Patch box order. Paper: b_{I_{u,v}} = (u/h, v/w, (u+1)/h, (v+1)/w) with u a row index §3.3. Code: patch boxes are (column/r, row/r, (column+1)/r, (row+1)/r), so the first coordinate is x layout_encoder.py. The code keeps the (x0, y0, x1, y1) order that object boxes use.
• Attention heads. Paper's hyperparameter table: 4 heads for OaCA tab. 4. Code: heads = channels // num_head_channels, with num_head_channels = 64 LayoutDiffusion_large.yaml layout_diffusion_unet.py. Block widths are 256 times the channel multiplier layout_diffusion_unet.py. The COCO-Stuff 256 config attends at downsampling factors 8, 16 and 32 LayoutDiffusion_large.yaml. By the code's rule those widths (512, 1024, 1024) give 8, 16 and 16 heads.
• Padding mask. The paper does not describe a key-padding mask. The code has one, but use_key_padding_mask: False in all four configs read LayoutDiffusion_large.yaml layout_encoder.py. Padding and invalid object tokens therefore take part in both LFM and OaCA attention.
• Learning rate, VG 256. Paper's table: 1e-5 tab. 4. VG 256 config: lr: 2e-5 LayoutDiffusion_large.yaml. The paper's table also has 2e-5 for COCO 128 and VG 128, and the COCO 128 config matches. Which VG 256 value produced the reported model is not stated.
• Batch size and learning rate in the text. Paper text: "batch size 24, learning rate 1e-5" §4.2. The table gives 32 or 64 per run tab. 4. The COCO-Stuff 256 config uses 4 per GPU on 8 GPUs, total 32 LayoutDiffusion_large.yaml.
• Hardware. Paper text: "32 NVIDIA 3090s" §4.2. Appendix caption: "All trained on eight RTX 3090" tab. 4. The shipped training script uses 8 processes train_on_image_space.bash.
• Augmentation. Not in the paper. The shipped configs turn on left-right flip and min-IoU random crop for training LayoutDiffusion_large.yaml.
• Object limits. The COCO loader ignores images outside the 3 to 8 object range coco.py. The VG loader keeps at most 10 objects and fills spare slots with orphaned objects vg.py. The paper gives 3 to 30 boxes for VG §4.1 and does not say how extra boxes are handled.
• Null class id. In training, the null layout uses class num_classes − 1 for every slot except the whole-layout token, which uses class 0 train_util.py. That gives 184 for COCO-Stuff and 179 for VG.
• Padded length. The VG loader sets the layout length to max objects + 2 vg.py. The configs match this (10 and 12).
• Sampler. The paper's sampler is DPM-Solver §3.4. The shipped sample script runs dpm_solver with 25 steps sample.bash. The model function passed to DPM-Solver returns the guided mean only, not the learned variance classifier-free_sample_for_layout.py.
• Not checked. The VG 128 config, the GPU training configuration beyond the shipped script, and the evaluation scripts for FID, IS, DS, CAS and YOLOScore were not compared line by line. The README points to external repositories for FID, IS, DS and YOLOScore README.md.

Pitfalls

• Classifier-free scale. The same letter s means different things in the main text and the code. Use the code's form and default when you reproduce the released model, and state the scale you used. Appendix figure captions use 1.0, 0.6 and 0.5 for different panels, not one value app. A.4 app. A.5.
• Variance learning. Do not implement the paper's fixed-variance setting and expect the released weights to match. The code learns the variance with a variational term gaussian_diffusion.py.
• Null layout. The unconditional branch must use the same null layout as training: whole-layout token at index 0 with class 0 and box (0, 0, 1, 1), and padding for everything else. The code also sets is_valid_obj to 0 for the padding slots train_util.py.
• Padding attention. Because the shipped config does not mask padding, padding tokens affect the output. Changing the padded length or the padding class changes the model's inputs, not just its speed. Keep the code's length (max objects + 2).
• Attention scale. The paper writes sqrt(d_k) for OaCA without defining d_k §3.3. The code divides each of query and key by the fourth root of the per-head width, which includes the positional channels layout_diffusion_unet.py. The effective divisor is the square root of the per-head width of the query.
• Box convention. The paper's patch box formula orders the coordinates (row, column, row, column) with u as a row index, but the object boxes are (x0, y0, x1, y1) §3.1 §3.3. Use the code's order for patches too, or the patch positions will not line up with the object boxes.
• Heads and width. The head count changes with the block width in the code. Do not hard-code 4 heads, as the paper's table would suggest.
• Statistical noise. Table 2 gives FID for the full model (16.57) and for LFM with plain cross-attention (16.46), so the plain variant is 0.11 better on FID. No FID ± is reported, and the text does not discuss it tab. 2 §4.5. Treat FID differences of this size as not established. The rule of thumb to use is ours, not the paper's.
• Benchmark protocol. COCO FID for the headline uses 3,097 layouts with 5 samples each tab. 1. The LDM comparison uses 2,048 images tab. 3. The released README reports 15.61 for the first and 31.68 for the second README.md. The two numbers are not comparable.
• VG sample count. Table 1's caption says 5,096 test layouts, while the README says 5,097 tab. 1 README.md. Check which set you use.
• Dataset filtering. Results depend on the COCO-Stuff filtering. The appendix shows 25,210 training images with the deprecated 2017 subset and crowd filter, against 24,972 or 74,777 with other filters app. D.1 tab. 7. Use the filter the config encodes.
• Training from scratch. The authors say the model is trained from scratch on box-labelled data, and that combining it with text-to-image pretraining is future work §5. Expect no text prompts and no pretrained image prior.
• Cost. The paper reports 569M parameters, 0.308 images/s and 216.55 V100-days for the conditional stage tab. 3. The LDM comparison uses V100 days, not wall-clock time on one card.
• CUDA assumption. LayoutTransformerEncoder.__init__ moves the patch-box tensors to CUDA layout_encoder.py. CPU-only runs fail there.
• Latent variant. The repository also has latent-space configs (latent_LayoutDiffusion_*.yaml) that need a separate VAE download README.md. The paper's results are for image space only.
• Released weights. The COCO-Stuff 256 sample config names the 1,150,000-iteration checkpoint, which matches the appendix LayoutDiffusion_large.yaml. The VG 256 sample config names a 130,000-iteration checkpoint, while the appendix says 1.45M iterations LayoutDiffusion_large.yaml tab. 4. The README's VG links carry no iteration count README.md.

Lineage

• Builds on: DDPM and ADM-G, the guided-diffusion UNet and its training recipe §2 app. B.1. Classifier-free guidance, Ho and Salimans §3.4 app. B.2. DPM-Solver [arXiv:2206.00927], used for sampling §3.4 app. C.2. Transformer encoders from CLIP-style ViT code for the layout encoder layout_encoder.py.
• Compared against: LostGAN-v2, PLGAN and Grid2Im (GANs), and LDM §4.4 tab. 1 tab. 3. None of these has a card on this site.
• Related cards: GLIGEN adds box grounding to a frozen text-to-image diffusion model with gated layers. MultiDiffusion and DenseDiffusion add layout control to frozen models without training. LayoutGPT plans layouts with an LLM and then renders them. The eDiff-I card covers a different text-to-image cascade.
• Superseded by: not known from content already on this site.