timesteps designed according to VAE compression

if spatial_compression_ratio >= 16:
mask_conditions_bs = mask_conditions.size()[0]
mask_conditions[:, :, 1:, :, :] = 1
if not mask_conditions[:, :, 0, :, :].any():
noisy_latents = (1 - mask_conditions) * inpaint_latents[:, -vae.latent_channels:] + mask_conditions * noisy_latents

    temp_ts = (mask_conditions[:, 0, :, ::2, ::2] * timesteps[:, None, None, None]).flatten(1)
    timesteps = torch.cat([temp_ts, temp_ts.new_ones(mask_conditions_bs, seq_len - temp_ts.size(1)) * timesteps[:, None,]], dim = 1)
else:
    timesteps = mask_conditions.new_ones(mask_conditions_bs, seq_len) * timesteps[:, None,]

The model code for training 5b has timesteps designed according to VAE compression, but no corresponding operation is seen during inference. Is this as expected?

Start by locating the training model code for 5b and the corresponding inference entry point. Compare how timesteps are constructed with and without VAE compression, including the mask_conditions path shown in the issue. Done means establishing whether inference intentionally differs and, if not, identifying the matching behavior needed.