Composer: Creative and Controllable Image Synthesis with Composable Conditions

In plain words

Image generators that work from a text prompt can make striking pictures, but a designer often needs more control: the exact layout, the palette, the shape of one object. Adding one kind of control at a time, such as a sketch or a depth map, is a partial answer.

Composer takes a different route. Before training, each training image is taken apart into several simpler pieces: its caption, its overall look, its colors, its outline, its objects, its depth and a grayscale version. The model is then trained to rebuild the picture from some random subset of those pieces. Because it practises rebuilding from partial pieces, it can later be handed pieces from different pictures and still produce a coherent image. Pieces that are left out simply are not used, and the user can swap, keep or mix them to steer the output.

The point is that many combinations of a small set of pieces give a large space of possible images, which is what makes the approach flexible for creative work.

TL;DR

Composer is a diffusion image model trained to reconstruct images from random subsets of eight representations: caption, CLIP embedding, color histogram, sketch, instance masks, depth, grayscale and masked image §2.2. At inference, the user chooses which representations to keep, change or interpolate, which gives control over layout, palette, style and an editable region §3.2. The authors report a COCO FID of 9.2 and a CLIP score of 0.28 for text-to-image generation §3.5. The official repository currently contains no model code or weights README.md.

Problem

Text-conditioned generators give limited control, and they struggle to meet specifications for semantics, shape, style and color together §1. Earlier extensions add one condition type at a time, such as segmentation maps, scene graphs, sketches, depth maps or inpainting masks, or fine-tune per subject §1. The authors argue that composition of several conditions is what widens the control space, since with several representations per image the number of combinations grows exponentially §1.

The approach assumes a very large training corpus and substantial compute. The base model is trained on about 1B images from a filtered mix of public datasets §3.1. The hardware, training time and total compute are not stated.

Method

Composer is a multi-conditional diffusion model with a UNet backbone. It follows the GLIDE architecture and changes only how conditions enter §2.3. Each training iteration has two phases. In the decomposition phase, the images in the batch are split into representations by vision models. In the composition phase, the model is optimized to reconstruct each image from a subset of its representations §1.

1. Diffusion objective §2.1 eq. 1. Training minimises the mean squared error between the true noise and the noise predicted from the noisy image a_t x_0 + sigma_t eps and the conditions c. The timestep is written as t ~ U(0, 1). Classifier-free guidance combines a conditional and an unconditional noise prediction §2.1 eq. 2. Guidance with two different condition sets is defined in §2.1 eq. 3.

2. Decomposition §2.2. All representations are extracted on the fly during training:

• Caption: the title or description from the training data, or a captioning model if none exists. It is embedded as sentence and word embeddings with CLIP ViT-L/14@336px §2.2.
• Semantics and style: the CLIP image embedding from the same model, as in unCLIP §2.2.
• Color: a smoothed CIELab histogram, quantized to 11 hue, 5 saturation and 5 lightness values, with smoothing sigma 10 §2.2.
• Sketch: an edge detector followed by a sketch simplification step §2.2.
• Instances: instance masks from YOLOv5 §2.2.
• Depth: a monocular depth estimate §2.2.
• Intensity: a grayscale image, using RGB channel weights drawn at random from a predefined set §2.2.
• Masking: four channels, the masked RGB image and a binary mask marking the editable region §2.2.

3. Conditioning inside the model §2.3 app. B fig. 7 fig. 8.

• Global conditions (CLIP sentence embedding, image embedding, color palette) are projected and added to the timestep embedding. The image embedding and palette are also projected into eight extra tokens, which are concatenated with the CLIP word embeddings and used as the cross-attention context.
• Localized conditions (sketch, segmentation, depth, intensity, masked image) are projected by stacked convolutions to the spatial size of the noisy latent x_t. The embeddings are summed and concatenated to x_t before the UNet.
• A missing condition is handled by omitting its term: global conditions are added or masked in cross-attention, local ones are summed §2.3.

4. Joint training §2.3. Each condition is dropped independently with probability 0.5. All conditions are dropped with probability 0.1, and all are kept with probability 0.1. Intensity images use dropout 0.7, because they carry most of the image information and would otherwise overpower the other conditions §2.3.

5. Cascade and prior §2.3 §3.1 tab. 1. A 64x64 base model is followed by two upsamplers, to 256x256 and then 1024x1024. The architectures of the upsamplers are modified from unCLIP, with more channels in the low-resolution layers and self-attention blocks added §2.3. An optional prior maps a caption to an image embedding. The authors say it improves diversity for some combinations of conditions §2.3.

6. Inference operations §2.1 §3.2 §3.3.

• Guidance with two condition sets eq. 3: conditions in the second set are emphasised with weight omega, conditions in the first set are suppressed with weight 1 - omega, and shared conditions get weight 1.0 §2.1.
• Bidirectional guidance: invert an image to noise under one condition set, then sample from that noise under another set. The difference between the two sets defines the edit §2.1.
• Reconfiguration: the same inversion-and-resampling, applied to a modified set of representations, changes one aspect while keeping the rest §3.2.
• Variations and interpolation use subsets of representations, or interpolation in the global embedding space §3.2 fig. 2.
• Editable region: a masked image is added as a condition, so changes stay inside the mask §3.2 fig. 4.
• Reformulated tasks §3.3 fig. 5: palette colorization by reconfiguring the palette; style transfer by taking style from one image and content from another; image translation by reconfiguration with a text target; pose transfer from a segmentation map and an image embedding; virtual try-on from a masked body image and a garment embedding. The authors describe try-on as of moderate quality §3.3.

# Training step  [§2.2] [§2.3] [§3.1]
# Input:  images (B, 3, H, W)
# Output: scalar denoising loss

def training_step(images):
    # 1. Extract the representations of every image  [§2.2]
    conditions = decompose(images)

    # 2. Drop conditions per sample: 0.1 all, 0.1 none, else  [§2.3]
    #    independently with p = 0.5 (0.7 for intensity)
    conditions = drop_conditions(conditions)

    # 3. Noise the image at a sampled timestep  [§2.1]
    timesteps = sample_uniform(0, 1, size=B)
    noise = sample_gaussian(shape=images.shape)
    noisy = alpha(timesteps) * images + sigma(timesteps) * noise

    # 4. Predict the noise and take the squared error  [§2.1] [eq. 1]
    predicted = unet(noisy, conditions, timesteps)
    return mean((noise - predicted) ** 2)

# Classifier-free guidance step  [§2.1] [eq. 2]
# Weight omega is applied to the conditional prediction

def guided_noise(noisy, conditions, t, omega):
    # Prediction with all conditions present  [eq. 2]
    conditional = unet(noisy, conditions, t)
    # Prediction with conditions dropped  [eq. 2]
    unconditional = unet(noisy, None, t)
    return omega * conditional + (1 - omega) * unconditional

# Reconfiguration: invert, then resample with new conditions  [§2.1] [§3.2]
# Inputs: image, old conditions c_i, new conditions c_j

def reconfigure(image, old_conditions, new_conditions):
    # Deterministic DDIM inversion of the image under c_i  [§2.1]
    latent = ddim_invert(image, old_conditions)
    # DDIM sampling from that latent under c_j  [§2.1] [§3.2]
    return ddim_sample(latent, new_conditions)

Key equations, as numbered in the paper's LaTeX source (eqnarray rows):

\mathcal{L}_{\text{simple}} = \mathbb{E}\left(\|\bm{\epsilon} - \bm{\epsilon}_{\theta}(a_t \mathbf{x}_0 + \sigma_t \bm{\epsilon}, \mathbf{c})\|_2^2\right) eq. 1

\hat{\bm{\epsilon}}_{\theta}(\mathbf{x}_t, \mathbf{c}) = \omega \bm{\epsilon}_{\theta}(\mathbf{x}_t, \mathbf{c}) + (1 - \omega) \bm{\epsilon}_{\theta}(\mathbf{x}_t) eq. 2

\hat{\bm{\epsilon}}_{\theta}(\mathbf{x}_t, \mathbf{c}) = \omega \bm{\epsilon}_{\theta}(\mathbf{x}_t, \mathbf{c}_2) + (1 - \omega) \bm{\epsilon}_{\theta}(\mathbf{x}_t, \mathbf{c}_1) eq. 3

Hyperparameters

Model sizes, schedule and optimiser settings for each stage tab. 1 app. A:

The table caption says DPM-Solver++ is the sampler for all diffusion models tab. 1. Section 3.5 states the step counts for the prior, base and 64 to 256 models, and applies a guidance weight of 3.0 to the first two §3.5. The 256 to 1024 entry (10 steps) is not repeated in the text.

Other settings, from the body of the paper:

Results

Author-reported. Setup: COCO text-to-image generation, using the caption as the only condition, compared against state-of-the-art text-to-image models §1 §3.5. The paper does not name the baselines in the text and gives no baseline numbers.

The other tasks in the paper are shown qualitatively, without metrics: variations and interpolations fig. 2, reconfiguration fig. 3, region editing fig. 4, reformulated tasks fig. 5, compositional generation fig. 6, conflict resolution between conditions fig. 9, colorization fig. 10 and style transfer fig. 11. For reconstructions, the paper says that when more conditions are used, it reconstructs images more accurately than unCLIP §3.2, without a number.

Paper vs code

The official repository at the recorded commit contains a README, an MIT license (copyright line "damo-vilab", 2023) and images, and no model or training code README.md LICENSE. Its TODO list still has these items open: training and inference code, pretrained models, a Gradio UI, and a lighter variant built on Stable Diffusion 2.1 README.md.

• The README describes Composer as a 5 billion parameter model [code: README.md#L10]. The paper's stated sizes are 2B base, 1.1B and 300M upsamplers, and 1B prior, which sum to 4.4B §3.1 tab. 1. The paper does not state a total, so the difference is not resolved.
• The abstract says code and models "will be made available" (abstract). Hence every hyperparameter above is paper-only and cannot be checked against code.
• The project page linked from the arXiv listing at a damo-vilab address returned 404 when checked. The current project page is ali-vilab/composer-page README.md. The code field points to ali-vilab/composer, the official repository.

Pitfalls

• Timestep sampling. Equation 1 writes t ~ U(0, 1), but the table lists 1000 diffusion steps and two noise schedules §2.1 tab. 1. How continuous t maps onto those steps is not stated.
• Dropout order. The 0.1 keep-all branch and the 0.5 per-condition dropout are described in one sentence. Whether per-condition dropout also applies inside the keep-all branch, or in what order the branches are drawn, is not stated §2.3.
• Guidance convention. Equation 2 gives weight omega to the conditional prediction and 1 - omega to the unconditional one. With that form, the stated scale of 3.0 §3.5 means omega = 3.0, which is the usual "scale" convention, but the paper does not say so in words. Code is not available to confirm.
• Upsampler conditioning. The upsamplers are called "unconditional" diffusion models §2.3. The text does not say whether they receive the low-resolution image as input. Assume they do, and verify against the figures only.
• Masked-region editing. The paper says pixels outside the editable region are not modified §1, but the mechanism that guarantees this (for example, compositing after sampling) is not described §3.2.
• Intensity channel weights. The set of RGB weights for the grayscale representation is "predefined" but not listed §2.2. Implementers have to choose their own set.
• Color bins. The histogram quantization (11, 5, 5) gives 275 bins if multiplied out; this count is our arithmetic, not the paper's. The smoothing kernel behind "sigma 10" is not defined §2.2.
• Cost. The base model alone has 2B parameters §3.1, and sampling runs four stages with 100, 50, 20 and 10 steps in the table tab. 1. Memory use and latency are not stated.
• Evaluation details. The FID and CLIP score protocols (reference set, number of generated samples, CLIP model) are not stated §3.5.
• Inversion steps. Reconfiguration relies on DDIM inversion §2.1 §3.2. The number of inversion steps is not stated.
• Known weaknesses. Without the global embedding, sketch- and depth-only generation often produces dark images. Text can be down-weighted when image and text embeddings disagree §5. The authors also note misuse risk and say they plan a filtered version before public release §5.
• Try-on. The paper calls the virtual try-on results "moderate quality" §3.3.

Lineage

• Builds on: GLIDE, the UNet and conditioning base §2.3 [Nichol2021GLIDETP in the paper's bibliography]; unCLIP / "Hierarchical Text-Conditional Image Generation with CLIP Latents", for CLIP embeddings, the upsamplers and the prior §2.3 [ramesh2022 in the bibliography]; DDIM, for inversion and sampling §2.1 [Song2020DenoisingDI]; DPM-Solver++ tab. 1. The paper's bibliography gives no arXiv identifiers, so no wiki links are added.
• Compared against: unCLIP, qualitatively, for reconstructions §3.2. The COCO comparison names no baselines §3.5.
• Superseded by: not known from content in this wiki.