How to add pixel-level loss

[Zhao-Xinda]

How can I add pixel-level loss in stage2, such as the MSE of the entire image? I see that the loss in the code is just diffusion noise. How can I change it to the MSE of the entire image?
Thanks！