(
xt,
v_pred,
sigma,
sigma_next,
cache_dtype=torch.float32,
return_denoised: bool = False,
)
| 85 | |
| 86 | @staticmethod |
| 87 | def step( |
| 88 | xt, |
| 89 | v_pred, |
| 90 | sigma, |
| 91 | sigma_next, |
| 92 | cache_dtype=torch.float32, |
| 93 | return_denoised: bool = False, |
| 94 | ): |
| 95 | dtype = xt.dtype |
| 96 | xt = xt.to(cache_dtype) |
| 97 | sigma = sigma.to(cache_dtype) |
| 98 | v_pred = v_pred.to(cache_dtype) |
| 99 | sigma_next = sigma_next.to(cache_dtype) |
| 100 | """ |
| 101 | # xt = x0 + sigma * (epsilon - x0) this line is pseudo code |
| 102 | x0 = xt - v_pred * sigma |
| 103 | epsilon = (xt - x0) / sigma + x0 |
| 104 | xt_next = x0 + sigma_next * (epsilon - x0) |
| 105 | # the above three lines are equivalent to the following line |
| 106 | """ |
| 107 | xt_next = (xt + v_pred * (sigma_next - sigma)).to(dtype) |
| 108 | if not return_denoised: |
| 109 | return xt_next |
| 110 | else: |
| 111 | x0 = (xt - v_pred * sigma).to(dtype) |
| 112 | return xt_next, x0 |
| 113 | |
| 114 | def retrieve_inference_timesteps_and_sigma_given_shift_scale( |
| 115 | self, num_inference_steps: int, shift_scale: int = 3, device=None): |
no outgoing calls
no test coverage detected