Question about Implementation of EDMEulerScheduler

[KeepNoob]

First of all as shown in HuggingFace [https://huggingface.co/docs/diffusers/api/schedulers/edm_euler](official doc) the EDMEulerScheduler is a implementation of Algorithm 2 from the Elucidating the Design Space of Diffusion-Based Generative Models paper. Here is my question in the step function shown in below, the procedure end in prev_sample = sample + derivative * dt which is exactly line 8 in Algorithm 2 in the paper. But then the function does not apply second order correction? Second, in this line pred_original_sample = self.precondition_outputs(sample, model_output, sigma_hat), it follows line 7 the $ D _{\theta} (\hat{x} _i ; \hat{t} _{i} ) $ , but model_output is obtained outside of step function, and it is $F _\theta (c _{in} \sigma) x ; c _{noise} (\sigma) $, not the $F _\theta (c _{in} \sigma) \hat{x} ; c _{noise} (\sigma) $.

        eps = noise * s_noise
        sigma_hat = sigma * (gamma + 1)

        if gamma > 0:
            sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5

        # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
        pred_original_sample = self.precondition_outputs(sample, model_output, sigma_hat)

        # 2. Convert to an ODE derivative
        derivative = (sample - pred_original_sample) / sigma_hat

        dt = self.sigmas[self.step_index + 1] - sigma_hat

        prev_sample = sample + derivative * dt

        # Cast sample back to model compatible dtype
        prev_sample = prev_sample.to(model_output.dtype)

        # upon completion increase step index by one
        self._step_index += 1

        if not return_dict:
            return (prev_sample,)

        return EDMEulerSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)

[sayakpaul]

Cc: @patil-suraj