


"""simple docstring"""


import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers.testing_utils import require_vision
from transformers.utils import is_vision_available


if is_vision_available():
	from PIL import Image

	from transformers import AutoProcessor, BlipaProcessor, BlipImageProcessor, GPTaTokenizer, PreTrainedTokenizerFast
@require_vision
class   _UpperCamelCase  (   unittest.TestCase		):
							'''simple docstring'''







							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			tempfile.mkdtemp()

										__lowerCAmelCase				    =			BlipImageProcessor()
										__lowerCAmelCase				    =			GPTaTokenizer.from_pretrained("hf-internal-testing/tiny-random-GPT2Model"	)

										__lowerCAmelCase				    =			BlipaProcessor(__a    ,   __a	)

										processor.save_pretrained(self.tmpdirname	)

							def 			snake_case				(     self    ,   **__a	):
										return AutoProcessor.from_pretrained(self.tmpdirname    ,   **__a	).tokenizer

							def 			snake_case				(     self    ,   **__a	):
										return AutoProcessor.from_pretrained(self.tmpdirname    ,   **__a	).image_processor

							def 			snake_case				(     self	):
										shutil.rmtree(self.tmpdirname	)

							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			[np.random.randint(2_55    ,   size=(3, 30, 4_00)    ,   dtype=np.uinta	)]

										__lowerCAmelCase				    =			[Image.fromarray(np.moveaxis(__a    ,   0    ,   -1	)	) for x in image_inputs]

										return image_inputs

							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			BlipaProcessor(tokenizer=self.get_tokenizer()    ,   image_processor=self.get_image_processor()	)
										processor.save_pretrained(self.tmpdirname	)

										__lowerCAmelCase				    =			self.get_tokenizer(bos_token="(BOS)"    ,   eos_token="(EOS)"	)
										__lowerCAmelCase				    =			self.get_image_processor(do_normalize=__a    ,   padding_value=1.0	)

										__lowerCAmelCase				    =			BlipaProcessor.from_pretrained(
										    self.tmpdirname    ,   bos_token="(BOS)"    ,   eos_token="(EOS)"    ,   do_normalize=__a    ,   padding_value=1.0	)

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

										self.assertEqual(processor.image_processor.to_json_string()    ,   image_processor_add_kwargs.to_json_string()	)
										self.assertIsInstance(processor.image_processor    ,   __a	)

							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			self.get_image_processor()
										__lowerCAmelCase				    =			self.get_tokenizer()

										__lowerCAmelCase				    =			BlipaProcessor(tokenizer=__a    ,   image_processor=__a	)

										__lowerCAmelCase				    =			self.prepare_image_inputs()

										__lowerCAmelCase				    =			image_processor(__a    ,   return_tensors="np"	)
										__lowerCAmelCase				    =			processor(images=__a    ,   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	):
										__lowerCAmelCase				    =			self.get_image_processor()
										__lowerCAmelCase				    =			self.get_tokenizer()

										__lowerCAmelCase				    =			BlipaProcessor(tokenizer=__a    ,   image_processor=__a	)

										__lowerCAmelCase				    =			"lower newer"

										__lowerCAmelCase				    =			processor(text=__a	)

										__lowerCAmelCase				    =			tokenizer(__a    ,   return_token_type_ids=__a	)

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

							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			self.get_image_processor()
										__lowerCAmelCase				    =			self.get_tokenizer()

										__lowerCAmelCase				    =			BlipaProcessor(tokenizer=__a    ,   image_processor=__a	)

										__lowerCAmelCase				    =			"lower newer"
										__lowerCAmelCase				    =			self.prepare_image_inputs()

										__lowerCAmelCase				    =			processor(text=__a    ,   images=__a	)

										self.assertListEqual(list(inputs.keys()	)    ,   ["pixel_values", "input_ids", "attention_mask"]	)

										# test if it raises when no input is passed
										with pytest.raises(__a	):
													processor()

							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			self.get_image_processor()
										__lowerCAmelCase				    =			self.get_tokenizer()

										__lowerCAmelCase				    =			BlipaProcessor(tokenizer=__a    ,   image_processor=__a	)

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

										__lowerCAmelCase				    =			processor.batch_decode(__a	)
										__lowerCAmelCase				    =			tokenizer.batch_decode(__a	)

										self.assertListEqual(__a    ,   __a	)



							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			self.get_image_processor()
										__lowerCAmelCase				    =			self.get_tokenizer()

										__lowerCAmelCase				    =			BlipaProcessor(tokenizer=__a    ,   image_processor=__a	)

										__lowerCAmelCase				    =			"lower newer"
										__lowerCAmelCase				    =			self.prepare_image_inputs()

										__lowerCAmelCase				    =			processor(text=__a    ,   images=__a	)

										# For now the processor supports only ['pixel_values', 'input_ids', 'attention_mask']
										self.assertListEqual(list(inputs.keys()	)    ,   ["pixel_values", "input_ids", "attention_mask"]	)







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      (_UpperCAmelCase		,   unittest.TestCase ):
 _UpperCamelCase			  =  ConsistencyModelPipeline
 _UpperCamelCase			  =  UNCONDITIONAL_IMAGE_GENERATION_PARAMS
 _UpperCamelCase			  =  UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

 # Override required_optional_params to remove num_images_per_prompt
 _UpperCamelCase			  =  frozenset(
     [
         """num_inference_steps""",
         """generator""",
         """latents""",
         """output_type""",
         """return_dict""",
         """callback""",
         """callback_steps""",
     ] )
 @property
 def 	UpperCamelCase__  (							self					)					->List[str]:



       '''simple docstring'''





       __lowerCAmelCase					:						List[Any]					=					UNetaDModel.from_pretrained(
           '''diffusers/consistency-models-test'''      ,							subfolder='''test_unet'''      ,							)
       return unet
 @property
 def 	UpperCamelCase__  (							self					)					->Optional[int]:



       '''simple docstring'''





       __lowerCAmelCase					:						List[str]					=					UNetaDModel.from_pretrained(
           '''diffusers/consistency-models-test'''      ,							subfolder='''test_unet_class_cond'''      ,							)
       return unet
 def 	UpperCamelCase__  (							self      ,							A_=False					)					->Dict:



       '''simple docstring'''





       if class_cond:
             __lowerCAmelCase					:						List[str]					=					self.dummy_cond_unet
       else:
             __lowerCAmelCase					:						Optional[Any]					=					self.dummy_uncond_unet

       # Default to CM multistep sampler
       __lowerCAmelCase					:						List[str]					=					CMStochasticIterativeScheduler(
           num_train_timesteps=40      ,							sigma_min=0.002      ,							sigma_max=80.0      ,							)

       __lowerCAmelCase					:						Dict					=					{
           '''unet''': unet,
           '''scheduler''': scheduler,
       }

       return components
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_=0					)					->Tuple:



       '''simple docstring'''





       if str(A_					).startswith('''mps'''					):
             __lowerCAmelCase					:						str					=					torch.manual_seed(A_					)
       else:
             __lowerCAmelCase					:						Dict					=					torch.Generator(device=A_					).manual_seed(A_					)

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

       return inputs
 def 	UpperCamelCase__  (							self					)					->Tuple:



       '''simple docstring'''





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

       __lowerCAmelCase					:						Any					=					self.get_dummy_inputs(A_					)
       __lowerCAmelCase					:						int					=					pipe(**A_					).images
       assert image.shape == (1, 32, 32, 3)

       __lowerCAmelCase					:						Tuple					=					image[0, -3:, -3:, -1]
       __lowerCAmelCase					:						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 	UpperCamelCase__  (							self					)					->Tuple:



       '''simple docstring'''





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

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

       __lowerCAmelCase					:						Optional[int]					=					image[0, -3:, -3:, -1]
       __lowerCAmelCase					:						List[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 	UpperCamelCase__  (							self					)					->Optional[Any]:



       '''simple docstring'''





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

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

       __lowerCAmelCase					:						Any					=					image[0, -3:, -3:, -1]
       __lowerCAmelCase					:						List[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 	UpperCamelCase__  (							self					)					->Optional[Any]:



       '''simple docstring'''





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

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

       __lowerCAmelCase					:						List[Any]					=					image[0, -3:, -3:, -1]
       __lowerCAmelCase					:						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







@slow
@require_torch_gpu
class   __lowercase      (unittest.TestCase ):
 def 	UpperCamelCase__  (							self					)					->int:



       '''simple docstring'''





       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
 def 	UpperCamelCase__  (							self      ,							A_=0      ,							A_=False      ,							A_="cpu"      ,							A_=torch.floataa      ,							A_=(1, 3, 64, 64)					)					->str:



       '''simple docstring'''





       __lowerCAmelCase					:						Dict					=					torch.manual_seed(A_					)

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

       if get_fixed_latents:
             __lowerCAmelCase					:						List[str]					=					self.get_fixed_latents(seed=A_      ,							device=A_      ,							dtype=A_      ,							shape=A_					)
             __lowerCAmelCase					:						Union[str, Any]					=					latents

       return inputs
 def 	UpperCamelCase__  (							self      ,							A_=0      ,							A_="cpu"      ,							A_=torch.floataa      ,							A_=(1, 3, 64, 64)					)					->Optional[int]:



       '''simple docstring'''





       if type(A_					) == str:
             __lowerCAmelCase					:						int					=					torch.device(A_					)
       __lowerCAmelCase					:						Optional[Any]					=					torch.Generator(device=A_					).manual_seed(A_					)
       __lowerCAmelCase					:						Union[str, Any]					=					randn_tensor(A_      ,							generator=A_      ,							device=A_      ,							dtype=A_					)
       return latents
 def 	UpperCamelCase__  (							self					)					->Optional[Any]:



       '''simple docstring'''





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

       __lowerCAmelCase					:						str					=					self.get_inputs()
       __lowerCAmelCase					:						Any					=					pipe(**A_					).images
       assert image.shape == (1, 64, 64, 3)

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

       __lowerCAmelCase					:						Optional[int]					=					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 	UpperCamelCase__  (							self					)					->Any:



       '''simple docstring'''





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

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

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

       __lowerCAmelCase					:						List[Any]					=					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 	UpperCamelCase__  (							self					)					->Optional[Any]:



       '''simple docstring'''





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

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

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

       __lowerCAmelCase					:						Optional[int]					=					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 	UpperCamelCase__  (							self					)					->List[str]:



       '''simple docstring'''





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

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

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

       __lowerCAmelCase					:						Any					=					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


