Text-to-Image System Design

Publish on: 2026-05-31 Classify at: AI/System Design

Words: 281 Read:≈ 2min Views:

Pinterest Canvas

Generate images conditioned on both text prompts and reference images.

Inference
1. Embed the text prompt to obtain the text embedding
2. Feed the noisy target latent and the text embedding into the DiT together to obtain the clean target latent
3. The clean target latent passes through the VAE decoder to produce the image
Training
- Training data: (text prompt, real image)
- Only the DiT is trained. Both the text encoder and the VAE use pretrained weights, and are then frozen
- The real image passes through the VAE encoder to obtain the real latent. During training, x_t is generated from the real latent

Stage 2 starts from the trained Stage 1 text-to-image model.

The reference image passes through the VAE encoder and becomes image latent tokens, which are then concatenated with the noisy target latent tokens to form image tokens
The text embedding acts as a condition and learns its own W_q/k/v; the resulting Q/K/V projections are concatenated with the image tokens’ Q/K/V projections to form the final Q/K/V used for self-attention
The DiT block outputs the hidden states of all tokens:
- hidden states of the target noisy latent tokens
- hidden states of the reference image tokens
- hidden states of the text tokens
However, in the end only the hidden states corresponding to the target noisy latent tokens pass through the final output head, and get mapped into the velocity v.