hkproj / pytorch-stable-diffusion

Stable Diffusion implemented from scratch in PyTorch
https://www.youtube.com/watch?v=ZBKpAp_6TGI
MIT License
607 stars 138 forks source link

Question about the output of the Unet #3

Closed lixiaoqingnnz closed 12 months ago

lixiaoqingnnz commented 12 months ago

Hi, thank you so much for your work and sharing.

I want to retrain the Unet from the beginning, may I ask what is the output of the Unet according to this ddpm scheduler? Whether it is the random noise or the noise computed with the specific time step? I saw other codes, they use the random noise as the output of the model, but my model only works when I used the later one.

Highly appreciate your help!

hkproj commented 12 months ago

Hi! The input of the Unet is image + noise + timestep, the output of the Unet is noise (how much noise was added to the input). If you want to train your own Unet, I recommend training a Diffusion Model, and not a Latent Difusion Model. Please check my implementation of the DDPM paper at: https://github.com/hkproj/pytorch-ddpm This implementation also has a training code.