





'''simple docstring'''




import json
import os
from typing import Dict, List, Optional, Tuple

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


_SCREAMING_SNAKE_CASE							     =			logging.get_logger(__name__)


_SCREAMING_SNAKE_CASE							     =			{
    """vocab_file""": """vocab.json""",
    """tokenizer_config_file""": """tokenizer_config.json""",
    """merges_file""": """merges.txt""",
}

_SCREAMING_SNAKE_CASE							     =			{
    """vocab_file""": {
        """facebook/s2t-wav2vec2-large-en-de""": (
            """https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/vocab.json"""
        ),
    },
    """tokenizer_config_file""": {
        """facebook/s2t-wav2vec2-large-en-de""": (
            """https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/tokenizer_config.json"""
        ),
    },
    """merges_file""": {
        """facebook/s2t-wav2vec2-large-en-de""": (
            """https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/merges.txt"""
        ),
    },
}

_SCREAMING_SNAKE_CASE							     =			"""</w>"""
_SCREAMING_SNAKE_CASE							     =			"""@@ """





def 	__lowerCamelCase			(      __lowerCAmelCase       :							Any		)		->  str:
   snake_case 						=   set()
   snake_case 						=   word[0]
   for char in word[1:]:
      pairs.add((prev_char, char)		)
      snake_case 						=   char
   return pairs


# Speech2Text2 has no max input length
_SCREAMING_SNAKE_CASE							     =			{"""facebook/s2t-wav2vec2-large-en-de""": 1024}


class 			_lowerCAmelCase		(			__lowercase	):






     """simple docstring"""






     snake_case_						      =   VOCAB_FILES_NAMES
     snake_case_						      =   PRETRAINED_VOCAB_FILES_MAP
     snake_case_						      =   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     snake_case_						      =   ['''input_ids''', '''attention_mask''']



     def __init__(   self							:			List[str]   ,       __snake_case							:			List[Any]   ,       __snake_case							:			int="<s>"   ,       __snake_case							:			int="<pad>"   ,       __snake_case							:			int="</s>"   ,       __snake_case							:			Optional[int]="<unk>"   ,       __snake_case							:			Any=False   ,       __snake_case							:			Optional[int]=None   ,       **__snake_case							:			Tuple   ,       )->	List[str]:
        super().__init__(
            unk_token=_lowerCAmelCase   ,       bos_token=_lowerCAmelCase   ,       eos_token=_lowerCAmelCase   ,       pad_token=_lowerCAmelCase   ,       do_lower_case=_lowerCAmelCase   ,       **_lowerCAmelCase   ,       )

        snake_case 						=   do_lower_case

        with open(_lowerCAmelCase   ,       encoding="""utf-8"""						) as vocab_handle:
           snake_case 						=   json.load(_lowerCAmelCase						)
        snake_case 						=   {v: k for k, v in self.encoder.items()}

        if merges_file is None:
           logger.info(f'''No merges files provided. {self.__class__.__name__} can only be used for decoding.'''						)

           snake_case 						=   None
           snake_case 						=   None
        else:
           with open(_lowerCAmelCase   ,       encoding="""utf-8"""						) as merges_handle:
              snake_case 						=   merges_handle.read().split("""\n"""						)[:-1]

           snake_case 						=   [tuple(merge.split()[:2]						) for merge in merges]
           snake_case 						=   dict(zip(_lowerCAmelCase   ,       range(len(_lowerCAmelCase						)						)						)						)
           snake_case 						=   {}



     @property
     def      lowerCAmelCase			(   self							:			List[Any]						)->	List[str]:
        return len(self.decoder						)



     def      lowerCAmelCase			(   self							:			Optional[Any]						)->	int:
        return dict(self.encoder   ,       **self.added_tokens_encoder						)



     def      lowerCAmelCase			(   self							:			str   ,       __snake_case							:			Tuple						)->	int:
        snake_case 						=   tuple(token[:-1]						) + (token[-1] + BPE_TOKEN_MERGES,)
        if token in self.cache:
           return self.cache[token]
        snake_case 						=   get_pairs(_lowerCAmelCase						)

        if not pairs:
           return token

        while True:
           snake_case 						=   min(_lowerCAmelCase   ,       key=lambda __snake_case						: self.bpe_ranks.get(_lowerCAmelCase   ,       float("""inf"""						)						)						)
           if bigram not in self.bpe_ranks:
              break
           snake_case 						=   bigram
           snake_case 						=   []
           snake_case 						=   0
           while i < len(_lowerCAmelCase						):
              try:
                 snake_case 						=   word.index(_lowerCAmelCase   ,       _lowerCAmelCase						)
              except ValueError:
                 new_word.extend(word[i:]						)
                 break
              else:
                 new_word.extend(word[i:j]						)
                 snake_case 						=   j

              if word[i] == first and i < len(_lowerCAmelCase						) - 1 and word[i + 1] == second:
                 new_word.append(first + second						)
                 i += 2
              else:
                 new_word.append(word[i]						)
                 i += 1
           snake_case 						=   tuple(_lowerCAmelCase						)
           snake_case 						=   new_word
           if len(_lowerCAmelCase						) == 1:
              break
           else:
              snake_case 						=   get_pairs(_lowerCAmelCase						)
        snake_case 						=   ''' '''.join(_lowerCAmelCase						)
        if word == "\n  " + BPE_TOKEN_MERGES:
           snake_case 						=   '''\n''' + BPE_TOKEN_MERGES

        if word.endswith(_lowerCAmelCase						):
           snake_case 						=   word.replace(_lowerCAmelCase   ,       """"""						)

        snake_case 						=   word.replace(""" """   ,       _lowerCAmelCase						)
        snake_case 						=   word
        return word



     def      lowerCAmelCase			(   self							:			Any   ,       __snake_case							:			List[str]						)->	Optional[Any]:




        if self.bpe_ranks is None:
           raise ValueError(
               """This tokenizer was instantiated without a `merges.txt` file, so"""
               """ that it can only be used for decoding, not for encoding."""
               """Make sure to provide `merges.txt` file at instantiation to enable """
               """encoding."""						)

        if self.do_lower_case:
           snake_case 						=   text.lower()

        snake_case 						=   text.split()

        snake_case 						=   []
        for token in text:
           if token:
              split_tokens.extend(list(self.bpe(_lowerCAmelCase						).split(""" """						)						)						)

        return split_tokens



     def      lowerCAmelCase			(   self							:			List[str]   ,       __snake_case							:			Optional[Any]						)->	List[Any]:
        return self.encoder.get(_lowerCAmelCase   ,       self.encoder.get(self.unk_token						)						)



     def      lowerCAmelCase			(   self							:			Optional[Any]   ,       __snake_case							:			Any						)->	Optional[int]:
        snake_case 						=   self.decoder.get(_lowerCAmelCase   ,       self.unk_token						)
        return result



     def      lowerCAmelCase			(   self							:			List[Any]   ,       __snake_case							:			str						)->	List[Any]:
        snake_case 						=   ''' '''.join(_lowerCAmelCase						)

        # make sure @@ tokens are concatenated
        snake_case 						=   ''''''.join(string.split(_lowerCAmelCase						)						)

        return string



     def      lowerCAmelCase			(   self							:			Dict   ,       __snake_case							:			Optional[Any]   ,       __snake_case							:			List[str] = None						)->	List[Any]:




        if not os.path.isdir(_lowerCAmelCase						):
           logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''						)
           return
        snake_case 						=   os.path.join(
            _lowerCAmelCase   ,       (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]						)
        snake_case 						=   os.path.join(
            _lowerCAmelCase   ,       (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""merges_file"""]						)

        with open(_lowerCAmelCase   ,       """w"""   ,       encoding="""utf-8"""						) as f:
           f.write(json.dumps(self.encoder   ,       indent=2   ,       sort_keys=_lowerCAmelCase   ,       ensure_ascii=_lowerCAmelCase						) + """\n"""						)

        snake_case 						=   0
        if self.bpe_ranks is None:
           return (vocab_file,)

        with open(_lowerCAmelCase   ,       """w"""   ,       encoding="""utf-8"""						) as writer:
           for bpe_tokens, token_index in sorted(self.bpe_ranks.items()   ,       key=lambda __snake_case						: kv[1]						):
              if index != token_index:
                 logger.warning(
                     f'''Saving vocabulary to {merges_file}: BPE merge indices are not consecutive.'''
                     """ Please check that the tokenizer is not corrupted!"""						)
                 snake_case 						=   token_index
              writer.write(""" """.join(_lowerCAmelCase						) + """\n"""						)
              index += 1

        return (vocab_file, merges_file)








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 				a			(      __lowercase   ,unittest.TestCase			):
				SCREAMING_SNAKE_CASE__     :   List[Any]		      =     ConsistencyModelPipeline
				SCREAMING_SNAKE_CASE__     :   Any		      =     UNCONDITIONAL_IMAGE_GENERATION_PARAMS
				SCREAMING_SNAKE_CASE__     :   str		      =     UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

				# Override required_optional_params to remove num_images_per_prompt
				SCREAMING_SNAKE_CASE__     :   Optional[int]		      =     frozenset(
				    [
				        '''num_inference_steps''',
				        '''generator''',
				        '''latents''',
				        '''output_type''',
				        '''return_dict''',
				        '''callback''',
				        '''callback_steps''',
				    ]			)






				@property
				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				int           =   UNetaDModel.from_pretrained(
										    '''diffusers/consistency-models-test'''			, subfolder='''test_unet'''			, )
										return unet






				@property
				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[Any]           =   UNetaDModel.from_pretrained(
										    '''diffusers/consistency-models-test'''			, subfolder='''test_unet_class_cond'''			, )
										return unet






				def 			snake_case_    (  self			, _lowerCAmelCase=False					):

										"""simple docstring"""




										if class_cond:
																__SCREAMING_SNAKE_CASE:				List[Any]           =   self.dummy_cond_unet
										else:
																__SCREAMING_SNAKE_CASE:				Tuple           =   self.dummy_uncond_unet

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

										__SCREAMING_SNAKE_CASE:				List[str]           =   {
										    '''unet''': unet,
										    '''scheduler''': scheduler,
										}

										return components






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase=0					):

										"""simple docstring"""




										if str(_lowerCAmelCase					).startswith('''mps'''					):
																__SCREAMING_SNAKE_CASE:				List[str]           =   torch.manual_seed(_lowerCAmelCase					)
										else:
																__SCREAMING_SNAKE_CASE:				List[str]           =   torch.Generator(device=_lowerCAmelCase					).manual_seed(_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   {
										    '''batch_size''': 1,
										    '''num_inference_steps''': None,
										    '''timesteps''': [22, 0],
										    '''generator''': generator,
										    '''output_type''': '''np''',
										}

										return inputs






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Dict           =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   self.get_dummy_components()
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   ConsistencyModelPipeline(**_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Any           =   pipe.to(_lowerCAmelCase					)
										pipe.set_progress_bar_config(disable=_lowerCAmelCase					)

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

										__SCREAMING_SNAKE_CASE:				int           =   image[0, -3:, -3:, -1]
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]					)

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






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Dict           =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__SCREAMING_SNAKE_CASE:				str           =   self.get_dummy_components(class_cond=_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Any           =   ConsistencyModelPipeline(**_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Any           =   pipe.to(_lowerCAmelCase					)
										pipe.set_progress_bar_config(disable=_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				Dict           =   self.get_dummy_inputs(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Dict           =   0
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   pipe(**_lowerCAmelCase					).images
										assert image.shape == (1, 32, 32, 3)

										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   image[0, -3:, -3:, -1]
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]					)

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






				def 			snake_case_    (  self					):

										"""simple docstring"""




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

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

										__SCREAMING_SNAKE_CASE:				Tuple           =   image[0, -3:, -3:, -1]
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]					)

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






				def 			snake_case_    (  self					):

										"""simple docstring"""




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

										__SCREAMING_SNAKE_CASE:				Tuple           =   self.get_dummy_inputs(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Any           =   1
										__SCREAMING_SNAKE_CASE:				List[Any]           =   None
										__SCREAMING_SNAKE_CASE:				Tuple           =   0
										__SCREAMING_SNAKE_CASE:				List[str]           =   pipe(**_lowerCAmelCase					).images
										assert image.shape == (1, 32, 32, 3)

										__SCREAMING_SNAKE_CASE:				List[Any]           =   image[0, -3:, -3:, -1]
										__SCREAMING_SNAKE_CASE:				Any           =   np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]					)

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





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






				def 			snake_case_    (  self					):

										"""simple docstring"""




										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()






				def 			snake_case_    (  self			, _lowerCAmelCase=0			, _lowerCAmelCase=False			, _lowerCAmelCase="cpu"			, _lowerCAmelCase=torch.floataa			, _lowerCAmelCase=(1, 3, 64, 64)					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				str           =   torch.manual_seed(_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				Dict           =   {
										    '''num_inference_steps''': None,
										    '''timesteps''': [22, 0],
										    '''class_labels''': 0,
										    '''generator''': generator,
										    '''output_type''': '''np''',
										}

										if get_fixed_latents:
																__SCREAMING_SNAKE_CASE:				Dict           =   self.get_fixed_latents(seed=_lowerCAmelCase			, device=_lowerCAmelCase			, dtype=_lowerCAmelCase			, shape=_lowerCAmelCase					)
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   latents

										return inputs






				def 			snake_case_    (  self			, _lowerCAmelCase=0			, _lowerCAmelCase="cpu"			, _lowerCAmelCase=torch.floataa			, _lowerCAmelCase=(1, 3, 64, 64)					):

										"""simple docstring"""




										if type(_lowerCAmelCase					) == str:
																__SCREAMING_SNAKE_CASE:				List[Any]           =   torch.device(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Dict           =   torch.Generator(device=_lowerCAmelCase					).manual_seed(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				int           =   randn_tensor(_lowerCAmelCase			, generator=_lowerCAmelCase			, device=_lowerCAmelCase			, dtype=_lowerCAmelCase					)
										return latents






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   UNetaDModel.from_pretrained('''diffusers/consistency_models'''			, subfolder='''diffusers_cd_imagenet64_l2'''					)
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   CMStochasticIterativeScheduler(
										    num_train_timesteps=40			, sigma_min=0.002			, sigma_max=80.0			, )
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   ConsistencyModelPipeline(unet=_lowerCAmelCase			, scheduler=_lowerCAmelCase					)
										pipe.to(torch_device=_lowerCAmelCase					)
										pipe.set_progress_bar_config(disable=_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				int           =   self.get_inputs()
										__SCREAMING_SNAKE_CASE:				List[Any]           =   pipe(**_lowerCAmelCase					).images
										assert image.shape == (1, 64, 64, 3)

										__SCREAMING_SNAKE_CASE:				Dict           =   image[0, -3:, -3:, -1]

										__SCREAMING_SNAKE_CASE:				Any           =   np.array([0.0888, 0.0881, 0.0666, 0.0479, 0.0292, 0.0195, 0.0201, 0.0163, 0.0254]					)

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






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[Any]           =   UNetaDModel.from_pretrained('''diffusers/consistency_models'''			, subfolder='''diffusers_cd_imagenet64_l2'''					)
										__SCREAMING_SNAKE_CASE:				str           =   CMStochasticIterativeScheduler(
										    num_train_timesteps=40			, sigma_min=0.002			, sigma_max=80.0			, )
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   ConsistencyModelPipeline(unet=_lowerCAmelCase			, scheduler=_lowerCAmelCase					)
										pipe.to(torch_device=_lowerCAmelCase					)
										pipe.set_progress_bar_config(disable=_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				List[str]           =   self.get_inputs()
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   1
										__SCREAMING_SNAKE_CASE:				List[Any]           =   None
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   pipe(**_lowerCAmelCase					).images
										assert image.shape == (1, 64, 64, 3)

										__SCREAMING_SNAKE_CASE:				Dict           =   image[0, -3:, -3:, -1]

										__SCREAMING_SNAKE_CASE:				Any           =   np.array([0.0340, 0.0152, 0.0063, 0.0267, 0.0221, 0.0107, 0.0416, 0.0186, 0.0217]					)

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






				@require_torch_a
				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Tuple           =   UNetaDModel.from_pretrained('''diffusers/consistency_models'''			, subfolder='''diffusers_cd_imagenet64_l2'''					)
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   CMStochasticIterativeScheduler(
										    num_train_timesteps=40			, sigma_min=0.002			, sigma_max=80.0			, )
										__SCREAMING_SNAKE_CASE:				List[Any]           =   ConsistencyModelPipeline(unet=_lowerCAmelCase			, scheduler=_lowerCAmelCase					)
										pipe.to(torch_device=_lowerCAmelCase			, torch_dtype=torch.floataa					)
										pipe.set_progress_bar_config(disable=_lowerCAmelCase					)

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

										__SCREAMING_SNAKE_CASE:				List[Any]           =   image[0, -3:, -3:, -1]

										__SCREAMING_SNAKE_CASE:				Optional[int]           =   np.array([0.1875, 0.1428, 0.1289, 0.2151, 0.2092, 0.1477, 0.1877, 0.1641, 0.1353]					)

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






				@require_torch_a
				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   UNetaDModel.from_pretrained('''diffusers/consistency_models'''			, subfolder='''diffusers_cd_imagenet64_l2'''					)
										__SCREAMING_SNAKE_CASE:				List[str]           =   CMStochasticIterativeScheduler(
										    num_train_timesteps=40			, sigma_min=0.002			, sigma_max=80.0			, )
										__SCREAMING_SNAKE_CASE:				List[Any]           =   ConsistencyModelPipeline(unet=_lowerCAmelCase			, scheduler=_lowerCAmelCase					)
										pipe.to(torch_device=_lowerCAmelCase			, torch_dtype=torch.floataa					)
										pipe.set_progress_bar_config(disable=_lowerCAmelCase					)

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

										__SCREAMING_SNAKE_CASE:				Tuple           =   image[0, -3:, -3:, -1]

										__SCREAMING_SNAKE_CASE:				Tuple           =   np.array([0.1663, 0.1948, 0.2275, 0.1680, 0.1204, 0.1245, 0.1858, 0.1338, 0.2095]					)

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



