


'''simple docstring'''




import tempfile

import torch

from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    UniPCMultistepScheduler,
)

from .test_schedulers import SchedulerCommonTest






class 							__snake_case    (				__magic_name__       ):
	__lowerCAmelCase						=     (DPMSolverMultistepScheduler,)
	__lowerCAmelCase						=     (('''num_inference_steps''', 25),)





	def      _snake_case      (		self			,					**UpperCamelCase_		)			->		Optional[Any]:
		snake_case__  					=      {
		    'num_train_timesteps': 1000,
		    'beta_start': 0.0_0_0_1,
		    'beta_end': 0.0_2,
		    'beta_schedule': 'linear',
		    'solver_order': 2,
		    'prediction_type': 'epsilon',
		    'thresholding': False,
		    'sample_max_value': 1.0,
		    'algorithm_type': 'dpmsolver++',
		    'solver_type': 'midpoint',
		    'lower_order_final': False,
		    'lambda_min_clipped': -float('inf'		),
		    'variance_type': None,
		}

		config.update(**UpperCamelCase_		)
		return config





	def      _snake_case      (		self			,					UpperCamelCase_=0			,					**UpperCamelCase_		)			->		Any:
		snake_case__  					=      dict(self.forward_default_kwargs		)
		snake_case__  					=      kwargs.pop('num_inference_steps'			,					UpperCamelCase_		)
		snake_case__  					=      self.dummy_sample
		snake_case__  					=      0.1 * sample
		snake_case__  					=      [residual + 0.2, residual + 0.1_5, residual + 0.1_0]

		for scheduler_class in self.scheduler_classes:
			snake_case__  					=      self.get_scheduler_config(**UpperCamelCase_		)
			snake_case__  					=      scheduler_class(**UpperCamelCase_		)
			scheduler.set_timesteps(UpperCamelCase_		)
			# copy over dummy past residuals
			snake_case__  					=      dummy_past_residuals[: scheduler.config.solver_order]

			with tempfile.TemporaryDirectory() as tmpdirname:
				scheduler.save_config(UpperCamelCase_		)
				snake_case__  					=      scheduler_class.from_pretrained(UpperCamelCase_		)
				new_scheduler.set_timesteps(UpperCamelCase_		)
				# copy over dummy past residuals
				snake_case__  					=      dummy_past_residuals[: new_scheduler.config.solver_order]

			snake_case__					,     snake_case__  					=      sample, sample
			for t in range(UpperCamelCase_			,					time_step + scheduler.config.solver_order + 1		):
				snake_case__  					=      scheduler.step(UpperCamelCase_			,					UpperCamelCase_			,					UpperCamelCase_			,					**UpperCamelCase_		).prev_sample
				snake_case__  					=      new_scheduler.step(UpperCamelCase_			,					UpperCamelCase_			,					UpperCamelCase_			,					**UpperCamelCase_		).prev_sample

				assert torch.sum(torch.abs(output - new_output		)		) < 1E-5, "Scheduler outputs are not identical"





	def      _snake_case      (		self		)			->		Dict:
		pass





	def      _snake_case      (		self			,					UpperCamelCase_=0			,					**UpperCamelCase_		)			->		Union[str, Any]:
		snake_case__  					=      dict(self.forward_default_kwargs		)
		snake_case__  					=      kwargs.pop('num_inference_steps'			,					UpperCamelCase_		)
		snake_case__  					=      self.dummy_sample
		snake_case__  					=      0.1 * sample
		snake_case__  					=      [residual + 0.2, residual + 0.1_5, residual + 0.1_0]

		for scheduler_class in self.scheduler_classes:
			snake_case__  					=      self.get_scheduler_config()
			snake_case__  					=      scheduler_class(**UpperCamelCase_		)
			scheduler.set_timesteps(UpperCamelCase_		)

			# copy over dummy past residuals (must be after setting timesteps)
			snake_case__  					=      dummy_past_residuals[: scheduler.config.solver_order]

			with tempfile.TemporaryDirectory() as tmpdirname:
				scheduler.save_config(UpperCamelCase_		)
				snake_case__  					=      scheduler_class.from_pretrained(UpperCamelCase_		)
				# copy over dummy past residuals
				new_scheduler.set_timesteps(UpperCamelCase_		)

				# copy over dummy past residual (must be after setting timesteps)
				snake_case__  					=      dummy_past_residuals[: new_scheduler.config.solver_order]

			snake_case__  					=      scheduler.step(UpperCamelCase_			,					UpperCamelCase_			,					UpperCamelCase_			,					**UpperCamelCase_		).prev_sample
			snake_case__  					=      new_scheduler.step(UpperCamelCase_			,					UpperCamelCase_			,					UpperCamelCase_			,					**UpperCamelCase_		).prev_sample

			assert torch.sum(torch.abs(output - new_output		)		) < 1E-5, "Scheduler outputs are not identical"





	def      _snake_case      (		self			,					UpperCamelCase_=None			,					**UpperCamelCase_		)			->		Any:
		if scheduler is None:
			snake_case__  					=      self.scheduler_classes[0]
			snake_case__  					=      self.get_scheduler_config(**UpperCamelCase_		)
			snake_case__  					=      scheduler_class(**UpperCamelCase_		)

		snake_case__  					=      10
		snake_case__  					=      self.dummy_model()
		snake_case__  					=      self.dummy_sample_deter
		scheduler.set_timesteps(UpperCamelCase_		)

		for i, t in enumerate(scheduler.timesteps		):
			snake_case__  					=      model(UpperCamelCase_			,					UpperCamelCase_		)
			snake_case__  					=      scheduler.step(UpperCamelCase_			,					UpperCamelCase_			,					UpperCamelCase_		).prev_sample

		return sample





	def      _snake_case      (		self		)			->		Any:
		snake_case__  					=      dict(self.forward_default_kwargs		)

		snake_case__  					=      kwargs.pop('num_inference_steps'			,					UpperCamelCase_		)

		for scheduler_class in self.scheduler_classes:
			snake_case__  					=      self.get_scheduler_config()
			snake_case__  					=      scheduler_class(**UpperCamelCase_		)

			snake_case__  					=      self.dummy_sample
			snake_case__  					=      0.1 * sample

			if num_inference_steps is not None and hasattr(UpperCamelCase_			,					'set_timesteps'		):
				scheduler.set_timesteps(UpperCamelCase_		)
			elif num_inference_steps is not None and not hasattr(UpperCamelCase_			,					'set_timesteps'		):
				snake_case__  					=      num_inference_steps

			# copy over dummy past residuals (must be done after set_timesteps)
			snake_case__  					=      [residual + 0.2, residual + 0.1_5, residual + 0.1_0]
			snake_case__  					=      dummy_past_residuals[: scheduler.config.solver_order]

			snake_case__  					=      scheduler.timesteps[5]
			snake_case__  					=      scheduler.timesteps[6]

			snake_case__  					=      scheduler.step(UpperCamelCase_			,					UpperCamelCase_			,					UpperCamelCase_			,					**UpperCamelCase_		).prev_sample
			snake_case__  					=      scheduler.step(UpperCamelCase_			,					UpperCamelCase_			,					UpperCamelCase_			,					**UpperCamelCase_		).prev_sample

			self.assertEqual(output_a.shape			,					sample.shape		)
			self.assertEqual(output_a.shape			,					output_a.shape		)





	def      _snake_case      (		self		)			->		Optional[int]:
		for timesteps in [25, 50, 100, 999, 1000]:
			self.check_over_configs(num_train_timesteps=UpperCamelCase_		)





	def      _snake_case      (		self		)			->		Optional[int]:
		self.check_over_configs(thresholding=UpperCamelCase_		)
		for order in [1, 2, 3]:
			for solver_type in ["midpoint", "heun"]:
				for threshold in [0.5, 1.0, 2.0]:
					for prediction_type in ["epsilon", "sample"]:
						self.check_over_configs(
						    thresholding=UpperCamelCase_			,					prediction_type=UpperCamelCase_			,					sample_max_value=UpperCamelCase_			,					algorithm_type='dpmsolver++'			,					solver_order=UpperCamelCase_			,					solver_type=UpperCamelCase_			,					)





	def      _snake_case      (		self		)			->		List[str]:
		for prediction_type in ["epsilon", "v_prediction"]:
			self.check_over_configs(prediction_type=UpperCamelCase_		)





	def      _snake_case      (		self		)			->		Tuple:
		for algorithm_type in ["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"]:
			for solver_type in ["midpoint", "heun"]:
				for order in [1, 2, 3]:
					for prediction_type in ["epsilon", "sample"]:
						if algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]:
							if order == 3:
								continue
						else:
							self.check_over_configs(
							    solver_order=UpperCamelCase_			,					solver_type=UpperCamelCase_			,					prediction_type=UpperCamelCase_			,					algorithm_type=UpperCamelCase_			,					)
						snake_case__  					=      self.full_loop(
						    solver_order=UpperCamelCase_			,					solver_type=UpperCamelCase_			,					prediction_type=UpperCamelCase_			,					algorithm_type=UpperCamelCase_			,					)
						assert not torch.isnan(UpperCamelCase_		).any(), "Samples have nan numbers"





	def      _snake_case      (		self		)			->		List[str]:
		self.check_over_configs(lower_order_final=UpperCamelCase_		)
		self.check_over_configs(lower_order_final=UpperCamelCase_		)





	def      _snake_case      (		self		)			->		List[str]:
		self.check_over_configs(lambda_min_clipped=-float('inf'		)		)
		self.check_over_configs(lambda_min_clipped=-5.1		)





	def      _snake_case      (		self		)			->		Optional[Any]:
		self.check_over_configs(variance_type=UpperCamelCase_		)
		self.check_over_configs(variance_type='learned_range'		)





	def      _snake_case      (		self		)			->		Optional[int]:
		for num_inference_steps in [1, 2, 3, 5, 10, 50, 100, 999, 1000]:
			self.check_over_forward(num_inference_steps=UpperCamelCase_			,					time_step=0		)





	def      _snake_case      (		self		)			->		int:
		snake_case__  					=      self.full_loop()
		snake_case__  					=      torch.mean(torch.abs(UpperCamelCase_		)		)

		assert abs(result_mean.item() - 0.3_3_0_1		) < 1E-3





	def      _snake_case      (		self		)			->		int:
		snake_case__  					=      self.full_loop(thresholding=UpperCamelCase_			,					dynamic_thresholding_ratio=0.8_7			,					sample_max_value=0.5		)
		snake_case__  					=      torch.mean(torch.abs(UpperCamelCase_		)		)

		assert abs(result_mean.item() - 1.1_3_6_4		) < 1E-3





	def      _snake_case      (		self		)			->		Tuple:
		snake_case__  					=      self.full_loop(prediction_type='v_prediction'		)
		snake_case__  					=      torch.mean(torch.abs(UpperCamelCase_		)		)

		assert abs(result_mean.item() - 0.2_2_5_1		) < 1E-3





	def      _snake_case      (		self		)			->		Dict:
		snake_case__  					=      self.full_loop(prediction_type='v_prediction'			,					use_karras_sigmas=UpperCamelCase_		)
		snake_case__  					=      torch.mean(torch.abs(UpperCamelCase_		)		)

		assert abs(result_mean.item() - 0.2_0_9_6		) < 1E-3





	def      _snake_case      (		self		)			->		str:
		# make sure that iterating over schedulers with same config names gives same results
		# for defaults
		snake_case__  					=      DPMSolverMultistepScheduler(**self.get_scheduler_config()		)
		snake_case__  					=      self.full_loop(scheduler=UpperCamelCase_		)
		snake_case__  					=      torch.mean(torch.abs(UpperCamelCase_		)		)

		assert abs(result_mean.item() - 0.3_3_0_1		) < 1E-3

		snake_case__  					=      DPMSolverSinglestepScheduler.from_config(scheduler.config		)
		snake_case__  					=      UniPCMultistepScheduler.from_config(scheduler.config		)
		snake_case__  					=      DEISMultistepScheduler.from_config(scheduler.config		)
		snake_case__  					=      DPMSolverMultistepScheduler.from_config(scheduler.config		)

		snake_case__  					=      self.full_loop(scheduler=UpperCamelCase_		)
		snake_case__  					=      torch.mean(torch.abs(UpperCamelCase_		)		)

		assert abs(result_mean.item() - 0.3_3_0_1		) < 1E-3





	def      _snake_case      (		self		)			->		Any:
		snake_case__  					=      self.scheduler_classes[0]
		snake_case__  					=      self.get_scheduler_config(thresholding=UpperCamelCase_			,					dynamic_thresholding_ratio=0		)
		snake_case__  					=      scheduler_class(**UpperCamelCase_		)

		snake_case__  					=      10
		snake_case__  					=      self.dummy_model()
		snake_case__  					=      self.dummy_sample_deter.half()
		scheduler.set_timesteps(UpperCamelCase_		)

		for i, t in enumerate(scheduler.timesteps		):
			snake_case__  					=      model(UpperCamelCase_			,					UpperCamelCase_		)
			snake_case__  					=      scheduler.step(UpperCamelCase_			,					UpperCamelCase_			,					UpperCamelCase_		).prev_sample

		assert sample.dtype == torch.floataa





	def      _snake_case      (		self			,					**UpperCamelCase_		)			->		int:
		for scheduler_class in self.scheduler_classes:
			snake_case__  					=      self.get_scheduler_config(**UpperCamelCase_		)
			snake_case__  					=      scheduler_class(**UpperCamelCase_		)

			scheduler.set_timesteps(scheduler.config.num_train_timesteps		)
			assert len(scheduler.timesteps.unique()		) == scheduler.num_inference_steps










'''simple docstring'''




import shutil
import tempfile
import unittest

import pytest

from transformers import WhisperTokenizer, is_speech_available
from transformers.testing_utils import require_sentencepiece, require_torch, require_torchaudio

from .test_feature_extraction_whisper import floats_list


if is_speech_available():
			from transformers import WhisperFeatureExtractor, WhisperProcessor


a__				: Tuple									=    5_0358
a__				: str									=    5_0362






@require_torch
@require_torchaudio
@require_sentencepiece
class 							__snake_case    (				unittest.TestCase       ):





	def      _snake_case      (		self		)			->		Dict:
		snake_case__  					=      'openai/whisper-small.en'
		snake_case__  					=      tempfile.mkdtemp()





	def      _snake_case      (		self			,					**UpperCamelCase_		)			->		Optional[int]:
		return WhisperTokenizer.from_pretrained(self.checkpoint			,					**UpperCamelCase_		)





	def      _snake_case      (		self			,					**UpperCamelCase_		)			->		Any:
		return WhisperFeatureExtractor.from_pretrained(self.checkpoint			,					**UpperCamelCase_		)





	def      _snake_case      (		self		)			->		Tuple:
		shutil.rmtree(self.tmpdirname		)





	def      _snake_case      (		self		)			->		Any:
		snake_case__  					=      self.get_tokenizer()
		snake_case__  					=      self.get_feature_extractor()

		snake_case__  					=      WhisperProcessor(tokenizer=UpperCamelCase_			,					feature_extractor=UpperCamelCase_		)

		processor.save_pretrained(self.tmpdirname		)
		snake_case__  					=      WhisperProcessor.from_pretrained(self.tmpdirname		)

		self.assertEqual(processor.tokenizer.get_vocab()			,					tokenizer.get_vocab()		)
		self.assertIsInstance(processor.tokenizer			,					UpperCamelCase_		)

		self.assertEqual(processor.feature_extractor.to_json_string()			,					feature_extractor.to_json_string()		)
		self.assertIsInstance(processor.feature_extractor			,					UpperCamelCase_		)





	def      _snake_case      (		self		)			->		Dict:
		snake_case__  					=      WhisperProcessor(tokenizer=self.get_tokenizer()			,					feature_extractor=self.get_feature_extractor()		)
		processor.save_pretrained(self.tmpdirname		)

		snake_case__  					=      self.get_tokenizer(bos_token='(BOS)'			,					eos_token='(EOS)'		)
		snake_case__  					=      self.get_feature_extractor(do_normalize=UpperCamelCase_			,					padding_value=1.0		)

		snake_case__  					=      WhisperProcessor.from_pretrained(
		    self.tmpdirname			,					bos_token='(BOS)'			,					eos_token='(EOS)'			,					do_normalize=UpperCamelCase_			,					padding_value=1.0		)

		self.assertEqual(processor.tokenizer.get_vocab()			,					tokenizer_add_kwargs.get_vocab()		)
		self.assertIsInstance(processor.tokenizer			,					UpperCamelCase_		)

		self.assertEqual(processor.feature_extractor.to_json_string()			,					feature_extractor_add_kwargs.to_json_string()		)
		self.assertIsInstance(processor.feature_extractor			,					UpperCamelCase_		)





	def      _snake_case      (		self		)			->		List[str]:
		snake_case__  					=      self.get_feature_extractor()
		snake_case__  					=      self.get_tokenizer()

		snake_case__  					=      WhisperProcessor(tokenizer=UpperCamelCase_			,					feature_extractor=UpperCamelCase_		)

		snake_case__  					=      floats_list((3, 1000)		)

		snake_case__  					=      feature_extractor(UpperCamelCase_			,					return_tensors='np'		)
		snake_case__  					=      processor(UpperCamelCase_			,					return_tensors='np'		)

		for key in input_feat_extract.keys():
			self.assertAlmostEqual(input_feat_extract[key].sum()			,					input_processor[key].sum()			,					delta=1E-2		)





	def      _snake_case      (		self		)			->		Union[str, Any]:
		snake_case__  					=      self.get_feature_extractor()
		snake_case__  					=      self.get_tokenizer()

		snake_case__  					=      WhisperProcessor(tokenizer=UpperCamelCase_			,					feature_extractor=UpperCamelCase_		)

		snake_case__  					=      'This is a test string'

		snake_case__  					=      processor(text=UpperCamelCase_		)

		snake_case__  					=      tokenizer(UpperCamelCase_		)

		for key in encoded_tok.keys():
			self.assertListEqual(encoded_tok[key]			,					encoded_processor[key]		)





	def      _snake_case      (		self		)			->		Tuple:
		snake_case__  					=      self.get_feature_extractor()
		snake_case__  					=      self.get_tokenizer()

		snake_case__  					=      WhisperProcessor(tokenizer=UpperCamelCase_			,					feature_extractor=UpperCamelCase_		)

		snake_case__  					=      [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

		snake_case__  					=      processor.batch_decode(UpperCamelCase_		)
		snake_case__  					=      tokenizer.batch_decode(UpperCamelCase_		)

		self.assertListEqual(UpperCamelCase_			,					UpperCamelCase_		)





	def      _snake_case      (		self		)			->		List[str]:
		snake_case__  					=      self.get_feature_extractor()
		snake_case__  					=      self.get_tokenizer()

		snake_case__  					=      WhisperProcessor(tokenizer=UpperCamelCase_			,					feature_extractor=UpperCamelCase_		)

		self.assertListEqual(
		    processor.model_input_names			,					feature_extractor.model_input_names			,					msg='`processor` and `feature_extractor` model input names do not match'			,					)





	def      _snake_case      (		self		)			->		Dict:
		snake_case__  					=      self.get_feature_extractor()
		snake_case__  					=      self.get_tokenizer()

		snake_case__  					=      WhisperProcessor(tokenizer=UpperCamelCase_			,					feature_extractor=UpperCamelCase_		)
		snake_case__  					=      processor.get_decoder_prompt_ids(task='transcribe'			,					no_timestamps=UpperCamelCase_		)

		self.assertIsInstance(UpperCamelCase_			,					UpperCamelCase_		)
		for ids in forced_decoder_ids:
			self.assertIsInstance(UpperCamelCase_			,					(list, tuple)		)

		snake_case__  					=      [TRANSCRIBE, NOTIMESTAMPS]
		self.assertListEqual([ids[-1] for ids in forced_decoder_ids]			,					UpperCamelCase_		)





	def      _snake_case      (		self		)			->		List[str]:
		snake_case__  					=      WhisperProcessor(tokenizer=self.get_tokenizer()			,					feature_extractor=self.get_feature_extractor()		)
		snake_case__  					=      processor.get_prompt_ids('Mr. Quilter'		)
		snake_case__  					=      processor.tokenizer.decode(UpperCamelCase_		)

		self.assertListEqual(prompt_ids.tolist()			,					[5_0360, 1770, 13, 2264, 346, 353]		)
		self.assertEqual(UpperCamelCase_			,					'<|startofprev|> Mr. Quilter'		)





	def      _snake_case      (		self		)			->		Dict:
		snake_case__  					=      WhisperProcessor(tokenizer=self.get_tokenizer()			,					feature_extractor=self.get_feature_extractor()		)
		snake_case__  					=      processor.get_prompt_ids(''		)
		snake_case__  					=      processor.tokenizer.decode(UpperCamelCase_		)

		self.assertListEqual(prompt_ids.tolist()			,					[5_0360, 220]		)
		self.assertEqual(UpperCamelCase_			,					'<|startofprev|> '		)





	def      _snake_case      (		self		)			->		int:
		snake_case__  					=      WhisperProcessor(tokenizer=self.get_tokenizer()			,					feature_extractor=self.get_feature_extractor()		)

		def _test_prompt_error_raised_helper(UpperCamelCase_			,					UpperCamelCase_		):
			with pytest.raises(UpperCamelCase_		) as excinfo:
				processor.get_prompt_ids(UpperCamelCase_		)
			snake_case__  					=      F'''Encountered text in the prompt corresponding to disallowed special token: {special_token}.'''
			self.assertEqual(UpperCamelCase_			,					str(excinfo.value		)		)

		_test_prompt_error_raised_helper('<|startofprev|> test'			,					'<|startofprev|>'		)
		_test_prompt_error_raised_helper('test <|notimestamps|>'			,					'<|notimestamps|>'		)
		_test_prompt_error_raised_helper('test <|zh|> test <|transcribe|>'			,					'<|zh|>'		)



