fix artifacts from trying to use f16

src/editor.py

@@ -44,7 +44,7 @@ class ImageEditorDemo:
         img_size = (512,512)
         VQAE_SCALE = 8
         latents_size = (1, 4, img_size[0] // VQAE_SCALE, img_size[1] // VQAE_SCALE)
-        noise = [randn_tensor(latents_size, dtype=torch.float32, device=torch.device(device), generator=g_cpu) for i
+        noise = [randn_tensor(latents_size, dtype=torch.float16, device=torch.device(device), generator=g_cpu) for i
                  in range(cfg.num_inversion_steps)]
         pipe_inversion.scheduler.set_noise_list(noise)
         pipe_inversion.scheduler_inference.set_noise_list(noise)

src/euler_scheduler.py

@@ -419,172 +419,4 @@ class MyEulerAncestralDiscreteScheduler(EulerAncestralDiscreteScheduler):
             sigma = sigma.unsqueeze(-1)
 
         noisy_samples = original_samples + noise * sigma
-        return noisy_samples
-
-    # def update_noise_for_friendly_inversion(
-    #     self,
-    #     model_output: torch.FloatTensor,
-    #     timestep: Union[float, torch.FloatTensor],
-    #     z_t: torch.FloatTensor,
-    #     z_tp1: torch.FloatTensor,
-    #     return_dict: bool = True,
-    # ) -> Union[EulerAncestralDiscreteSchedulerOutput, Tuple]:
-    #     if (
-    #         isinstance(timestep, int)
-    #         or isinstance(timestep, torch.IntTensor)
-    #         or isinstance(timestep, torch.LongTensor)
-    #     ):
-    #         raise ValueError(
-    #             (
-    #                 "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
-    #                 " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
-    #                 " one of the `scheduler.timesteps` as a timestep."
-    #             ),
-    #         )
-
-    #     if not self.is_scale_input_called:
-    #         logger.warning(
-    #             "The `scale_model_input` function should be called before `step` to ensure correct denoising. "
-    #             "See `StableDiffusionPipeline` for a usage example."
-    #         )
-
-    #     self._init_step_index(timestep.view((1)))
-
-    #     sigma = self.sigmas[self.step_index]
-
-    #     sigma_from = self.sigmas[self.step_index]
-    #     sigma_to = self.sigmas[self.step_index+1]
-    #     # sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5
-    #     sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2).abs() / sigma_from**2) ** 0.5
-    #     # sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5
-    #     sigma_down = sigma_to**2 / sigma_from
-
-    #     # 2. Conv = (sample - pred_original_sample) / sigma
-    #     derivative = model_output
-
-    #     dt = sigma_down - sigma
-    #     # dt = sigma_down - sigma_from
-
-    #     prev_sample = z_t - derivative * dt
-
-    #     if sigma_up > 0:
-    #         self.noise_list[self.step_index] = (prev_sample - z_tp1) / sigma_up
-
-    #     prev_sample = prev_sample - self.noise_list[self.step_index] * sigma_up
-
-
-    #     if not return_dict:
-    #         return (prev_sample,)
-
-    #     return EulerAncestralDiscreteSchedulerOutput(
-    #         prev_sample=prev_sample, pred_original_sample=None
-    #     )
-
-
-    # def step_friendly_inversion(
-    #     self,
-    #     model_output: torch.FloatTensor,
-    #     timestep: Union[float, torch.FloatTensor],
-    #     sample: torch.FloatTensor,
-    #     generator: Optional[torch.Generator] = None,
-    #     return_dict: bool = True,
-    #     expected_next_sample: torch.FloatTensor = None,
-    # ) -> Union[EulerAncestralDiscreteSchedulerOutput, Tuple]:
-    #     """
-    #     Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
-    #     process from the learned model outputs (most often the predicted noise).
-
-    #     Args:
-    #         model_output (`torch.FloatTensor`):
-    #             The direct output from learned diffusion model.
-    #         timestep (`float`):
-    #             The current discrete timestep in the diffusion chain.
-    #         sample (`torch.FloatTensor`):
-    #             A current instance of a sample created by the diffusion process.
-    #         generator (`torch.Generator`, *optional*):
-    #             A random number generator.
-    #         return_dict (`bool`):
-    #             Whether or not to return a
-    #             [`~schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteSchedulerOutput`] or tuple.
-
-    #     Returns:
-    #         [`~schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteSchedulerOutput`] or `tuple`:
-    #             If return_dict is `True`,
-    #             [`~schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteSchedulerOutput`] is returned,
-    #             otherwise a tuple is returned where the first element is the sample tensor.
-
-    #     """
-
-    #     if (
-    #         isinstance(timestep, int)
-    #         or isinstance(timestep, torch.IntTensor)
-    #         or isinstance(timestep, torch.LongTensor)
-    #     ):
-    #         raise ValueError(
-    #             (
-    #                 "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
-    #                 " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
-    #                 " one of the `scheduler.timesteps` as a timestep."
-    #             ),
-    #         )
-
-    #     if not self.is_scale_input_called:
-    #         logger.warning(
-    #             "The `scale_model_input` function should be called before `step` to ensure correct denoising. "
-    #             "See `StableDiffusionPipeline` for a usage example."
-    #         )
-
-    #     self._init_step_index(timestep.view((1)))
-
-    #     sigma = self.sigmas[self.step_index]
-
-    #     # Upcast to avoid precision issues when computing prev_sample
-        sample = sample.to(torch.float32)
-
-    #     # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
-    #     if self.config.prediction_type == "epsilon":
-    #         pred_original_sample = sample - sigma * model_output
-    #     elif self.config.prediction_type == "v_prediction":
-    #         # * c_out + input * c_skip
-    #         pred_original_sample = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1))
-    #     elif self.config.prediction_type == "sample":
-    #         raise NotImplementedError("prediction_type not implemented yet: sample")
-    #     else:
-    #         raise ValueError(
-    #             f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`"
-    #         )
-
-    #     sigma_from = self.sigmas[self.step_index]
-    #     sigma_to = self.sigmas[self.step_index + 1]
-    #     sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5
-    #     sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5
-
-    #     # 2. Convert to an ODE derivative
-    #     # derivative = (sample - pred_original_sample) / sigma
-    #     derivative = model_output
-
-    #     dt = sigma_down - sigma
-
-    #     prev_sample = sample + derivative * dt
-
-    #     device = model_output.device
-    #     # noise = randn_tensor(model_output.shape, dtype=model_output.dtype, device=device, generator=generator)
-    #     # prev_sample = prev_sample + noise * sigma_up
-
-    #     if sigma_up > 0:
-    #         self.noise_list[self.step_index] = (expected_next_sample - prev_sample) / sigma_up
-
-    #     prev_sample = prev_sample + self.noise_list[self.step_index] * sigma_up
-
-    #     # 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 EulerAncestralDiscreteSchedulerOutput(
-    #         prev_sample=prev_sample, pred_original_sample=pred_original_sample
-    #     )
+        return noisy_samples