

'''simple docstring'''


import gc
import unittest

import numpy as np
import torch
from torch.backends.cuda import sdp_kernel

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)
from diffusers.utils import randn_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_a, require_torch_gpu

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class 					lowercase						(  a_,					unittest.TestCase ):
							_lowerCamelCase		:      Optional[int]=   ConsistencyModelPipeline
							_lowerCamelCase		:      Optional[int]=   UNCONDITIONAL_IMAGE_GENERATION_PARAMS
							_lowerCamelCase		:      Tuple=   UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

							# Override required_optional_params to remove num_images_per_prompt
							_lowerCamelCase		:      Tuple=   frozenset(
							    [
							        "num_inference_steps",
							        "generator",
							        "latents",
							        "output_type",
							        "return_dict",
							        "callback",
							        "callback_steps",
							    ] )


							@property
							def 						_snake_case							(			self)							->				int:
								UpperCAmelCase_		:      Union[str, Any]       	=					UNetaDModel.from_pretrained(
								    'diffusers/consistency-models-test'    ,			subfolder='test_unet'    ,			)
								return unet


							@property
							def 						_snake_case							(			self)							->				List[str]:
								UpperCAmelCase_		:      str       	=					UNetaDModel.from_pretrained(
								    'diffusers/consistency-models-test'    ,			subfolder='test_unet_class_cond'    ,			)
								return unet


							def 						_snake_case							(			self    ,			_snake_case=False)							->				Dict:
								if class_cond:
									UpperCAmelCase_		:      int       	=					self.dummy_cond_unet
								else:
									UpperCAmelCase_		:      Optional[int]       	=					self.dummy_uncond_unet

								# Default to CM multistep sampler
								UpperCAmelCase_		:      int       	=					CMStochasticIterativeScheduler(
								    num_train_timesteps=40    ,			sigma_min=0.002    ,			sigma_max=80.0    ,			)

								UpperCAmelCase_		:      Union[str, Any]       	=					{
								    'unet': unet,
								    'scheduler': scheduler,
								}

								return components


							def 						_snake_case							(			self    ,			_snake_case    ,			_snake_case=0)							->				Tuple:
								if str(_snake_case).startswith('mps'):
									UpperCAmelCase_		:      int       	=					torch.manual_seed(_snake_case)
								else:
									UpperCAmelCase_		:      Union[str, Any]       	=					torch.Generator(device=_snake_case).manual_seed(_snake_case)

								UpperCAmelCase_		:      Dict       	=					{
								    'batch_size': 1,
								    'num_inference_steps': None,
								    'timesteps': [22, 0],
								    'generator': generator,
								    'output_type': 'np',
								}

								return inputs


							def 						_snake_case							(			self)							->				str:
								UpperCAmelCase_		:      Optional[Any]       	=					'cpu'  # ensure determinism for the device-dependent torch.Generator
								UpperCAmelCase_		:      Any       	=					self.get_dummy_components()
								UpperCAmelCase_		:      Dict       	=					ConsistencyModelPipeline(**_snake_case)
								UpperCAmelCase_		:      Optional[Any]       	=					pipe.to(_snake_case)
								pipe.set_progress_bar_config(disable=_snake_case)

								UpperCAmelCase_		:      Optional[int]       	=					self.get_dummy_inputs(_snake_case)
								UpperCAmelCase_		:      Union[str, Any]       	=					pipe(**_snake_case).images
								assert image.shape == (1, 32, 32, 3)

								UpperCAmelCase_		:      int       	=					image[0, -3:, -3:, -1]
								UpperCAmelCase_		:      str       	=					np.array([0.3_572, 0.6_273, 0.4_031, 0.3_961, 0.4_321, 0.5_730, 0.5_266, 0.4_780, 0.5_004])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3


							def 						_snake_case							(			self)							->				Optional[Any]:
								UpperCAmelCase_		:      List[str]       	=					'cpu'  # ensure determinism for the device-dependent torch.Generator
								UpperCAmelCase_		:      Optional[Any]       	=					self.get_dummy_components(class_cond=_snake_case)
								UpperCAmelCase_		:      Union[str, Any]       	=					ConsistencyModelPipeline(**_snake_case)
								UpperCAmelCase_		:      Dict       	=					pipe.to(_snake_case)
								pipe.set_progress_bar_config(disable=_snake_case)

								UpperCAmelCase_		:      int       	=					self.get_dummy_inputs(_snake_case)
								UpperCAmelCase_		:      Optional[int]       	=					0
								UpperCAmelCase_		:      Any       	=					pipe(**_snake_case).images
								assert image.shape == (1, 32, 32, 3)

								UpperCAmelCase_		:      int       	=					image[0, -3:, -3:, -1]
								UpperCAmelCase_		:      Union[str, Any]       	=					np.array([0.3_572, 0.6_273, 0.4_031, 0.3_961, 0.4_321, 0.5_730, 0.5_266, 0.4_780, 0.5_004])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3


							def 						_snake_case							(			self)							->				int:
								UpperCAmelCase_		:      Tuple       	=					'cpu'  # ensure determinism for the device-dependent torch.Generator
								UpperCAmelCase_		:      Union[str, Any]       	=					self.get_dummy_components()
								UpperCAmelCase_		:      Dict       	=					ConsistencyModelPipeline(**_snake_case)
								UpperCAmelCase_		:      Optional[Any]       	=					pipe.to(_snake_case)
								pipe.set_progress_bar_config(disable=_snake_case)

								UpperCAmelCase_		:      List[Any]       	=					self.get_dummy_inputs(_snake_case)
								UpperCAmelCase_		:      Optional[Any]       	=					1
								UpperCAmelCase_		:      Union[str, Any]       	=					None
								UpperCAmelCase_		:      Union[str, Any]       	=					pipe(**_snake_case).images
								assert image.shape == (1, 32, 32, 3)

								UpperCAmelCase_		:      str       	=					image[0, -3:, -3:, -1]
								UpperCAmelCase_		:      Union[str, Any]       	=					np.array([0.5_004, 0.5_004, 0.4_994, 0.5_008, 0.4_976, 0.5_018, 0.4_990, 0.4_982, 0.4_987])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3





							def 						_snake_case							(			self)							->				Any:
								UpperCAmelCase_		:      int       	=					'cpu'  # ensure determinism for the device-dependent torch.Generator
								UpperCAmelCase_		:      List[str]       	=					self.get_dummy_components(class_cond=_snake_case)
								UpperCAmelCase_		:      Optional[int]       	=					ConsistencyModelPipeline(**_snake_case)
								UpperCAmelCase_		:      Tuple       	=					pipe.to(_snake_case)
								pipe.set_progress_bar_config(disable=_snake_case)

								UpperCAmelCase_		:      Tuple       	=					self.get_dummy_inputs(_snake_case)
								UpperCAmelCase_		:      List[str]       	=					1
								UpperCAmelCase_		:      str       	=					None
								UpperCAmelCase_		:      int       	=					0
								UpperCAmelCase_		:      Dict       	=					pipe(**_snake_case).images
								assert image.shape == (1, 32, 32, 3)

								UpperCAmelCase_		:      Optional[int]       	=					image[0, -3:, -3:, -1]
								UpperCAmelCase_		:      int       	=					np.array([0.5_004, 0.5_004, 0.4_994, 0.5_008, 0.4_976, 0.5_018, 0.4_990, 0.4_982, 0.4_987])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3



@slow
@require_torch_gpu
class 					lowercase						(  unittest.TestCase ):


							def 						_snake_case							(			self)							->				Optional[Any]:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()


							def 						_snake_case							(			self    ,			_snake_case=0    ,			_snake_case=False    ,			_snake_case="cpu"    ,			_snake_case=torch.floataa    ,			_snake_case=(1, 3, 64, 64))							->				List[str]:
								UpperCAmelCase_		:      str       	=					torch.manual_seed(_snake_case)

								UpperCAmelCase_		:      Union[str, Any]       	=					{
								    'num_inference_steps': None,
								    'timesteps': [22, 0],
								    'class_labels': 0,
								    'generator': generator,
								    'output_type': 'np',
								}

								if get_fixed_latents:
									UpperCAmelCase_		:      Any       	=					self.get_fixed_latents(seed=_snake_case    ,			device=_snake_case    ,			dtype=_snake_case    ,			shape=_snake_case)
									UpperCAmelCase_		:      int       	=					latents

								return inputs


							def 						_snake_case							(			self    ,			_snake_case=0    ,			_snake_case="cpu"    ,			_snake_case=torch.floataa    ,			_snake_case=(1, 3, 64, 64))							->				Any:
								if type(_snake_case) == str:
									UpperCAmelCase_		:      Any       	=					torch.device(_snake_case)
								UpperCAmelCase_		:      Optional[int]       	=					torch.Generator(device=_snake_case).manual_seed(_snake_case)
								UpperCAmelCase_		:      List[Any]       	=					randn_tensor(_snake_case    ,			generator=_snake_case    ,			device=_snake_case    ,			dtype=_snake_case)
								return latents


							def 						_snake_case							(			self)							->				List[Any]:
								UpperCAmelCase_		:      Dict       	=					UNetaDModel.from_pretrained('diffusers/consistency_models'    ,			subfolder='diffusers_cd_imagenet64_l2')
								UpperCAmelCase_		:      str       	=					CMStochasticIterativeScheduler(
								    num_train_timesteps=40    ,			sigma_min=0.002    ,			sigma_max=80.0    ,			)
								UpperCAmelCase_		:      List[Any]       	=					ConsistencyModelPipeline(unet=_snake_case    ,			scheduler=_snake_case)
								pipe.to(torch_device=_snake_case)
								pipe.set_progress_bar_config(disable=_snake_case)

								UpperCAmelCase_		:      List[Any]       	=					self.get_inputs()
								UpperCAmelCase_		:      Dict       	=					pipe(**_snake_case).images
								assert image.shape == (1, 64, 64, 3)

								UpperCAmelCase_		:      Optional[int]       	=					image[0, -3:, -3:, -1]

								UpperCAmelCase_		:      Tuple       	=					np.array([0.0_888, 0.0_881, 0.0_666, 0.0_479, 0.0_292, 0.0_195, 0.0_201, 0.0_163, 0.0_254])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 2e-2


							def 						_snake_case							(			self)							->				str:
								UpperCAmelCase_		:      List[Any]       	=					UNetaDModel.from_pretrained('diffusers/consistency_models'    ,			subfolder='diffusers_cd_imagenet64_l2')
								UpperCAmelCase_		:      str       	=					CMStochasticIterativeScheduler(
								    num_train_timesteps=40    ,			sigma_min=0.002    ,			sigma_max=80.0    ,			)
								UpperCAmelCase_		:      Any       	=					ConsistencyModelPipeline(unet=_snake_case    ,			scheduler=_snake_case)
								pipe.to(torch_device=_snake_case)
								pipe.set_progress_bar_config(disable=_snake_case)

								UpperCAmelCase_		:      str       	=					self.get_inputs()
								UpperCAmelCase_		:      Optional[int]       	=					1
								UpperCAmelCase_		:      Optional[Any]       	=					None
								UpperCAmelCase_		:      Tuple       	=					pipe(**_snake_case).images
								assert image.shape == (1, 64, 64, 3)

								UpperCAmelCase_		:      str       	=					image[0, -3:, -3:, -1]

								UpperCAmelCase_		:      Dict       	=					np.array([0.0_340, 0.0_152, 0.0_063, 0.0_267, 0.0_221, 0.0_107, 0.0_416, 0.0_186, 0.0_217])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 2e-2


							@require_torch_a
							def 						_snake_case							(			self)							->				Optional[Any]:
								UpperCAmelCase_		:      int       	=					UNetaDModel.from_pretrained('diffusers/consistency_models'    ,			subfolder='diffusers_cd_imagenet64_l2')
								UpperCAmelCase_		:      List[Any]       	=					CMStochasticIterativeScheduler(
								    num_train_timesteps=40    ,			sigma_min=0.002    ,			sigma_max=80.0    ,			)
								UpperCAmelCase_		:      Optional[int]       	=					ConsistencyModelPipeline(unet=_snake_case    ,			scheduler=_snake_case)
								pipe.to(torch_device=_snake_case    ,			torch_dtype=torch.floataa)
								pipe.set_progress_bar_config(disable=_snake_case)

								UpperCAmelCase_		:      Dict       	=					self.get_inputs(get_fixed_latents=_snake_case    ,			device=_snake_case)
								# Ensure usage of flash attention in torch 2.0
								with sdp_kernel(enable_flash=_snake_case    ,			enable_math=_snake_case    ,			enable_mem_efficient=_snake_case):
									UpperCAmelCase_		:      int       	=					pipe(**_snake_case).images
								assert image.shape == (1, 64, 64, 3)

								UpperCAmelCase_		:      List[str]       	=					image[0, -3:, -3:, -1]

								UpperCAmelCase_		:      List[str]       	=					np.array([0.1_875, 0.1_428, 0.1_289, 0.2_151, 0.2_092, 0.1_477, 0.1_877, 0.1_641, 0.1_353])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3





							@require_torch_a
							def 						_snake_case							(			self)							->				List[Any]:
								UpperCAmelCase_		:      Dict       	=					UNetaDModel.from_pretrained('diffusers/consistency_models'    ,			subfolder='diffusers_cd_imagenet64_l2')
								UpperCAmelCase_		:      Optional[Any]       	=					CMStochasticIterativeScheduler(
								    num_train_timesteps=40    ,			sigma_min=0.002    ,			sigma_max=80.0    ,			)
								UpperCAmelCase_		:      str       	=					ConsistencyModelPipeline(unet=_snake_case    ,			scheduler=_snake_case)
								pipe.to(torch_device=_snake_case    ,			torch_dtype=torch.floataa)
								pipe.set_progress_bar_config(disable=_snake_case)

								UpperCAmelCase_		:      Any       	=					self.get_inputs(get_fixed_latents=_snake_case    ,			device=_snake_case)
								UpperCAmelCase_		:      str       	=					1
								UpperCAmelCase_		:      Optional[int]       	=					None
								# Ensure usage of flash attention in torch 2.0
								with sdp_kernel(enable_flash=_snake_case    ,			enable_math=_snake_case    ,			enable_mem_efficient=_snake_case):
									UpperCAmelCase_		:      List[Any]       	=					pipe(**_snake_case).images
								assert image.shape == (1, 64, 64, 3)

								UpperCAmelCase_		:      int       	=					image[0, -3:, -3:, -1]

								UpperCAmelCase_		:      Dict       	=					np.array([0.1_663, 0.1_948, 0.2_275, 0.1_680, 0.1_204, 0.1_245, 0.1_858, 0.1_338, 0.2_095])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3









'''simple docstring'''


import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin
from .feature_extraction_wavaveca import WavaVecaFeatureExtractor
from .tokenization_wavaveca import WavaVecaCTCTokenizer



class 					lowercase						(  a_ ):
							_lowerCamelCase		:      List[str]=   "Wav2Vec2FeatureExtractor"
							_lowerCamelCase		:      Any=   "AutoTokenizer"


							def __init__(			self    ,			_snake_case    ,			_snake_case)							->				List[str]:
								super().__init__(_snake_case    ,			_snake_case)
								UpperCAmelCase_		:      int       	=					self.feature_extractor
								UpperCAmelCase_		:      str       	=					False


							@classmethod
							def 						_snake_case							(			cls    ,			_snake_case    ,			**_snake_case)							->				List[Any]:
								try:
									return super().from_pretrained(_snake_case    ,			**_snake_case)
								except OSError:
									warnings.warn(
									    F"""Loading a tokenizer inside {cls.__name__} from a config that does not"""
									    ' include a `tokenizer_class` attribute is deprecated and will be '
									    'removed in v5. Please add `\'tokenizer_class\': \'Wav2Vec2CTCTokenizer\'`'
									    ' attribute to either your `config.json` or `tokenizer_config.json` '
									    'file to suppress this warning: '    ,			_snake_case    ,			)

									UpperCAmelCase_		:      Any       	=					WavaVecaFeatureExtractor.from_pretrained(_snake_case    ,			**_snake_case)
									UpperCAmelCase_		:      List[str]       	=					WavaVecaCTCTokenizer.from_pretrained(_snake_case    ,			**_snake_case)

									return cls(feature_extractor=_snake_case    ,			tokenizer=_snake_case)


							def __call__(			self    ,			*_snake_case    ,			**_snake_case)							->				Optional[Any]:
								# For backward compatibility
								if self._in_target_context_manager:
									return self.current_processor(*_snake_case    ,			**_snake_case)

								if "raw_speech" in kwargs:
									warnings.warn('Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.')
									UpperCAmelCase_		:      List[str]       	=					kwargs.pop('raw_speech')
								else:
									UpperCAmelCase_		:      Dict       	=					kwargs.pop('audio'    ,			_snake_case)
								UpperCAmelCase_		:      str       	=					kwargs.pop('sampling_rate'    ,			_snake_case)
								UpperCAmelCase_		:      List[str]       	=					kwargs.pop('text'    ,			_snake_case)
								if len(_snake_case) > 0:
									UpperCAmelCase_		:      List[str]       	=					args[0]
									UpperCAmelCase_		:      Tuple       	=					args[1:]

								if audio is None and text is None:
									raise ValueError('You need to specify either an `audio` or `text` input to process.')

								if audio is not None:
									UpperCAmelCase_		:      Union[str, Any]       	=					self.feature_extractor(_snake_case    ,			*_snake_case    ,			sampling_rate=_snake_case    ,			**_snake_case)
								if text is not None:
									UpperCAmelCase_		:      Optional[int]       	=					self.tokenizer(_snake_case    ,			**_snake_case)

								if text is None:
									return inputs
								elif audio is None:
									return encodings
								else:
									UpperCAmelCase_		:      List[str]       	=					encodings['input_ids']
									return inputs


							def 						_snake_case							(			self    ,			*_snake_case    ,			**_snake_case)							->				str:
								# For backward compatibility
								if self._in_target_context_manager:
									return self.current_processor.pad(*_snake_case    ,			**_snake_case)

								UpperCAmelCase_		:      Union[str, Any]       	=					kwargs.pop('input_features'    ,			_snake_case)
								UpperCAmelCase_		:      Dict       	=					kwargs.pop('labels'    ,			_snake_case)
								if len(_snake_case) > 0:
									UpperCAmelCase_		:      Optional[Any]       	=					args[0]
									UpperCAmelCase_		:      int       	=					args[1:]

								if input_features is not None:
									UpperCAmelCase_		:      Dict       	=					self.feature_extractor.pad(_snake_case    ,			*_snake_case    ,			**_snake_case)
								if labels is not None:
									UpperCAmelCase_		:      Union[str, Any]       	=					self.tokenizer.pad(_snake_case    ,			**_snake_case)

								if labels is None:
									return input_features
								elif input_features is None:
									return labels
								else:
									UpperCAmelCase_		:      Union[str, Any]       	=					labels['input_ids']
									return input_features


							def 						_snake_case							(			self    ,			*_snake_case    ,			**_snake_case)							->				Dict:
								return self.tokenizer.batch_decode(*_snake_case    ,			**_snake_case)


							def 						_snake_case							(			self    ,			*_snake_case    ,			**_snake_case)							->				Tuple:
								return self.tokenizer.decode(*_snake_case    ,			**_snake_case)





							@contextmanager
							def 						_snake_case							(			self)							->				Tuple:
								warnings.warn(
								    '`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '
								    'labels by using the argument `text` of the regular `__call__` method (either in the same call as '
								    'your audio inputs, or in a separate call.')
								UpperCAmelCase_		:      Optional[Any]       	=					True
								UpperCAmelCase_		:      Tuple       	=					self.tokenizer
								yield
								UpperCAmelCase_		:      Tuple       	=					self.feature_extractor
								UpperCAmelCase_		:      Tuple       	=					False






