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