
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 						__lowerCAmelCase	(	lowerCAmelCase_ , unittest.TestCase      ):



      """simple docstring"""







      A__	:			str        =    ConsistencyModelPipeline
      A__	:			Optional[Any]        =    UNCONDITIONAL_IMAGE_GENERATION_PARAMS
      A__	:			Optional[int]        =    UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

      # Override required_optional_params to remove num_images_per_prompt
      A__	:			List[Any]        =    frozenset(
          [
              '''num_inference_steps''',
              '''generator''',
              '''latents''',
              '''output_type''',
              '''return_dict''',
              '''callback''',
              '''callback_steps''',
          ]      )




      @property
      def 		snake_case_			(      self		:       str				):
          __lowercase		:			List[str]	    =       UNetaDModel.from_pretrained(
              '''diffusers/consistency-models-test'''				,       subfolder='''test_unet'''				,       )
          return unet




      @property
      def 		snake_case_			(      self		:       Dict				):
          __lowercase		:			Optional[int]	    =       UNetaDModel.from_pretrained(
              '''diffusers/consistency-models-test'''				,       subfolder='''test_unet_class_cond'''				,       )
          return unet




      def 		snake_case_			(      self		:       Any				,       _snake_case		:       List[str]=False				):
          if class_cond:
              __lowercase		:			Optional[int]	    =       self.dummy_cond_unet
          else:
              __lowercase		:			Dict	    =       self.dummy_uncond_unet

          # Default to CM multistep sampler
          __lowercase		:			Union[str, Any]	    =       CMStochasticIterativeScheduler(
              num_train_timesteps=40				,       sigma_min=0.0_02				,       sigma_max=80.0				,       )

          __lowercase		:			List[Any]	    =       {
              '''unet''': unet,
              '''scheduler''': scheduler,
          }

          return components




      def 		snake_case_			(      self		:       Any				,       _snake_case		:       Optional[int]				,       _snake_case		:       List[Any]=0				):
          if str(_snake_case				).startswith('''mps'''				):
              __lowercase		:			Optional[Any]	    =       torch.manual_seed(_snake_case				)
          else:
              __lowercase		:			List[Any]	    =       torch.Generator(device=_snake_case				).manual_seed(_snake_case				)

          __lowercase		:			Optional[Any]	    =       {
              '''batch_size''': 1,
              '''num_inference_steps''': None,
              '''timesteps''': [22, 0],
              '''generator''': generator,
              '''output_type''': '''np''',
          }

          return inputs




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

          __lowercase		:			Any	    =       self.get_dummy_inputs(_snake_case				)
          __lowercase		:			Tuple	    =       pipe(**_snake_case				).images
          assert image.shape == (1, 32, 32, 3)

          __lowercase		:			str	    =       image[0, -3:, -3:, -1]
          __lowercase		:			Any	    =       np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]				)

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




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

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

          __lowercase		:			Union[str, Any]	    =       image[0, -3:, -3:, -1]
          __lowercase		:			List[str]	    =       np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]				)

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




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

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

          __lowercase		:			Union[str, Any]	    =       image[0, -3:, -3:, -1]
          __lowercase		:			Optional[int]	    =       np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]				)

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




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

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

          __lowercase		:			Optional[int]	    =       image[0, -3:, -3:, -1]
          __lowercase		:			str	    =       np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]				)

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






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



      """simple docstring"""







      def 		snake_case_			(      self		:       Optional[int]				):
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()




      def 		snake_case_			(      self		:       int				,       _snake_case		:       Optional[Any]=0				,       _snake_case		:       Tuple=False				,       _snake_case		:       Dict="cpu"				,       _snake_case		:       Optional[int]=torch.floataa				,       _snake_case		:       Tuple=(1, 3, 64, 64)				):
          __lowercase		:			Tuple	    =       torch.manual_seed(_snake_case				)

          __lowercase		:			List[str]	    =       {
              '''num_inference_steps''': None,
              '''timesteps''': [22, 0],
              '''class_labels''': 0,
              '''generator''': generator,
              '''output_type''': '''np''',
          }

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

          return inputs




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




      def 		snake_case_			(      self		:       Optional[Any]				):
          __lowercase		:			List[Any]	    =       UNetaDModel.from_pretrained('''diffusers/consistency_models'''				,       subfolder='''diffusers_cd_imagenet64_l2'''				)
          __lowercase		:			Optional[Any]	    =       CMStochasticIterativeScheduler(
              num_train_timesteps=40				,       sigma_min=0.0_02				,       sigma_max=80.0				,       )
          __lowercase		:			Optional[int]	    =       ConsistencyModelPipeline(unet=_snake_case				,       scheduler=_snake_case				)
          pipe.to(torch_device=_snake_case				)
          pipe.set_progress_bar_config(disable=_snake_case				)

          __lowercase		:			Dict	    =       self.get_inputs()
          __lowercase		:			int	    =       pipe(**_snake_case				).images
          assert image.shape == (1, 64, 64, 3)

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

          __lowercase		:			Dict	    =       np.array([0.08_88, 0.08_81, 0.06_66, 0.04_79, 0.02_92, 0.01_95, 0.02_01, 0.01_63, 0.02_54]				)

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




      def 		snake_case_			(      self		:       Dict				):
          __lowercase		:			Union[str, Any]	    =       UNetaDModel.from_pretrained('''diffusers/consistency_models'''				,       subfolder='''diffusers_cd_imagenet64_l2'''				)
          __lowercase		:			Union[str, Any]	    =       CMStochasticIterativeScheduler(
              num_train_timesteps=40				,       sigma_min=0.0_02				,       sigma_max=80.0				,       )
          __lowercase		:			str	    =       ConsistencyModelPipeline(unet=_snake_case				,       scheduler=_snake_case				)
          pipe.to(torch_device=_snake_case				)
          pipe.set_progress_bar_config(disable=_snake_case				)

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

          __lowercase		:			Optional[Any]	    =       image[0, -3:, -3:, -1]

          __lowercase		:			str	    =       np.array([0.03_40, 0.01_52, 0.00_63, 0.02_67, 0.02_21, 0.01_07, 0.04_16, 0.01_86, 0.02_17]				)

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




      @require_torch_a
      def 		snake_case_			(      self		:       Dict				):
          __lowercase		:			int	    =       UNetaDModel.from_pretrained('''diffusers/consistency_models'''				,       subfolder='''diffusers_cd_imagenet64_l2'''				)
          __lowercase		:			Tuple	    =       CMStochasticIterativeScheduler(
              num_train_timesteps=40				,       sigma_min=0.0_02				,       sigma_max=80.0				,       )
          __lowercase		:			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				)

          __lowercase		:			Optional[Any]	    =       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				):
              __lowercase		:			List[str]	    =       pipe(**_snake_case				).images
          assert image.shape == (1, 64, 64, 3)

          __lowercase		:			Union[str, Any]	    =       image[0, -3:, -3:, -1]

          __lowercase		:			Dict	    =       np.array([0.18_75, 0.14_28, 0.12_89, 0.21_51, 0.20_92, 0.14_77, 0.18_77, 0.16_41, 0.13_53]				)

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




      @require_torch_a
      def 		snake_case_			(      self		:       Union[str, Any]				):
          __lowercase		:			Any	    =       UNetaDModel.from_pretrained('''diffusers/consistency_models'''				,       subfolder='''diffusers_cd_imagenet64_l2'''				)
          __lowercase		:			Union[str, Any]	    =       CMStochasticIterativeScheduler(
              num_train_timesteps=40				,       sigma_min=0.0_02				,       sigma_max=80.0				,       )
          __lowercase		:			Any	    =       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				)

          __lowercase		:			Optional[Any]	    =       self.get_inputs(get_fixed_latents=_snake_case				,       device=_snake_case				)
          __lowercase		:			List[Any]	    =       1
          __lowercase		:			Union[str, Any]	    =       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				):
              __lowercase		:			str	    =       pipe(**_snake_case				).images
          assert image.shape == (1, 64, 64, 3)

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

          __lowercase		:			int	    =       np.array([0.16_63, 0.19_48, 0.22_75, 0.16_80, 0.12_04, 0.12_45, 0.18_58, 0.13_38, 0.20_95]				)

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




import unittest

from datasets import load_dataset

from transformers import BloomTokenizerFast
from transformers.testing_utils import require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin


@require_tokenizers
class 						__lowerCAmelCase	(	lowerCAmelCase_ , unittest.TestCase      ):



      """simple docstring"""







      A__	:			Dict        =    None
      A__	:			Any        =    BloomTokenizerFast
      A__	:			Optional[Any]        =    BloomTokenizerFast
      A__	:			List[Any]        =    True
      A__	:			int        =    False
      A__	:			Any        =    '''tokenizer_file'''
      A__	:			Optional[int]        =    {'''bos_token''': '''<s>''', '''eos_token''': '''</s>''', '''unk_token''': '''<unk>''', '''pad_token''': '''<pad>'''}




      def 		snake_case_			(      self		:       Any				):
          super().setUp()
          __lowercase		:			int	    =       BloomTokenizerFast.from_pretrained('''bigscience/tokenizer'''				)
          tokenizer.save_pretrained(self.tmpdirname				)




      def 		snake_case_			(      self		:       Dict				,       **_snake_case		:       Any				):
          kwargs.update(self.special_tokens_map				)
          return BloomTokenizerFast.from_pretrained(self.tmpdirname				,       **_snake_case				)




      def 		snake_case_			(      self		:       Any				):
          __lowercase		:			Dict	    =       self.get_rust_tokenizer()

          __lowercase		:			Optional[int]	    =       ['''The quick brown fox</s>''', '''jumps over the lazy dog</s>''']
          __lowercase		:			Union[str, Any]	    =       [[2175, 2_3714, 7_3173, 14_4252, 2], [77, 13_2619, 3478, 368, 10_9586, 3_5433, 2]]

          __lowercase		:			Optional[int]	    =       tokenizer.batch_encode_plus(_snake_case				)['''input_ids''']
          self.assertListEqual(_snake_case				,       _snake_case				)

          __lowercase		:			List[Any]	    =       tokenizer.batch_decode(_snake_case				)
          self.assertListEqual(_snake_case				,       _snake_case				)




      def 		snake_case_			(      self		:       str				,       _snake_case		:       int=6				):
          for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
              with self.subTest(F'{tokenizer.__class__.__name__} ({pretrained_name})'				):
                  __lowercase		:			Dict	    =       self.rust_tokenizer_class.from_pretrained(_snake_case				,       **_snake_case				)
                  # tokenizer_r.pad_token = None # Hotfixing padding = None
                  # Simple input
                  __lowercase		:			List[Any]	    =       '''This is a simple input'''
                  __lowercase		:			List[str]	    =       ['''This is a simple input 1''', '''This is a simple input 2''']
                  __lowercase		:			Optional[Any]	    =       ('''This is a simple input''', '''This is a pair''')
                  __lowercase		:			Any	    =       [
                      ('''This is a simple input 1''', '''This is a simple input 2'''),
                      ('''This is a simple pair 1''', '''This is a simple pair 2'''),
                  ]

                  # Simple input tests
                  try:
                      tokenizer_r.encode(_snake_case				,       max_length=_snake_case				)
                      tokenizer_r.encode_plus(_snake_case				,       max_length=_snake_case				)

                      tokenizer_r.batch_encode_plus(_snake_case				,       max_length=_snake_case				)
                      tokenizer_r.encode(_snake_case				,       max_length=_snake_case				)
                      tokenizer_r.batch_encode_plus(_snake_case				,       max_length=_snake_case				)
                  except ValueError:
                      self.fail('''Bloom Tokenizer should be able to deal with padding'''				)

                  __lowercase		:			Dict	    =       None  # Hotfixing padding = None
                  self.assertRaises(_snake_case				,       tokenizer_r.encode				,       _snake_case				,       max_length=_snake_case				,       padding='''max_length'''				)

                  # Simple input
                  self.assertRaises(_snake_case				,       tokenizer_r.encode_plus				,       _snake_case				,       max_length=_snake_case				,       padding='''max_length'''				)

                  # Simple input
                  self.assertRaises(
                      _snake_case				,       tokenizer_r.batch_encode_plus				,       _snake_case				,       max_length=_snake_case				,       padding='''max_length'''				,       )

                  # Pair input
                  self.assertRaises(_snake_case				,       tokenizer_r.encode				,       _snake_case				,       max_length=_snake_case				,       padding='''max_length'''				)

                  # Pair input
                  self.assertRaises(_snake_case				,       tokenizer_r.encode_plus				,       _snake_case				,       max_length=_snake_case				,       padding='''max_length'''				)

                  # Pair input
                  self.assertRaises(
                      _snake_case				,       tokenizer_r.batch_encode_plus				,       _snake_case				,       max_length=_snake_case				,       padding='''max_length'''				,       )




      def 		snake_case_			(      self		:       Optional[int]				):
          __lowercase		:			List[Any]	    =       self.get_rust_tokenizer()
          __lowercase		:			int	    =       load_dataset('''xnli'''				,       '''all_languages'''				,       split='''test'''				,       streaming=_snake_case				)

          __lowercase		:			Optional[Any]	    =       next(iter(_snake_case				)				)['''premise''']  # pick up one data
          __lowercase		:			List[str]	    =       list(sample_data.values()				)

          __lowercase		:			List[str]	    =       list(map(tokenizer.encode				,       _snake_case				)				)
          __lowercase		:			Tuple	    =       [tokenizer.decode(_snake_case				,       clean_up_tokenization_spaces=_snake_case				) for x in output_tokens]
          self.assertListEqual(_snake_case				,       _snake_case				)




      def 		snake_case_			(      self		:       List[Any]				):
          # The test has to be overriden because BLOOM uses ALiBi positional embeddings that does not have
          # any sequence length constraints. This test of the parent class will fail since it relies on the
          # maximum sequence length of the positoonal embeddings.
          self.assertGreaterEqual(len(self.tokenizer_class.pretrained_vocab_files_map				)				,       1				)
          self.assertGreaterEqual(len(list(self.tokenizer_class.pretrained_vocab_files_map.values()				)[0]				)				,       1				)


