Update pipeline.py

pipeline.py

@@ -172,8 +172,18 @@ class KarrasEDMPipeline(DiffusionPipeline):
         # 5. Denoising loop
         # Implements the "EDM" column in Table 1 of the EDM paper
         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        # BEGIN CHANGE
+        t_init = timesteps[0]
+        scaled_sample_first_step = self.scheduler.scale_model_input(sample / self.scheduler.init_noise_sigma, t_init)
+        if self.scheduler.step_index is None:
+            self.scheduler._init_step_index(t_init)
+        sigma = self.scheduler.sigmas[self.scheduler.step_index]
+        sigma_input = self.scheduler.precondition_noise(sigma)
+        model_output = self.unet(scaled_sample_first_step, sigma_input.squeeze(), return_dict=False)[0]
+        sample = self.scheduler.step(model_output, t_init, sample).prev_sample
+        # END CHANGE
         with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
+            for i, t in enumerate(timesteps[1:]): # changed by giulio
                 # 1. Add noise (if necessary) and precondition the input sample.
                 scaled_sample = self.scheduler.scale_model_input(sample, t)