




import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest


class       __lowercase      (	__lowercase  ):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE   						=				(KDPMaDiscreteScheduler,)
     SCREAMING_SNAKE_CASE   						=				10




     def 		lowerCAmelCase_     (				self       :			List[Any] ,	**UpperCamelCase_       :			Optional[int]				):


         """simple docstring"""



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

         config.update(**__a				)
         return config




     def 		lowerCAmelCase_     (				self       :			Union[str, Any]				):


         """simple docstring"""



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




     def 		lowerCAmelCase_     (				self       :			str				):


         """simple docstring"""



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




     def 		lowerCAmelCase_     (				self       :			int				):


         """simple docstring"""



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




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



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




     def 		lowerCAmelCase_     (				self       :			Union[str, Any]				):


         """simple docstring"""



         __A         =	self.scheduler_classes[0]
         __A         =	self.get_scheduler_config(prediction_type="""v_prediction"""				)
         __A         =	scheduler_class(**__a				)

         scheduler.set_timesteps(self.num_inference_steps				)

         __A         =	self.dummy_model()
         __A         =	self.dummy_sample_deter * scheduler.init_noise_sigma
         __A         =	sample.to(__a				)

         for i, t in enumerate(scheduler.timesteps				):
             __A         =	scheduler.scale_model_input(__a ,	__a				)

             __A         =	model(__a ,	__a				)

             __A         =	scheduler.step(__a ,	__a ,	__a				)
             __A         =	output.prev_sample

         __A         =	torch.sum(torch.abs(__a				)				)
         __A         =	torch.mean(torch.abs(__a				)				)

         if torch_device in ["cpu", "mps"]:
             assert abs(result_sum.item() - 4.69_34e-07				) < 1e-2
             assert abs(result_mean.item() - 6.11_12e-10				) < 1e-3
         else:
             # CUDA
             assert abs(result_sum.item() - 4.6_93_42_86_50_17_09_72e-07				) < 1e-2
             assert abs(result_mean.item() - 0.0002				) < 1e-3




     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         if torch_device == "mps":
             return
         __A         =	self.scheduler_classes[0]
         __A         =	self.get_scheduler_config()
         __A         =	scheduler_class(**__a				)

         scheduler.set_timesteps(self.num_inference_steps				)

         __A         =	self.dummy_model()
         __A         =	self.dummy_sample_deter * scheduler.init_noise_sigma
         __A         =	sample.to(__a				)

         for i, t in enumerate(scheduler.timesteps				):
             __A         =	scheduler.scale_model_input(__a ,	__a				)

             __A         =	model(__a ,	__a				)

             __A         =	scheduler.step(__a ,	__a ,	__a				)
             __A         =	output.prev_sample

         __A         =	torch.sum(torch.abs(__a				)				)
         __A         =	torch.mean(torch.abs(__a				)				)

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





     def 		lowerCAmelCase_     (				self       :			Optional[Any]				):


         """simple docstring"""



         if torch_device == "mps":
             return
         __A         =	self.scheduler_classes[0]
         __A         =	self.get_scheduler_config()
         __A         =	scheduler_class(**__a				)

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

         __A         =	self.dummy_model()
         __A         =	self.dummy_sample_deter.to(__a				) * scheduler.init_noise_sigma

         for t in scheduler.timesteps:
             __A         =	scheduler.scale_model_input(__a ,	__a				)

             __A         =	model(__a ,	__a				)

             __A         =	scheduler.step(__a ,	__a ,	__a				)
             __A         =	output.prev_sample

         __A         =	torch.sum(torch.abs(__a				)				)
         __A         =	torch.mean(torch.abs(__a				)				)

         if str(__a				).startswith("""cpu"""				):
             # The following sum varies between 148 and 156 on mps. Why?
             assert abs(result_sum.item() - 20.4125				) < 1e-2
             assert abs(result_mean.item() - 0.0266				) < 1e-3
         else:
             # CUDA
             assert abs(result_sum.item() - 20.4125				) < 1e-2
             assert abs(result_mean.item() - 0.0266				) < 1e-3












from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding


class       __lowercase      (	lowercase_  ):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE   						=				["image_processor", "tokenizer"]
     SCREAMING_SNAKE_CASE   						=				"AutoImageProcessor"
     SCREAMING_SNAKE_CASE   						=				"AutoTokenizer"




     def __init__(				self       :			str ,	UpperCamelCase_       :			Any ,	UpperCamelCase_       :			str				):


         """simple docstring"""



         super().__init__(UpperCamelCase_ ,	UpperCamelCase_				)
         __A         =	self.image_processor




     def __call__(				self       :			Optional[Any] ,	UpperCamelCase_       :			Optional[Any]=None ,	UpperCamelCase_       :			Any=None ,	UpperCamelCase_       :			List[Any]=None ,	**UpperCamelCase_       :			int				):


         """simple docstring"""



         if text is None and images is None:
             raise ValueError("""You have to specify either text or images. Both cannot be none."""				)

         if text is not None:
             __A         =	self.tokenizer(UpperCamelCase_ ,	return_tensors=UpperCamelCase_ ,	**UpperCamelCase_				)

         if images is not None:
             __A         =	self.image_processor(UpperCamelCase_ ,	return_tensors=UpperCamelCase_ ,	**UpperCamelCase_				)

         if text is not None and images is not None:
             __A         =	image_features.pixel_values
             return encoding
         elif text is not None:
             return encoding
         else:
             return BatchEncoding(data=dict(**UpperCamelCase_				) ,	tensor_type=UpperCamelCase_				)




     def 		lowerCAmelCase_     (				self       :			List[str] ,	*UpperCamelCase_       :			Tuple ,	**UpperCamelCase_       :			List[str]				):


         """simple docstring"""



         return self.tokenizer.batch_decode(*UpperCamelCase_ ,	**UpperCamelCase_				)




     def 		lowerCAmelCase_     (				self       :			Any ,	*UpperCamelCase_       :			int ,	**UpperCamelCase_       :			Tuple				):


         """simple docstring"""



         return self.tokenizer.decode(*UpperCamelCase_ ,	**UpperCamelCase_				)





     @property
     def 		lowerCAmelCase_     (				self       :			Optional[Any]				):


         """simple docstring"""



         return ["input_ids", "attention_mask", "pixel_values"]




