Latent Scaling Round Trip

Problem Statement

Implement the latent normalisation that sits between Stable Diffusion's VAE and its U-Net, including the empirical derivation of the famous 0.18215 constant.

Background

A latent diffusion model runs the diffusion process on VAE latents, not pixels. But the forward process assumes unit-variance data: the schedule is built so that $x_t$ has variance $\bar{\alpha}_t \cdot \mathrm{Var}(x_0) + (1-\bar{\alpha}_t)$, which only lands on 1 at $t = T$ if $\mathrm{Var}(x_0) = 1$. A raw VAE latent has a standard deviation of roughly 5.5, so it must be normalised first:

$z_{\text{scaled}} = (z - \text{shift}) \cdot s, \qquad z = \frac{z_{\text{scaled}}}{s} + \text{shift}$

The scaling factor is estimated once from the training set as $s = 1/\mathrm{std}(z)$ -- which for SD 1.x came out at 0.18215, i.e. a latent std of about 5.49. SD 3 and Flux add a non-zero shift_factor to centre the latents too.

The pipeline is: encode, scale, diffuse; then at the end unscale and decode. Forgetting either half is a classic bug -- skip the scale and the U-Net sees data at 5x the variance it was trained on; skip the unscale and the VAE decodes a washed-out grey image.

Your Task

Implement three functions:

def compute_scaling_factor(z):
def scale_latents(z, scaling_factor, shift_factor=0.0):
def unscale_latents(z_scaled, scaling_factor, shift_factor=0.0):

compute_scaling_factor returns a Python float, the other two return arrays shaped like their input.

Input Format

• z, z_scaled: NumPy arrays of any shape.
• scaling_factor (float): non-zero.
• shift_factor (float): defaults to 0.0.

Output Format

A float, and two arrays.

Sample

z = np.arange(-4.0, 5.0)
f = compute_scaling_factor(z)
print(round(f, 4))
print(round(float(np.std(scale_latents(z, f))), 4))

The std of [-4 ... 4] is 2.582, so the factor is 1/2.582 = 0.3873, and scaling by it makes the std exactly 1.