


'''simple docstring'''


from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
						from PIL import Image

						from ..image_utils import load_image

if is_tf_available():
						import tensorflow as tf

						from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
						from ..tf_utils import stable_softmax

if is_torch_available():
						from ..models.auto.modeling_auto import MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING

lowerCamelCase__					=   logging.get_logger(__name__)






@add_end_docstrings(_UpperCAmelCase		)
class 		lowerCAmelCase__	(	_UpperCAmelCase		):




			def __init__(						self     :   Optional[int]		,     *lowerCamelCase__     :   Tuple		,     **lowerCamelCase__     :   Tuple       )      ->str:



					'''simple docstring'''







					super().__init__(*lowerCAmelCase__		,     **lowerCAmelCase__       )
					requires_backends(self		,     "vision"       )
					self.check_model_type(
					    TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
					    if self.framework == "tf"
					    else MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING       )




			def 							lowerCAmelCase__			(						self     :   Any		,     lowerCamelCase__     :   int=None       )      ->List[Any]:



					'''simple docstring'''







					_UpperCAmelCase						:				List[str]	    =					{}
					if top_k is not None:
							_UpperCAmelCase						:				List[Any]	    =					top_k
					return {}, {}, postprocess_params




			def __call__(						self     :   Optional[Any]		,     lowerCamelCase__     :   Dict		,     **lowerCamelCase__     :   Optional[int]       )      ->Tuple:



					'''simple docstring'''







					return super().__call__(lowerCAmelCase__		,     **lowerCAmelCase__       )




			def 							lowerCAmelCase__			(						self     :   List[str]		,     lowerCamelCase__     :   Optional[Any]       )      ->Union[str, Any]:



					'''simple docstring'''







					_UpperCAmelCase						:				int	    =					load_image(lowerCAmelCase__       )
					_UpperCAmelCase						:				str	    =					self.image_processor(images=lowerCAmelCase__		,     return_tensors=self.framework       )
					return model_inputs




			def 							lowerCAmelCase__			(						self     :   Any		,     lowerCamelCase__     :   Optional[Any]       )      ->List[Any]:



					'''simple docstring'''







					_UpperCAmelCase						:				List[Any]	    =					self.model(**lowerCAmelCase__       )
					return model_outputs






			def 							lowerCAmelCase__			(						self     :   int		,     lowerCamelCase__     :   Union[str, Any]		,     lowerCamelCase__     :   str=5       )      ->List[Any]:



					'''simple docstring'''







					if top_k > self.model.config.num_labels:
							_UpperCAmelCase						:				Union[str, Any]	    =					self.model.config.num_labels

					if self.framework == "pt":
							_UpperCAmelCase						:				Dict	    =					model_outputs.logits.softmax(-1       )[0]
							_UpperCAmelCase    ,			_UpperCAmelCase						:				List[Any]	    =					probs.topk(lowerCAmelCase__       )
					elif self.framework == "tf":
							_UpperCAmelCase						:				Union[str, Any]	    =					stable_softmax(model_outputs.logits		,     axis=-1       )[0]
							_UpperCAmelCase						:				Dict	    =					tf.math.top_k(lowerCAmelCase__		,     k=lowerCAmelCase__       )
							_UpperCAmelCase    ,			_UpperCAmelCase						:				Optional[Any]	    =					topk.values.numpy(), topk.indices.numpy()
					else:
							raise ValueError(F"""Unsupported framework: {self.framework}"""       )

					_UpperCAmelCase						:				Optional[Any]	    =					scores.tolist()
					_UpperCAmelCase						:				Dict	    =					ids.tolist()
					return [{"score": score, "label": self.model.config.idalabel[_id]} for score, _id in zip(lowerCAmelCase__		,     lowerCAmelCase__       )]




import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest




class    _UpperCamelCase      (						_UpperCAmelCase ):


			"""simple docstring"""






			__a  : Optional[Any]			 =						(KDPMaDiscreteScheduler,)
			__a  : Dict			 =						10





			def        _SCREAMING_SNAKE_CASE			(   self      ,   **lowerCAmelCase__						)     ->					int:

							'''simple docstring'''


							__lowercase 			=							{
							    '''num_train_timesteps''': 11_00,
							    '''beta_start''': 0.0001,
							    '''beta_end''': 0.02,
							    '''beta_schedule''': '''linear''',
							}

							config.update(**lowerCAmelCase__						)
							return config





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					List[str]:

							'''simple docstring'''


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





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Any:

							'''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        _SCREAMING_SNAKE_CASE			(   self						)     ->					Union[str, Any]:

							'''simple docstring'''


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





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Union[str, Any]:

							'''simple docstring'''


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





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Tuple:

							'''simple docstring'''


							__lowercase 			=							self.scheduler_classes[0]
							__lowercase 			=							self.get_scheduler_config(prediction_type='''v_prediction'''						)
							__lowercase 			=							scheduler_class(**lowerCAmelCase__						)

							scheduler.set_timesteps(self.num_inference_steps						)

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

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

											__lowercase 			=							model(lowerCAmelCase__      ,   lowerCAmelCase__						)

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

							__lowercase 			=							torch.sum(torch.abs(lowerCAmelCase__						)						)
							__lowercase 			=							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.693_4286_5017_0972E-07						) < 1E-2
											assert abs(result_mean.item() - 0.0002						) < 1E-3





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					str:

							'''simple docstring'''


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

							scheduler.set_timesteps(self.num_inference_steps						)

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

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

											__lowercase 			=							model(lowerCAmelCase__      ,   lowerCAmelCase__						)

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

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

							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        _SCREAMING_SNAKE_CASE			(   self						)     ->					Dict:

							'''simple docstring'''


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

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

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

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

											__lowercase 			=							model(lowerCAmelCase__      ,   lowerCAmelCase__						)

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

							__lowercase 			=							torch.sum(torch.abs(lowerCAmelCase__						)						)
							__lowercase 			=							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() - 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