





import gc
import os
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
from diffusers.models.attention_processor import (
    Attention,
    AttnProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)
from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device



def   __lowercase  ( _SCREAMING_SNAKE_CASE      )					->  List[str]:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE						=		{}
      for name in unet.attn_processors.keys():
            SCREAMING_SNAKE_CASE						=		None if name.endswith("""attn1.processor"""      ) else unet.config.cross_attention_dim
            if name.startswith("""mid_block"""      ):
                  SCREAMING_SNAKE_CASE						=		unet.config.block_out_channels[-1]
            elif name.startswith("""up_blocks"""      ):
                  SCREAMING_SNAKE_CASE						=		int(name[len("""up_blocks."""      )]      )
                  SCREAMING_SNAKE_CASE						=		list(reversed(unet.config.block_out_channels      )      )[block_id]
            elif name.startswith("""down_blocks"""      ):
                  SCREAMING_SNAKE_CASE						=		int(name[len("""down_blocks."""      )]      )
                  SCREAMING_SNAKE_CASE						=		unet.config.block_out_channels[block_id]
            SCREAMING_SNAKE_CASE						=		(
                LoRAAttnProcessor2_0 if hasattr(_UpperCAmelCase					,		"""scaled_dot_product_attention"""      ) else LoRAAttnProcessor
            )
            SCREAMING_SNAKE_CASE						=		lora_attn_processor_class(
                hidden_size=_UpperCAmelCase					,		cross_attention_dim=_UpperCAmelCase      )
      SCREAMING_SNAKE_CASE						=		AttnProcsLayers(_UpperCAmelCase      )
      return lora_attn_procs, unet_lora_layers



def   __lowercase  ( _SCREAMING_SNAKE_CASE      )					->  Tuple:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE						=		{}
      SCREAMING_SNAKE_CASE						=		(
          LoRAAttnProcessor2_0 if hasattr(_UpperCAmelCase					,		"""scaled_dot_product_attention"""      ) else LoRAAttnProcessor
      )
      for name, module in text_encoder.named_modules():
            if name.endswith(_UpperCAmelCase      ):
                  SCREAMING_SNAKE_CASE						=		lora_attn_processor_class(
                      hidden_size=module.out_proj.out_features					,		cross_attention_dim=_UpperCAmelCase      )
      return text_lora_attn_procs



def   __lowercase  ( _SCREAMING_SNAKE_CASE      )					->  Tuple:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE						=		create_text_encoder_lora_attn_procs(_UpperCAmelCase      )
      SCREAMING_SNAKE_CASE						=		AttnProcsLayers(_UpperCAmelCase      )
      return text_encoder_lora_layers



def   __lowercase  ( _SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE=False      )					->  Optional[Any]:


      '''simple docstring'''


      for _, attn_proc in text_lora_attn_procs.items():
            # set up.weights
            for layer_name, layer_module in attn_proc.named_modules():
                  if layer_name.endswith("""_lora"""      ):
                        SCREAMING_SNAKE_CASE						=		(
                            torch.randn_like(layer_module.up.weight      )
                            if randn_weight
                            else torch.zeros_like(layer_module.up.weight      )
                        )
                        SCREAMING_SNAKE_CASE						=		torch.nn.Parameter(_UpperCAmelCase      )



class 						UpperCamelCase__					( unittest.TestCase  ):



     '''simple docstring'''

     def 							SCREAMING_SNAKE_CASE__    (				self		:       List[Any]			)	->						Dict:




           '''simple docstring'''


           torch.manual_seed(0			)
           SCREAMING_SNAKE_CASE						=		UNetaDConditionModel(
               block_out_channels=(32, 64)		,layers_per_block=2		,sample_size=32		,in_channels=4		,out_channels=4		,down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")		,up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")		,cross_attention_dim=32		,)
           SCREAMING_SNAKE_CASE						=		DDIMScheduler(
               beta_start=0.00085		,beta_end=0.012		,beta_schedule="""scaled_linear"""		,clip_sample=_lowerCamelCase		,set_alpha_to_one=_lowerCamelCase		,steps_offset=1		,)
           torch.manual_seed(0			)
           SCREAMING_SNAKE_CASE						=		AutoencoderKL(
               block_out_channels=[32, 64]		,in_channels=3		,out_channels=3		,down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]		,up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]		,latent_channels=4		,)
           SCREAMING_SNAKE_CASE						=		CLIPTextConfig(
               bos_token_id=0		,eos_token_id=2		,hidden_size=32		,intermediate_size=37		,layer_norm_eps=1e-0_5		,num_attention_heads=4		,num_hidden_layers=5		,pad_token_id=1		,vocab_size=1000		,)
           SCREAMING_SNAKE_CASE						=		CLIPTextModel(_lowerCamelCase			)
           SCREAMING_SNAKE_CASE						=		CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)

           SCREAMING_SNAKE_CASE						=		create_unet_lora_layers(_lowerCamelCase			)
           SCREAMING_SNAKE_CASE						=		create_text_encoder_lora_layers(_lowerCamelCase			)

           SCREAMING_SNAKE_CASE						=		{
               '''unet''': unet,
               '''scheduler''': scheduler,
               '''vae''': vae,
               '''text_encoder''': text_encoder,
               '''tokenizer''': tokenizer,
               '''safety_checker''': None,
               '''feature_extractor''': None,
           }
           SCREAMING_SNAKE_CASE						=		{
               '''unet_lora_layers''': unet_lora_layers,
               '''text_encoder_lora_layers''': text_encoder_lora_layers,
               '''unet_lora_attn_procs''': unet_lora_attn_procs,
           }
           return pipeline_components, lora_components






     def 							SCREAMING_SNAKE_CASE__    (				self		:       str			)	->						int:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		1
           SCREAMING_SNAKE_CASE						=		10
           SCREAMING_SNAKE_CASE						=		4
           SCREAMING_SNAKE_CASE						=		(32, 32)

           SCREAMING_SNAKE_CASE						=		torch.manual_seed(0			)
           SCREAMING_SNAKE_CASE						=		floats_tensor((batch_size, num_channels) + sizes			)
           SCREAMING_SNAKE_CASE						=		torch.randint(1		,_lowerCamelCase		,size=(batch_size, sequence_length)		,generator=_lowerCamelCase			)

           SCREAMING_SNAKE_CASE						=		{
               '''prompt''': '''A painting of a squirrel eating a burger''',
               '''generator''': generator,
               '''num_inference_steps''': 2,
               '''guidance_scale''': 6.0,
               '''output_type''': '''numpy''',
           }

           return noise, input_ids, pipeline_inputs

           # copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb






     def 							SCREAMING_SNAKE_CASE__    (				self		:       str			)	->						Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		77

           SCREAMING_SNAKE_CASE						=		torch.randint(2		,56		,size=(1, max_seq_length)		,generator=torch.manual_seed(0			)			)

           SCREAMING_SNAKE_CASE						=		{}
           SCREAMING_SNAKE_CASE						=		inputs
           return prepared_inputs






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Optional[int]		,lowerCamelCase__		:       List[Any]			)	->						Tuple:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		self.get_dummy_components()
           LoraLoaderMixin.save_lora_weights(
               save_directory=_lowerCamelCase		,unet_lora_layers=lora_components["""unet_lora_layers"""]		,text_encoder_lora_layers=lora_components["""text_encoder_lora_layers"""]		,)
           self.assertTrue(os.path.isfile(os.path.join(_lowerCamelCase		,"""pytorch_lora_weights.bin"""			)			)			)






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Any			)	->						Optional[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		self.get_dummy_components()
           SCREAMING_SNAKE_CASE						=		StableDiffusionPipeline(**_lowerCamelCase			)
           SCREAMING_SNAKE_CASE						=		sd_pipe.to(_lowerCamelCase			)
           sd_pipe.set_progress_bar_config(disable=_lowerCamelCase			)

           SCREAMING_SNAKE_CASE						=		self.get_dummy_inputs()

           SCREAMING_SNAKE_CASE						=		sd_pipe(**_lowerCamelCase			).images
           SCREAMING_SNAKE_CASE						=		original_images[0, -3:, -3:, -1]

           with tempfile.TemporaryDirectory() as tmpdirname:
                 LoraLoaderMixin.save_lora_weights(
                     save_directory=_lowerCamelCase		,unet_lora_layers=lora_components["""unet_lora_layers"""]		,text_encoder_lora_layers=lora_components["""text_encoder_lora_layers"""]		,)
                 self.assertTrue(os.path.isfile(os.path.join(_lowerCamelCase		,"""pytorch_lora_weights.bin"""			)			)			)
                 sd_pipe.load_lora_weights(_lowerCamelCase			)

           SCREAMING_SNAKE_CASE						=		sd_pipe(**_lowerCamelCase			).images
           SCREAMING_SNAKE_CASE						=		lora_images[0, -3:, -3:, -1]

           # Outputs shouldn't match.
           self.assertFalse(torch.allclose(torch.from_numpy(_lowerCamelCase			)		,torch.from_numpy(_lowerCamelCase			)			)			)






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Optional[int]			)	->						List[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		self.get_dummy_components()
           SCREAMING_SNAKE_CASE						=		StableDiffusionPipeline(**_lowerCamelCase			)
           SCREAMING_SNAKE_CASE						=		sd_pipe.to(_lowerCamelCase			)
           sd_pipe.set_progress_bar_config(disable=_lowerCamelCase			)

           SCREAMING_SNAKE_CASE						=		self.get_dummy_inputs()

           SCREAMING_SNAKE_CASE						=		sd_pipe(**_lowerCamelCase			).images
           SCREAMING_SNAKE_CASE						=		original_images[0, -3:, -3:, -1]

           with tempfile.TemporaryDirectory() as tmpdirname:
                 LoraLoaderMixin.save_lora_weights(
                     save_directory=_lowerCamelCase		,unet_lora_layers=lora_components["""unet_lora_layers"""]		,text_encoder_lora_layers=lora_components["""text_encoder_lora_layers"""]		,safe_serialization=_lowerCamelCase		,)
                 self.assertTrue(os.path.isfile(os.path.join(_lowerCamelCase		,"""pytorch_lora_weights.safetensors"""			)			)			)
                 sd_pipe.load_lora_weights(_lowerCamelCase			)

           SCREAMING_SNAKE_CASE						=		sd_pipe(**_lowerCamelCase			).images
           SCREAMING_SNAKE_CASE						=		lora_images[0, -3:, -3:, -1]

           # Outputs shouldn't match.
           self.assertFalse(torch.allclose(torch.from_numpy(_lowerCamelCase			)		,torch.from_numpy(_lowerCamelCase			)			)			)






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Optional[int]			)	->						Union[str, Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		self.get_dummy_components()
           SCREAMING_SNAKE_CASE						=		lora_components['''unet_lora_attn_procs''']
           SCREAMING_SNAKE_CASE						=		StableDiffusionPipeline(**_lowerCamelCase			)
           SCREAMING_SNAKE_CASE						=		sd_pipe.to(_lowerCamelCase			)
           sd_pipe.set_progress_bar_config(disable=_lowerCamelCase			)

           SCREAMING_SNAKE_CASE						=		self.get_dummy_inputs()

           SCREAMING_SNAKE_CASE						=		sd_pipe(**_lowerCamelCase			).images
           SCREAMING_SNAKE_CASE						=		original_images[0, -3:, -3:, -1]

           with tempfile.TemporaryDirectory() as tmpdirname:
                 SCREAMING_SNAKE_CASE						=		sd_pipe.unet
                 unet.set_attn_processor(_lowerCamelCase			)
                 unet.save_attn_procs(_lowerCamelCase			)
                 self.assertTrue(os.path.isfile(os.path.join(_lowerCamelCase		,"""pytorch_lora_weights.bin"""			)			)			)
                 sd_pipe.load_lora_weights(_lowerCamelCase			)

           SCREAMING_SNAKE_CASE						=		sd_pipe(**_lowerCamelCase			).images
           SCREAMING_SNAKE_CASE						=		lora_images[0, -3:, -3:, -1]

           # Outputs shouldn't match.
           self.assertFalse(torch.allclose(torch.from_numpy(_lowerCamelCase			)		,torch.from_numpy(_lowerCamelCase			)			)			)






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Tuple			)	->						str:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		self.get_dummy_components()
           SCREAMING_SNAKE_CASE						=		StableDiffusionPipeline(**_lowerCamelCase			)

           SCREAMING_SNAKE_CASE						=		self.get_dummy_tokens()

           # inference without lora
           SCREAMING_SNAKE_CASE						=		pipe.text_encoder(**_lowerCamelCase			)[0]
           assert outputs_without_lora.shape == (1, 77, 32)

           # create lora_attn_procs with zeroed out up.weights
           SCREAMING_SNAKE_CASE						=		create_text_encoder_lora_attn_procs(pipe.text_encoder			)
           set_lora_up_weights(_lowerCamelCase		,randn_weight=_lowerCamelCase			)

           # monkey patch
           pipe._modify_text_encoder(_lowerCamelCase			)

           # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
           del text_attn_procs
           gc.collect()

           # inference with lora
           SCREAMING_SNAKE_CASE						=		pipe.text_encoder(**_lowerCamelCase			)[0]
           assert outputs_with_lora.shape == (1, 77, 32)

           assert torch.allclose(
               _lowerCamelCase		,_lowerCamelCase			), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

           # create lora_attn_procs with randn up.weights
           SCREAMING_SNAKE_CASE						=		create_text_encoder_lora_attn_procs(pipe.text_encoder			)
           set_lora_up_weights(_lowerCamelCase		,randn_weight=_lowerCamelCase			)

           # monkey patch
           pipe._modify_text_encoder(_lowerCamelCase			)

           # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
           del text_attn_procs
           gc.collect()

           # inference with lora
           SCREAMING_SNAKE_CASE						=		pipe.text_encoder(**_lowerCamelCase			)[0]
           assert outputs_with_lora.shape == (1, 77, 32)

           assert not torch.allclose(
               _lowerCamelCase		,_lowerCamelCase			), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Optional[int]			)	->						Any:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		self.get_dummy_components()
           SCREAMING_SNAKE_CASE						=		StableDiffusionPipeline(**_lowerCamelCase			)

           SCREAMING_SNAKE_CASE						=		self.get_dummy_tokens()

           # inference without lora
           SCREAMING_SNAKE_CASE						=		pipe.text_encoder(**_lowerCamelCase			)[0]
           assert outputs_without_lora.shape == (1, 77, 32)

           # create lora_attn_procs with randn up.weights
           SCREAMING_SNAKE_CASE						=		create_text_encoder_lora_attn_procs(pipe.text_encoder			)
           set_lora_up_weights(_lowerCamelCase		,randn_weight=_lowerCamelCase			)

           # monkey patch
           pipe._modify_text_encoder(_lowerCamelCase			)

           # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
           del text_attn_procs
           gc.collect()

           # inference with lora
           SCREAMING_SNAKE_CASE						=		pipe.text_encoder(**_lowerCamelCase			)[0]
           assert outputs_with_lora.shape == (1, 77, 32)

           assert not torch.allclose(
               _lowerCamelCase		,_lowerCamelCase			), "lora outputs should be different to without lora outputs"

           # remove monkey patch
           pipe._remove_text_encoder_monkey_patch()

           # inference with removed lora
           SCREAMING_SNAKE_CASE						=		pipe.text_encoder(**_lowerCamelCase			)[0]
           assert outputs_without_lora_removed.shape == (1, 77, 32)

           assert torch.allclose(
               _lowerCamelCase		,_lowerCamelCase			), "remove lora monkey patch should restore the original outputs"






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Tuple			)	->						Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		self.get_dummy_components()
           SCREAMING_SNAKE_CASE						=		StableDiffusionPipeline(**_lowerCamelCase			)
           SCREAMING_SNAKE_CASE						=		sd_pipe.to(_lowerCamelCase			)
           sd_pipe.set_progress_bar_config(disable=_lowerCamelCase			)

           SCREAMING_SNAKE_CASE						=		self.get_dummy_inputs()

           with tempfile.TemporaryDirectory() as tmpdirname:
                 LoraLoaderMixin.save_lora_weights(
                     save_directory=_lowerCamelCase		,unet_lora_layers=lora_components["""unet_lora_layers"""]		,text_encoder_lora_layers=lora_components["""text_encoder_lora_layers"""]		,)
                 self.assertTrue(os.path.isfile(os.path.join(_lowerCamelCase		,"""pytorch_lora_weights.bin"""			)			)			)
                 sd_pipe.load_lora_weights(_lowerCamelCase			)

           SCREAMING_SNAKE_CASE						=		sd_pipe(**_lowerCamelCase			).images
           SCREAMING_SNAKE_CASE						=		lora_images[0, -3:, -3:, -1]

           SCREAMING_SNAKE_CASE						=		sd_pipe(**_lowerCamelCase		,cross_attention_kwargs={"""scale""": 0.5}			).images
           SCREAMING_SNAKE_CASE						=		lora_images_with_scale[0, -3:, -3:, -1]

           # Outputs shouldn't match.
           self.assertFalse(
               torch.allclose(torch.from_numpy(_lowerCamelCase			)		,torch.from_numpy(_lowerCamelCase			)			)			)






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Tuple			)	->						Tuple:




           '''simple docstring'''


           with tempfile.TemporaryDirectory() as tmpdirname:
                 self.create_lora_weight_file(_lowerCamelCase			)

                 SCREAMING_SNAKE_CASE						=		self.get_dummy_components()
                 SCREAMING_SNAKE_CASE						=		StableDiffusionPipeline(**_lowerCamelCase			)
                 SCREAMING_SNAKE_CASE						=		sd_pipe.to(_lowerCamelCase			)
                 sd_pipe.set_progress_bar_config(disable=_lowerCamelCase			)

                 # check if vanilla attention processors are used
                 for _, module in sd_pipe.unet.named_modules():
                       if isinstance(_lowerCamelCase		,_lowerCamelCase			):
                             self.assertIsInstance(module.processor		,(AttnProcessor, AttnProcessor2_0)			)

            # load LoRA weight file
                 sd_pipe.load_lora_weights(_lowerCamelCase			)

                 # check if lora attention processors are used
                 for _, module in sd_pipe.unet.named_modules():
                       if isinstance(_lowerCamelCase		,_lowerCamelCase			):
                             SCREAMING_SNAKE_CASE						=		(
                                 LoRAAttnProcessor2_0 if hasattr(_lowerCamelCase		,"""scaled_dot_product_attention"""			) else LoRAAttnProcessor
                             )
                             self.assertIsInstance(module.processor		,_lowerCamelCase			)






     @unittest.skipIf(torch_device != """cuda"""		,"""This test is supposed to run on GPU"""			)
     def 							SCREAMING_SNAKE_CASE__    (				self		:       Optional[Any]			)	->						Tuple:




           '''simple docstring'''


           with tempfile.TemporaryDirectory() as tmpdirname:
                 self.create_lora_weight_file(_lowerCamelCase			)

                 SCREAMING_SNAKE_CASE						=		self.get_dummy_components()
                 SCREAMING_SNAKE_CASE						=		StableDiffusionPipeline(**_lowerCamelCase			)
                 SCREAMING_SNAKE_CASE						=		sd_pipe.to(_lowerCamelCase			)
                 sd_pipe.set_progress_bar_config(disable=_lowerCamelCase			)

                 # enable XFormers
                 sd_pipe.enable_xformers_memory_efficient_attention()

                 # check if xFormers attention processors are used
                 for _, module in sd_pipe.unet.named_modules():
                       if isinstance(_lowerCamelCase		,_lowerCamelCase			):
                             self.assertIsInstance(module.processor		,_lowerCamelCase			)

            # load LoRA weight file
                 sd_pipe.load_lora_weights(_lowerCamelCase			)

                 # check if lora attention processors are used
                 for _, module in sd_pipe.unet.named_modules():
                       if isinstance(_lowerCamelCase		,_lowerCamelCase			):
                             self.assertIsInstance(module.processor		,_lowerCamelCase			)






     @unittest.skipIf(torch_device != """cuda"""		,"""This test is supposed to run on GPU"""			)
     def 							SCREAMING_SNAKE_CASE__    (				self		:       Optional[Any]			)	->						int:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		self.get_dummy_components()
           SCREAMING_SNAKE_CASE						=		StableDiffusionPipeline(**_lowerCamelCase			)
           SCREAMING_SNAKE_CASE						=		sd_pipe.to(_lowerCamelCase			)
           sd_pipe.set_progress_bar_config(disable=_lowerCamelCase			)

           SCREAMING_SNAKE_CASE						=		self.get_dummy_inputs()

           # enable XFormers
           sd_pipe.enable_xformers_memory_efficient_attention()

           SCREAMING_SNAKE_CASE						=		sd_pipe(**_lowerCamelCase			).images
           SCREAMING_SNAKE_CASE						=		original_images[0, -3:, -3:, -1]

           with tempfile.TemporaryDirectory() as tmpdirname:
                 LoraLoaderMixin.save_lora_weights(
                     save_directory=_lowerCamelCase		,unet_lora_layers=lora_components["""unet_lora_layers"""]		,text_encoder_lora_layers=lora_components["""text_encoder_lora_layers"""]		,)
                 self.assertTrue(os.path.isfile(os.path.join(_lowerCamelCase		,"""pytorch_lora_weights.bin"""			)			)			)
                 sd_pipe.load_lora_weights(_lowerCamelCase			)

           SCREAMING_SNAKE_CASE						=		sd_pipe(**_lowerCamelCase			).images
           SCREAMING_SNAKE_CASE						=		lora_images[0, -3:, -3:, -1]

           # Outputs shouldn't match.
           self.assertFalse(torch.allclose(torch.from_numpy(_lowerCamelCase			)		,torch.from_numpy(_lowerCamelCase			)			)			)








"""simple docstring"""

import unittest

from transformers import (
    MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
    TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
    PreTrainedTokenizer,
    is_vision_available,
)
from transformers.pipelines import ImageClassificationPipeline, pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_torch_or_tf,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
		from PIL import Image
else:

		class 							lowercase   :







						@staticmethod
						def        a_	(				*_lowerCamelCase :       Tuple   ,  **_lowerCamelCase :       str							):



							"""simple docstring"""

							pass

@is_pipeline_test
@require_torch_or_tf
@require_vision
class 							lowercase   (     unittest.TestCase):
				__lowerCAmelCase     :     str								= MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
				__lowerCAmelCase     :     Union[str, Any]								= TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING







				def        a_	(				self :       Tuple   ,  _lowerCamelCase :       Any   ,  _lowerCamelCase :       str   ,  _lowerCamelCase :       Optional[Any]							):



					"""simple docstring"""

					A_	:					List[str]						 =		ImageClassificationPipeline(model=_lowerCamelCase   ,  image_processor=_lowerCamelCase   ,  top_k=2							)
					A_	:					Tuple						 =		[
					    Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							),
					    '''http://images.cocodataset.org/val2017/000000039769.jpg''',
					]
					return image_classifier, examples







				def        a_	(				self :       Optional[Any]   ,  _lowerCamelCase :       Optional[int]   ,  _lowerCamelCase :       Optional[Any]							):



					"""simple docstring"""

					A_	:					int						 =		image_classifier('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							)

					self.assertEqual(
					    _lowerCamelCase   ,  [
					        {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					        {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					    ]   ,  )

					import datasets

					A_	:					Optional[Any]						 =		datasets.load_dataset('''hf-internal-testing/fixtures_image_utils'''   ,  '''image'''   ,  split='''test'''							)

					# Accepts URL + PIL.Image + lists
					A_	:					Union[str, Any]						 =		image_classifier(
					    [
					        Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							),
					        '''http://images.cocodataset.org/val2017/000000039769.jpg''',
					        # RGBA
					        dataset[0]['''file'''],
					        # LA
					        dataset[1]['''file'''],
					        # L
					        dataset[2]['''file'''],
					    ]							)
					self.assertEqual(
					    _lowerCamelCase   ,  [
					        [
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					        ],
					        [
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					        ],
					        [
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					        ],
					        [
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					        ],
					        [
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					            {'''score''': ANY(_lowerCamelCase							), '''label''': ANY(_lowerCamelCase							)},
					        ],
					    ]   ,  )







				@require_torch
				def        a_	(				self :       Dict							):



					"""simple docstring"""

					A_	:					Optional[int]						 =		'''hf-internal-testing/tiny-random-vit'''
					A_	:					Tuple						 =		pipeline('''image-classification'''   ,  model=_lowerCamelCase							)

					A_	:					Optional[int]						 =		image_classifier('''http://images.cocodataset.org/val2017/000000039769.jpg'''							)
					self.assertEqual(
					    nested_simplify(_lowerCamelCase   ,  decimals=4							)   ,  [{'''label''': '''LABEL_1''', '''score''': 0.574}, {'''label''': '''LABEL_0''', '''score''': 0.426}]   ,  )

					A_	:					str						 =		image_classifier(
					    [
					        '''http://images.cocodataset.org/val2017/000000039769.jpg''',
					        '''http://images.cocodataset.org/val2017/000000039769.jpg''',
					    ]   ,  top_k=2   ,  )
					self.assertEqual(
					    nested_simplify(_lowerCamelCase   ,  decimals=4							)   ,  [
					        [{'''label''': '''LABEL_1''', '''score''': 0.574}, {'''label''': '''LABEL_0''', '''score''': 0.426}],
					        [{'''label''': '''LABEL_1''', '''score''': 0.574}, {'''label''': '''LABEL_0''', '''score''': 0.426}],
					    ]   ,  )







				@require_tf
				def        a_	(				self :       Tuple							):



					"""simple docstring"""

					A_	:					List[str]						 =		'''hf-internal-testing/tiny-random-vit'''
					A_	:					Optional[int]						 =		pipeline('''image-classification'''   ,  model=_lowerCamelCase   ,  framework='''tf'''							)

					A_	:					Tuple						 =		image_classifier('''http://images.cocodataset.org/val2017/000000039769.jpg'''							)
					self.assertEqual(
					    nested_simplify(_lowerCamelCase   ,  decimals=4							)   ,  [{'''label''': '''LABEL_1''', '''score''': 0.574}, {'''label''': '''LABEL_0''', '''score''': 0.426}]   ,  )

					A_	:					Optional[Any]						 =		image_classifier(
					    [
					        '''http://images.cocodataset.org/val2017/000000039769.jpg''',
					        '''http://images.cocodataset.org/val2017/000000039769.jpg''',
					    ]   ,  top_k=2   ,  )
					self.assertEqual(
					    nested_simplify(_lowerCamelCase   ,  decimals=4							)   ,  [
					        [{'''label''': '''LABEL_1''', '''score''': 0.574}, {'''label''': '''LABEL_0''', '''score''': 0.426}],
					        [{'''label''': '''LABEL_1''', '''score''': 0.574}, {'''label''': '''LABEL_0''', '''score''': 0.426}],
					    ]   ,  )







				def        a_	(				self :       Dict							):



					"""simple docstring"""

					A_	:					Any						 =		PreTrainedTokenizer()

					# Assert that the pipeline can be initialized with a feature extractor that is not in any mapping
					A_	:					List[Any]						 =		pipeline(
					    '''image-classification'''   ,  model='''hf-internal-testing/tiny-random-vit'''   ,  tokenizer=_lowerCamelCase							)

					self.assertIs(image_classifier.tokenizer   ,  _lowerCamelCase							)







				@slow
				@require_torch
				def        a_	(				self :       List[str]							):



					"""simple docstring"""

					A_	:					List[Any]						 =		pipeline('''image-classification'''   ,  model='''deepmind/vision-perceiver-conv'''							)
					A_	:					Optional[Any]						 =		image_classifier('''http://images.cocodataset.org/val2017/000000039769.jpg'''							)
					self.assertEqual(
					    nested_simplify(_lowerCamelCase   ,  decimals=4							)   ,  [
					        {'''score''': 0.4385, '''label''': '''tabby, tabby cat'''},
					        {'''score''': 0.321, '''label''': '''tiger cat'''},
					        {'''score''': 0.0502, '''label''': '''Egyptian cat'''},
					        {'''score''': 0.0137, '''label''': '''crib, cot'''},
					        {'''score''': 0.007, '''label''': '''radiator'''},
					    ]   ,  )

					A_	:					Dict						 =		pipeline('''image-classification'''   ,  model='''deepmind/vision-perceiver-fourier'''							)
					A_	:					str						 =		image_classifier('''http://images.cocodataset.org/val2017/000000039769.jpg'''							)
					self.assertEqual(
					    nested_simplify(_lowerCamelCase   ,  decimals=4							)   ,  [
					        {'''score''': 0.5658, '''label''': '''tabby, tabby cat'''},
					        {'''score''': 0.1309, '''label''': '''tiger cat'''},
					        {'''score''': 0.0722, '''label''': '''Egyptian cat'''},
					        {'''score''': 0.0707, '''label''': '''remote control, remote'''},
					        {'''score''': 0.0082, '''label''': '''computer keyboard, keypad'''},
					    ]   ,  )

					A_	:					Optional[int]						 =		pipeline('''image-classification'''   ,  model='''deepmind/vision-perceiver-learned'''							)
					A_	:					int						 =		image_classifier('''http://images.cocodataset.org/val2017/000000039769.jpg'''							)
					self.assertEqual(
					    nested_simplify(_lowerCamelCase   ,  decimals=4							)   ,  [
					        {'''score''': 0.3022, '''label''': '''tabby, tabby cat'''},
					        {'''score''': 0.2362, '''label''': '''Egyptian cat'''},
					        {'''score''': 0.1856, '''label''': '''tiger cat'''},
					        {'''score''': 0.0324, '''label''': '''remote control, remote'''},
					        {'''score''': 0.0096, '''label''': '''quilt, comforter, comfort, puff'''},
					    ]   ,  )





