






import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest



class  __UpperCAmelCase						(SCREAMING_SNAKE_CASE__       ):
			__snake_case							: Optional[Any]							=				(HeunDiscreteScheduler,)
			__snake_case							: Dict							=				10







			def 				UpperCamelCase (      self:    List[Any] ,      **UpperCAmelCase_:    List[Any]	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    {
									    "num_train_timesteps": 1_100,
									    "beta_start": 0.00_01,
									    "beta_end": 0.02,
									    "beta_schedule": "linear",
									}

									config.update(**lowerCAmelCase__	)
									return config







			def 				UpperCamelCase (      self:    List[Any]	):







									'''simple docstring'''




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







			def 				UpperCamelCase (      self:    str	):







									'''simple docstring'''




									for beta_start, beta_end in zip([0.0_00_01, 0.00_01, 0.0_01] ,      [0.00_02, 0.0_02, 0.02]	):
															self.check_over_configs(beta_start=lowerCAmelCase__ ,      beta_end=lowerCAmelCase__	)







			def 				UpperCamelCase (      self:    Dict	):







									'''simple docstring'''




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







			def 				UpperCamelCase (      self:    Any	):







									'''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 				UpperCamelCase (      self:    int	):







									'''simple docstring'''




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







			def 				UpperCamelCase (      self:    Any	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    self.scheduler_classes[0]
									_SCREAMING_SNAKE_CASE       =    self.get_scheduler_config()
									_SCREAMING_SNAKE_CASE       =    scheduler_class(**lowerCAmelCase__	)

									scheduler.set_timesteps(self.num_inference_steps	)

									_SCREAMING_SNAKE_CASE       =    self.dummy_model()
									_SCREAMING_SNAKE_CASE       =    self.dummy_sample_deter * scheduler.init_noise_sigma
									_SCREAMING_SNAKE_CASE       =    sample.to(lowerCAmelCase__	)

									for i, t in enumerate(scheduler.timesteps	):
															_SCREAMING_SNAKE_CASE       =    scheduler.scale_model_input(lowerCAmelCase__ ,      lowerCAmelCase__	)

															_SCREAMING_SNAKE_CASE       =    model(lowerCAmelCase__ ,      lowerCAmelCase__	)

															_SCREAMING_SNAKE_CASE       =    scheduler.step(lowerCAmelCase__ ,      lowerCAmelCase__ ,      lowerCAmelCase__	)
															_SCREAMING_SNAKE_CASE       =    output.prev_sample

									_SCREAMING_SNAKE_CASE       =    torch.sum(torch.abs(lowerCAmelCase__	)	)
									_SCREAMING_SNAKE_CASE       =    torch.mean(torch.abs(lowerCAmelCase__	)	)

									if torch_device in ["cpu", "mps"]:
															assert abs(result_sum.item() - 0.12_33	) < 1E-2
															assert abs(result_mean.item() - 0.00_02	) < 1E-3
									else:
															# CUDA
															assert abs(result_sum.item() - 0.12_33	) < 1E-2
															assert abs(result_mean.item() - 0.00_02	) < 1E-3







			def 				UpperCamelCase (      self:    Any	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    self.scheduler_classes[0]
									_SCREAMING_SNAKE_CASE       =    self.get_scheduler_config(prediction_type="""v_prediction"""	)
									_SCREAMING_SNAKE_CASE       =    scheduler_class(**lowerCAmelCase__	)

									scheduler.set_timesteps(self.num_inference_steps	)

									_SCREAMING_SNAKE_CASE       =    self.dummy_model()
									_SCREAMING_SNAKE_CASE       =    self.dummy_sample_deter * scheduler.init_noise_sigma
									_SCREAMING_SNAKE_CASE       =    sample.to(lowerCAmelCase__	)

									for i, t in enumerate(scheduler.timesteps	):
															_SCREAMING_SNAKE_CASE       =    scheduler.scale_model_input(lowerCAmelCase__ ,      lowerCAmelCase__	)

															_SCREAMING_SNAKE_CASE       =    model(lowerCAmelCase__ ,      lowerCAmelCase__	)

															_SCREAMING_SNAKE_CASE       =    scheduler.step(lowerCAmelCase__ ,      lowerCAmelCase__ ,      lowerCAmelCase__	)
															_SCREAMING_SNAKE_CASE       =    output.prev_sample

									_SCREAMING_SNAKE_CASE       =    torch.sum(torch.abs(lowerCAmelCase__	)	)
									_SCREAMING_SNAKE_CASE       =    torch.mean(torch.abs(lowerCAmelCase__	)	)

									if torch_device in ["cpu", "mps"]:
															assert abs(result_sum.item() - 4.6_934E-07	) < 1E-2
															assert abs(result_mean.item() - 6.1_112E-10	) < 1E-3
									else:
															# CUDA
															assert abs(result_sum.item() - 4.693_428_650_170_972E-07	) < 1E-2
															assert abs(result_mean.item() - 0.00_02	) < 1E-3







			def 				UpperCamelCase (      self:    Tuple	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    self.scheduler_classes[0]
									_SCREAMING_SNAKE_CASE       =    self.get_scheduler_config()
									_SCREAMING_SNAKE_CASE       =    scheduler_class(**lowerCAmelCase__	)

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

									_SCREAMING_SNAKE_CASE       =    self.dummy_model()
									_SCREAMING_SNAKE_CASE       =    self.dummy_sample_deter.to(lowerCAmelCase__	) * scheduler.init_noise_sigma

									for t in scheduler.timesteps:
															_SCREAMING_SNAKE_CASE       =    scheduler.scale_model_input(lowerCAmelCase__ ,      lowerCAmelCase__	)

															_SCREAMING_SNAKE_CASE       =    model(lowerCAmelCase__ ,      lowerCAmelCase__	)

															_SCREAMING_SNAKE_CASE       =    scheduler.step(lowerCAmelCase__ ,      lowerCAmelCase__ ,      lowerCAmelCase__	)
															_SCREAMING_SNAKE_CASE       =    output.prev_sample

									_SCREAMING_SNAKE_CASE       =    torch.sum(torch.abs(lowerCAmelCase__	)	)
									_SCREAMING_SNAKE_CASE       =    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.12_33	) < 1E-2
															assert abs(result_mean.item() - 0.00_02	) < 1E-3
									elif str(lowerCAmelCase__	).startswith("""mps"""	):
															# Larger tolerance on mps
															assert abs(result_mean.item() - 0.00_02	) < 1E-2
									else:
															# CUDA
															assert abs(result_sum.item() - 0.12_33	) < 1E-2
															assert abs(result_mean.item() - 0.00_02	) < 1E-3







			def 				UpperCamelCase (      self:    List[str]	):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE       =    self.scheduler_classes[0]
									_SCREAMING_SNAKE_CASE       =    self.get_scheduler_config()
									_SCREAMING_SNAKE_CASE       =    scheduler_class(**lowerCAmelCase__ ,      use_karras_sigmas=lowerCAmelCase__	)

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

									_SCREAMING_SNAKE_CASE       =    self.dummy_model()
									_SCREAMING_SNAKE_CASE       =    self.dummy_sample_deter.to(lowerCAmelCase__	) * scheduler.init_noise_sigma
									_SCREAMING_SNAKE_CASE       =    sample.to(lowerCAmelCase__	)

									for t in scheduler.timesteps:
															_SCREAMING_SNAKE_CASE       =    scheduler.scale_model_input(lowerCAmelCase__ ,      lowerCAmelCase__	)

															_SCREAMING_SNAKE_CASE       =    model(lowerCAmelCase__ ,      lowerCAmelCase__	)

															_SCREAMING_SNAKE_CASE       =    scheduler.step(lowerCAmelCase__ ,      lowerCAmelCase__ ,      lowerCAmelCase__	)
															_SCREAMING_SNAKE_CASE       =    output.prev_sample

									_SCREAMING_SNAKE_CASE       =    torch.sum(torch.abs(lowerCAmelCase__	)	)
									_SCREAMING_SNAKE_CASE       =    torch.mean(torch.abs(lowerCAmelCase__	)	)

									assert abs(result_sum.item() - 0.0_00_15	) < 1E-2
									assert abs(result_mean.item() - 1.9_869_554_535_034_695E-07	) < 1E-2







"""simple docstring"""
import json
import os
import tempfile

from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available


if is_torch_available():
    import numpy as np
    import torch

if is_vision_available():
    from PIL import Image






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		Any  ,	__lowerCAmelCase			:		Dict=False  ,	__lowerCAmelCase			:		Optional[Any]=False  ,	__lowerCAmelCase			:		int=False						)    ->       List[str]:
 assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

 _UpperCamelCase      :	Optional[int]					      =  []
 for i in range(image_processor_tester.batch_size						):
  if equal_resolution:
   _UpperCamelCase      :	List[str]					      =  image_processor_tester.max_resolution
  else:
   # To avoid getting image width/height 0
   _UpperCamelCase      :	int					      =  image_processor_tester.min_resolution
   if getattr(__lowerCAmelCase  ,	"size_divisor"  ,	__lowerCAmelCase						):
    # If `size_divisor` is defined, the image needs to have width/size >= `size_divisor`
    _UpperCamelCase      :	Optional[int]					      =  max(image_processor_tester.size_divisor  ,	__lowerCAmelCase						)
   _UpperCamelCase      ,	_UpperCamelCase      :	Optional[Any]					      =  np.random.choice(np.arange(__lowerCAmelCase  ,	image_processor_tester.max_resolution						)  ,	2						)
  image_inputs.append(
      np.random.randint(255  ,	size=(image_processor_tester.num_channels, width, height)  ,	dtype=np.uinta						)						)

 if not numpify and not torchify:
  # PIL expects the channel dimension as last dimension
  _UpperCamelCase      :	Optional[Any]					      =  [Image.fromarray(np.moveaxis(__lowerCAmelCase  ,	0  ,	-1						)						) for image in image_inputs]

 if torchify:
  _UpperCamelCase      :	str					      =  [torch.from_numpy(__lowerCAmelCase						) for image in image_inputs]

 return image_inputs






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		Optional[Any]  ,	__lowerCAmelCase			:		Tuple=10  ,	__lowerCAmelCase			:		str=10  ,	__lowerCAmelCase			:		Any=False  ,	__lowerCAmelCase			:		List[Any]=False						)    ->       Optional[int]:
 _UpperCamelCase      :	Optional[int]					      =  []
 for i in range(image_processor_tester.num_frames						):
  video.append(np.random.randint(255  ,	size=(image_processor_tester.num_channels, width, height)  ,	dtype=np.uinta						)						)

 if not numpify and not torchify:
  # PIL expects the channel dimension as last dimension
  _UpperCamelCase      :	Optional[int]					      =  [Image.fromarray(np.moveaxis(__lowerCAmelCase  ,	0  ,	-1						)						) for frame in video]

 if torchify:
  _UpperCamelCase      :	int					      =  [torch.from_numpy(__lowerCAmelCase						) for frame in video]

 return video






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		List[Any]  ,	__lowerCAmelCase			:		Union[str, Any]=False  ,	__lowerCAmelCase			:		Any=False  ,	__lowerCAmelCase			:		Dict=False						)    ->       Any:
 assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

 _UpperCamelCase      :	Optional[Any]					      =  []
 for i in range(image_processor_tester.batch_size						):
  if equal_resolution:
   _UpperCamelCase      :	List[Any]					      =  image_processor_tester.max_resolution
  else:
   _UpperCamelCase      ,	_UpperCamelCase      :	Optional[int]					      =  np.random.choice(
       np.arange(image_processor_tester.min_resolution  ,	image_processor_tester.max_resolution						)  ,	2						)
   _UpperCamelCase      :	List[Any]					      =  prepare_video(
       image_processor_tester=__lowerCAmelCase  ,	width=__lowerCAmelCase  ,	height=__lowerCAmelCase  ,	numpify=__lowerCAmelCase  ,	torchify=__lowerCAmelCase  ,	)
  video_inputs.append(__lowerCAmelCase						)

 return video_inputs







class  _SCREAMING_SNAKE_CASE       :



 '''simple docstring'''





 __UpperCAmelCase		   =							None






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	str					      =  self.image_processing_class(**self.image_processor_dict   )
  _UpperCamelCase      :	Optional[int]					      =  json.loads(image_processor.to_json_string()   )
  for key, value in self.image_processor_dict.items():
   self.assertEqual(obj[key]					,     lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Union[str, Any]					      =  self.image_processing_class(**self.image_processor_dict   )

  with tempfile.TemporaryDirectory() as tmpdirname:
   _UpperCamelCase      :	Tuple					      =  os.path.join(lowerCAmelCase__					,     "image_processor.json"   )
   image_processor_first.to_json_file(lowerCAmelCase__   )
   _UpperCamelCase      :	Tuple					      =  self.image_processing_class.from_json_file(lowerCAmelCase__   )

  self.assertEqual(image_processor_second.to_dict()					,     image_processor_first.to_dict()   )






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Dict					      =  self.image_processing_class(**self.image_processor_dict   )

  with tempfile.TemporaryDirectory() as tmpdirname:
   _UpperCamelCase      :	Optional[int]					      =  image_processor_first.save_pretrained(lowerCAmelCase__   )[0]
   check_json_file_has_correct_format(lowerCAmelCase__   )
   _UpperCamelCase      :	int					      =  self.image_processing_class.from_pretrained(lowerCAmelCase__   )

  self.assertEqual(image_processor_second.to_dict()					,     image_processor_first.to_dict()   )






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Tuple					      =  self.image_processing_class()
  self.assertIsNotNone(lowerCAmelCase__   )






 @require_torch
 @require_vision
 def lowercase_      (self   ):

  '''simple docstring'''







  if self.test_cast_dtype is not None:
   # Initialize image_processor
   _UpperCamelCase      :	Optional[int]					      =  self.image_processing_class(**self.image_processor_dict   )

   # create random PyTorch tensors
   _UpperCamelCase      :	str					      =  prepare_image_inputs(self.image_processor_tester					,     equal_resolution=lowerCAmelCase__					,     torchify=lowerCAmelCase__   )

   _UpperCamelCase      :	List[Any]					      =  image_processor(lowerCAmelCase__					,     return_tensors="pt"   )
   # for layoutLM compatiblity
   self.assertEqual(encoding.pixel_values.device					,     torch.device("cpu"   )   )
   self.assertEqual(encoding.pixel_values.dtype					,     torch.floataa   )

   _UpperCamelCase      :	List[Any]					      =  image_processor(lowerCAmelCase__					,     return_tensors="pt"   ).to(torch.floataa   )
   self.assertEqual(encoding.pixel_values.device					,     torch.device("cpu"   )   )
   self.assertEqual(encoding.pixel_values.dtype					,     torch.floataa   )

   _UpperCamelCase      :	str					      =  image_processor(lowerCAmelCase__					,     return_tensors="pt"   ).to("cpu"					,     torch.bfloataa   )
   self.assertEqual(encoding.pixel_values.device					,     torch.device("cpu"   )   )
   self.assertEqual(encoding.pixel_values.dtype					,     torch.bfloataa   )

   with self.assertRaises(lowerCAmelCase__   ):
    _UpperCamelCase      :	Optional[int]					      =  image_processor(lowerCAmelCase__					,     return_tensors="pt"   ).to(torch.bfloataa					,     "cpu"   )

   # Try with text + image feature
   _UpperCamelCase      :	Union[str, Any]					      =  image_processor(lowerCAmelCase__					,     return_tensors="pt"   )
   encoding.update({"input_ids": torch.LongTensor([[1, 2, 3], [4, 5, 6]]   )}   )
   _UpperCamelCase      :	Optional[int]					      =  encoding.to(torch.floataa   )

   self.assertEqual(encoding.pixel_values.device					,     torch.device("cpu"   )   )
   self.assertEqual(encoding.pixel_values.dtype					,     torch.floataa   )
   self.assertEqual(encoding.input_ids.dtype					,     torch.long   )



