(arr, timesteps, broadcast_shape)
| 24 | |
| 25 | |
| 26 | def _extract_into_tensor(arr, timesteps, broadcast_shape): |
| 27 | # from: https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/gaussian_diffusion.py#L895 """ |
| 28 | res = arr[timesteps].float() |
| 29 | dims_to_append = len(broadcast_shape) - len(res.shape) |
| 30 | return res[(...,) + (None,) * dims_to_append] |
| 31 | |
| 32 | |
| 33 | def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): |