I have trained a diffusion model to predict x_0 instead of predicting noise (follow this paper: [2303.09556] Efficient Diffusion Training via Min-SNR Weighting Strategy). Then I would like to use DDIMInverseScheduler(prediction_type=“sample”) to gradually add noise and then reconstruct the image using DDIMScheduler(prediction_type=“sample”). One thing that I observed is that in DDIMInverseScheduler from diffusers/src/diffusers/schedulers/scheduling_ddim_inverse.py at main · huggingface/diffusers · GitHub. At step t=0, they set beta_prod_t = 1 - alpha_prod_t = 0. So, line 349 is divided by 0 resulting in -inf or inf vector. The figure is shown below.
Is there any trick to avoid this issue? Or Is it just theoretically impossible to perform DDIM inversion with the predictive x_0 network?
Thank you very much