Clarification on where noise is added to VAE latent before diffusion

[RosNayak]

Hi, really impressive work!

While exploring the codebase, I’m trying to precisely understand the point at which noise is injected into the VAE latent space before diffusion.
From what I gather, encode_first_stage produces the latent representation of the input image, but it’s not immediately obvious where in the pipeline the noise is applied to this latent before it is passed to the UNet for denoising during the diffusion process.

init_cnt = model.get_first_stage_encoding(model.encode_first_stage(init_cnt)) cnt_z_enc, _ = sampler.encode_ddim(init_cnt.clone(), num_steps=ddim_inversion_steps, unconditional_conditioning=uc, \ end_step=time_idx_dict[ddim_inversion_steps-1-start_step], \ callback_ddim_timesteps=save_feature_timesteps, img_callback=ddim_sampler_callback)

Could you please clarify:

1. When is the latent noised?
2. Which module or function specifically handles this?

[jiwoogit]

Thanks for your interest.

We generate a noised latent (inversed latent) by applying DDIM inversion to the encoded image, which is handled by sampler.encode_ddim().

Specifically, the function at the following line adds the inverted noise at each timestep:

StyleID/ldm/models/diffusion/ddim.py

Line 328 in d4ea655

img, pred_x0 = self.reverse_ddim(img, t, t_next=steps[i+1] ,c=conditioning, unconditional_conditioning=unconditional_conditioning, unconditional_guidance_scale=unconditional_guidance_scale)