Lab 07 - UNet / DiT Conditioning

Overview

Lab 07: UNet / DiT Conditioning

Annotated code reading lab. Running code is optional.

Related handbook section

UNet / DiT Conditioning

This lab maps directly to the handbook section. Read the related handbook section first, then use the lab page and starter file to connect the concept to concrete variables, shapes, APIs, and interview-ready explanations.

Open handbook section Open starter file

Concept Goal

UNet / DiT Conditioning

Compare conditioning flow in convolutional UNet-style and Transformer/DiT-style denoisers.

Mental Model

Mechanism to keep in mind

`condition` is prepared once but read at many timesteps.
`cross_attention` or modulation injects it into the denoiser.
`prediction` keeps the same latent contract.

Annotated Code Preview

Starter preview

Excerpt from code/lab-07-unet-dit-conditioning/conditioning_flow.py. The linked starter file is the source of truth.

Open starter file

# UNet / DiT Conditioning
# Annotated reading material. Running this file is optional.
# Source-of-truth focus: Compare conditioning flow in convolutional UNet-style and Transformer/DiT-style denoisers.

condition = "text_encoder(prompt)"
latent_tokens = "patchify(noisy_latent)"  # DiT-style mental model
prediction = f"denoiser(latent={latent_tokens}, condition={condition}, timestep=t)"

# What to explain while reading:
# - condition is prepared once but read at many timesteps.
# - cross_attention or modulation injects it into the denoiser.
# - prediction keeps the same latent contract.
#
# Common traps:
# - Diffusion is not synonymous with UNet.
# - Conditioning is not just prompt string concatenation.

Line-by-line Explanation

What each block is doing

Setup / contract: `condition` is prepared once but read at many timesteps.
Main transition: `cross_attention` or modulation injects it into the denoiser.
Interview hook: `prediction` keeps the same latent contract.

What to Notice

Reading checkpoints

Backbone choice changes architecture, not denoising objective.
Text/image/video conditions need different encoders.
Condition shape must match the injection path.

Common Misunderstandings

What this lab prevents

Diffusion is not synonymous with UNet.
Conditioning is not just prompt string concatenation.

Interview Explanation

How to say it out loud

Compare conditioning flow in convolutional UNet-style and Transformer/DiT-style denoisers. Then explain the code by naming the state being transformed, the axis or shape that matters, and the tradeoff that would appear in a real system.

External intuition notes

Additional intuition

Use official docs and papers for API behavior and factual claims; use blogs only to improve the mental picture.
If support matrices, performance behavior or backend choices are version-sensitive, check current docs before repeating them.
A strong interview answer names the state object, the shape or axis it changes, and the tradeoff it creates.

InfraLens