


"""simple docstring"""





from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import center_crop, normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
  import PIL


SCREAMING_SNAKE_CASE				:     Any          =	logging.get_logger(__name__)


class       __lowerCamelCase	(   UpperCamelCase_    ):
   __UpperCamelCase									=   ['pixel_values']







   def __init__(self       ,	lowerCamelCase = True       ,	lowerCamelCase = None       ,	lowerCamelCase = PIL.Image.NEAREST       ,	lowerCamelCase = False       ,	lowerCamelCase = None       ,	lowerCamelCase = 1 / 255       ,	lowerCamelCase = False       ,	lowerCamelCase = True       ,	lowerCamelCase = True       ,	lowerCamelCase = None       ,	lowerCamelCase = None       ,	lowerCamelCase = True       ,	**lowerCamelCase       ,	):




      '''simple docstring'''




      super().__init__(**__UpperCamelCase				)
      _lowerCAmelCase					     =		size if size is not None else {'''height''': 346, '''width''': 346}
      _lowerCAmelCase					     =		get_size_dict(__UpperCamelCase				)
      _lowerCAmelCase					     =		crop_size if crop_size is not None else {'''height''': 289, '''width''': 289}
      _lowerCAmelCase					     =		get_size_dict(__UpperCamelCase       ,	param_name="""crop_size"""				)

      _lowerCAmelCase					     =		do_resize
      _lowerCAmelCase					     =		size
      _lowerCAmelCase					     =		resample
      _lowerCAmelCase					     =		do_center_crop
      _lowerCAmelCase					     =		crop_size
      _lowerCAmelCase					     =		do_rescale
      _lowerCAmelCase					     =		rescale_factor
      _lowerCAmelCase					     =		rescale_offset
      _lowerCAmelCase					     =		do_normalize
      _lowerCAmelCase					     =		image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
      _lowerCAmelCase					     =		image_std if image_std is not None else IMAGENET_STANDARD_STD
      _lowerCAmelCase					     =		include_top







   def 							A__						(self       ,	lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase = PIL.Image.NEAREST       ,	lowerCamelCase = None       ,	**lowerCamelCase       ,	):




      '''simple docstring'''




      _lowerCAmelCase					     =		get_size_dict(__UpperCamelCase				)
      if "height" not in size or "width" not in size:
         raise ValueError(f"""The size dictionary must have keys \'height\' and \'width\'. Got {size.keys()}"""				)
      return resize(
          __UpperCamelCase       ,	size=(size["""height"""], size["""width"""])       ,	resample=__UpperCamelCase       ,	data_format=__UpperCamelCase       ,	**__UpperCamelCase				)







   def 							A__						(self       ,	lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase = None       ,	**lowerCamelCase       ,	):




      '''simple docstring'''




      _lowerCAmelCase					     =		get_size_dict(__UpperCamelCase				)
      if "height" not in size or "width" not in size:
         raise ValueError(f"""The size dictionary must have keys \'height\' and \'width\'. Got {size.keys()}"""				)
      return center_crop(__UpperCamelCase       ,	size=(size["""height"""], size["""width"""])       ,	data_format=__UpperCamelCase       ,	**__UpperCamelCase				)







   def 							A__						(self       ,	lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase = True       ,	lowerCamelCase = None       ,	**lowerCamelCase       ,	):




      '''simple docstring'''




      if offset:
         _lowerCAmelCase					     =		(image - 127.5) * scale
         if data_format is not None:
            _lowerCAmelCase					     =		to_channel_dimension_format(__UpperCamelCase       ,	__UpperCamelCase				)
         _lowerCAmelCase					     =		rescaled_image.astype(np.floataa				)
      else:
         _lowerCAmelCase					     =		rescale(__UpperCamelCase       ,	scale=__UpperCamelCase       ,	data_format=__UpperCamelCase       ,	**__UpperCamelCase				)
      return rescaled_image







   def 							A__						(self       ,	lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase = None       ,	**lowerCamelCase       ,	):




      '''simple docstring'''




      return normalize(__UpperCamelCase       ,	mean=__UpperCamelCase       ,	std=__UpperCamelCase       ,	data_format=__UpperCamelCase       ,	**__UpperCamelCase				)







   def 							A__						(self       ,	lowerCamelCase       ,	lowerCamelCase = None       ,	lowerCamelCase = None       ,	lowerCamelCase=None       ,	lowerCamelCase = None       ,	lowerCamelCase = None       ,	lowerCamelCase = None       ,	lowerCamelCase = None       ,	lowerCamelCase = None       ,	lowerCamelCase = None       ,	lowerCamelCase = None       ,	lowerCamelCase = None       ,	lowerCamelCase = None       ,	lowerCamelCase = None       ,	lowerCamelCase = ChannelDimension.FIRST       ,	**lowerCamelCase       ,	):




      '''simple docstring'''




      _lowerCAmelCase					     =		do_resize if do_resize is not None else self.do_resize
      _lowerCAmelCase					     =		resample if resample is not None else self.resample
      _lowerCAmelCase					     =		do_center_crop if do_center_crop is not None else self.do_center_crop
      _lowerCAmelCase					     =		do_rescale if do_rescale is not None else self.do_rescale
      _lowerCAmelCase					     =		rescale_factor if rescale_factor is not None else self.rescale_factor
      _lowerCAmelCase					     =		rescale_offset if rescale_offset is not None else self.rescale_offset
      _lowerCAmelCase					     =		do_normalize if do_normalize is not None else self.do_normalize
      _lowerCAmelCase					     =		image_mean if image_mean is not None else self.image_mean
      _lowerCAmelCase					     =		image_std if image_std is not None else self.image_std
      _lowerCAmelCase					     =		include_top if include_top is not None else self.include_top

      _lowerCAmelCase					     =		size if size is not None else self.size
      _lowerCAmelCase					     =		get_size_dict(__UpperCamelCase				)
      _lowerCAmelCase					     =		crop_size if crop_size is not None else self.crop_size
      _lowerCAmelCase					     =		get_size_dict(__UpperCamelCase       ,	param_name="""crop_size"""				)

      _lowerCAmelCase					     =		make_list_of_images(__UpperCamelCase				)

      if not valid_images(__UpperCamelCase				):
         raise ValueError(
             """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
             """torch.Tensor, tf.Tensor or jax.ndarray."""				)

      if do_resize and size is None or resample is None:
         raise ValueError("""Size and resample must be specified if do_resize is True."""				)

      if do_center_crop and crop_size is None:
         raise ValueError("""Crop size must be specified if do_center_crop is True."""				)

      if do_rescale and rescale_factor is None:
         raise ValueError("""Rescale factor must be specified if do_rescale is True."""				)

      if do_normalize and (image_mean is None or image_std is None):
         raise ValueError("""Image mean and std must be specified if do_normalize is True."""				)

      # All transformations expect numpy arrays.
      _lowerCAmelCase					     =		[to_numpy_array(__UpperCamelCase				) for image in images]

      if do_resize:
         _lowerCAmelCase					     =		[self.resize(image=__UpperCamelCase       ,	size=__UpperCamelCase       ,	resample=__UpperCamelCase				) for image in images]

      if do_center_crop:
         _lowerCAmelCase					     =		[self.center_crop(image=__UpperCamelCase       ,	size=__UpperCamelCase				) for image in images]

      if do_rescale:
         _lowerCAmelCase					     =		[self.rescale(image=__UpperCamelCase       ,	scale=__UpperCamelCase       ,	offset=__UpperCamelCase				) for image in images]

      if do_normalize:
         _lowerCAmelCase					     =		[self.normalize(image=__UpperCamelCase       ,	mean=__UpperCamelCase       ,	std=__UpperCamelCase				) for image in images]

      if include_top:
         _lowerCAmelCase					     =		[self.normalize(image=__UpperCamelCase       ,	mean=[0, 0, 0]       ,	std=__UpperCamelCase				) for image in images]

      _lowerCAmelCase					     =		[to_channel_dimension_format(__UpperCamelCase       ,	__UpperCamelCase				) for image in images]

      _lowerCAmelCase					     =		{'''pixel_values''': images}
      return BatchFeature(data=__UpperCamelCase       ,	tensor_type=__UpperCamelCase				)




"""simple docstring"""





import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest


class       __lowerCamelCase	(   __lowercase    ):
   __UpperCamelCase									=   (HeunDiscreteScheduler,)
   __UpperCamelCase									=   10







   def 							A__						(self       ,	**lowerCamelCase				):




      '''simple docstring'''




      _lowerCAmelCase					     =		{
          """num_train_timesteps""": 1_100,
          """beta_start""": 0.0001,
          """beta_end""": 0.02,
          """beta_schedule""": """linear""",
      }

      config.update(**lowerCamelCase				)
      return config







   def 							A__						(self				):




      '''simple docstring'''




      for timesteps in [10, 50, 100, 1_000]:
         self.check_over_configs(num_train_timesteps=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




      for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001]       ,	[0.0002, 0.002, 0.02]				):
         self.check_over_configs(beta_start=lowerCamelCase       ,	beta_end=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




      for schedule in ["linear", "scaled_linear", "exp"]:
         self.check_over_configs(beta_schedule=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




      for clip_sample_range in [1.0, 2.0, 3.0]:
         self.check_over_configs(clip_sample_range=lowerCamelCase       ,	clip_sample=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




      for prediction_type in ["epsilon", "v_prediction", "sample"]:
         self.check_over_configs(prediction_type=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.scheduler_classes[0]
      _lowerCAmelCase					     =		self.get_scheduler_config()
      _lowerCAmelCase					     =		scheduler_class(**lowerCamelCase				)

      scheduler.set_timesteps(self.num_inference_steps				)

      _lowerCAmelCase					     =		self.dummy_model()
      _lowerCAmelCase					     =		self.dummy_sample_deter * scheduler.init_noise_sigma
      _lowerCAmelCase					     =		sample.to(lowerCamelCase				)

      for i, t in enumerate(scheduler.timesteps				):
         _lowerCAmelCase					     =		scheduler.scale_model_input(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		model(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		scheduler.step(lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		torch.mean(torch.abs(lowerCamelCase				)				)

      if torch_device in ["cpu", "mps"]:
         assert abs(result_sum.item() - 0.1233				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3
      else:
         # CUDA
         assert abs(result_sum.item() - 0.1233				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.scheduler_classes[0]
      _lowerCAmelCase					     =		self.get_scheduler_config(prediction_type="""v_prediction"""				)
      _lowerCAmelCase					     =		scheduler_class(**lowerCamelCase				)

      scheduler.set_timesteps(self.num_inference_steps				)

      _lowerCAmelCase					     =		self.dummy_model()
      _lowerCAmelCase					     =		self.dummy_sample_deter * scheduler.init_noise_sigma
      _lowerCAmelCase					     =		sample.to(lowerCamelCase				)

      for i, t in enumerate(scheduler.timesteps				):
         _lowerCAmelCase					     =		scheduler.scale_model_input(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		model(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		scheduler.step(lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		torch.mean(torch.abs(lowerCamelCase				)				)

      if torch_device in ["cpu", "mps"]:
         assert abs(result_sum.item() - 4.6934e-07				) < 1e-2
         assert abs(result_mean.item() - 6.1112e-10				) < 1e-3
      else:
         # CUDA
         assert abs(result_sum.item() - 4.693428650170972e-07				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.scheduler_classes[0]
      _lowerCAmelCase					     =		self.get_scheduler_config()
      _lowerCAmelCase					     =		scheduler_class(**lowerCamelCase				)

      scheduler.set_timesteps(self.num_inference_steps       ,	device=lowerCamelCase				)

      _lowerCAmelCase					     =		self.dummy_model()
      _lowerCAmelCase					     =		self.dummy_sample_deter.to(lowerCamelCase				) * scheduler.init_noise_sigma

      for t in scheduler.timesteps:
         _lowerCAmelCase					     =		scheduler.scale_model_input(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		model(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		scheduler.step(lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		torch.mean(torch.abs(lowerCamelCase				)				)

      if str(lowerCamelCase				).startswith("""cpu"""				):
         # The following sum varies between 148 and 156 on mps. Why?
         assert abs(result_sum.item() - 0.1233				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3
      elif str(lowerCamelCase				).startswith("""mps"""				):
         # Larger tolerance on mps
         assert abs(result_mean.item() - 0.0002				) < 1e-2
      else:
         # CUDA
         assert abs(result_sum.item() - 0.1233				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.scheduler_classes[0]
      _lowerCAmelCase					     =		self.get_scheduler_config()
      _lowerCAmelCase					     =		scheduler_class(**lowerCamelCase       ,	use_karras_sigmas=lowerCamelCase				)

      scheduler.set_timesteps(self.num_inference_steps       ,	device=lowerCamelCase				)

      _lowerCAmelCase					     =		self.dummy_model()
      _lowerCAmelCase					     =		self.dummy_sample_deter.to(lowerCamelCase				) * scheduler.init_noise_sigma
      _lowerCAmelCase					     =		sample.to(lowerCamelCase				)

      for t in scheduler.timesteps:
         _lowerCAmelCase					     =		scheduler.scale_model_input(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		model(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		scheduler.step(lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		torch.mean(torch.abs(lowerCamelCase				)				)

      assert abs(result_sum.item() - 0.0_0015				) < 1e-2
      assert abs(result_mean.item() - 1.9869554535034695e-07				) < 1e-2