



"""simple docstring"""






import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

@skip_mps
class 			__SCREAMING_SNAKE_CASE      ( lowerCAmelCase_						,  lowerCAmelCase_						,  lowerCAmelCase_						,  unittest.TestCase  ):


					'''simple docstring'''





					_a       =				StableDiffusionModelEditingPipeline
					_a       =				TEXT_TO_IMAGE_PARAMS
					_a       =				TEXT_TO_IMAGE_BATCH_PARAMS
					_a       =				TEXT_TO_IMAGE_IMAGE_PARAMS
					_a       =				TEXT_TO_IMAGE_IMAGE_PARAMS



					def 						snake_case			(			self				:    Optional[int]	)->   str:
								torch.manual_seed(0	)
								lowerCamelCase__						:						Optional[int]						      =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,    )
								lowerCamelCase__						:						Any						      =DDIMScheduler()
								torch.manual_seed(0	)
								lowerCamelCase__						:						Union[str, Any]						      =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,    )
								torch.manual_seed(0	)
								lowerCamelCase__						:						List[Any]						      =CLIPTextConfig(
								    bos_token_id=0,    eos_token_id=2,    hidden_size=32,    intermediate_size=37,    layer_norm_eps=1E-05,    num_attention_heads=4,    num_hidden_layers=5,    pad_token_id=1,    vocab_size=1000,    )
								lowerCamelCase__						:						List[Any]						      =CLIPTextModel(lowerCamelCase	)
								lowerCamelCase__						:						Tuple						      =CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''	)

								lowerCamelCase__						:						Any						      ={
								    '''unet''': unet,
								    '''scheduler''': scheduler,
								    '''vae''': vae,
								    '''text_encoder''': text_encoder,
								    '''tokenizer''': tokenizer,
								    '''safety_checker''': None,
								    '''feature_extractor''': None,
								}
								return components



					def 						snake_case			(			self				:    Tuple,    lowerCamelCase				:    Optional[int],    lowerCamelCase				:    List[Any]=0	)->   Dict:
								lowerCamelCase__						:						Union[str, Any]						      =torch.manual_seed(lowerCamelCase	)
								lowerCamelCase__						:						Tuple						      ={
								    '''prompt''': '''A field of roses''',
								    '''generator''': generator,
								    # Setting height and width to None to prevent OOMs on CPU.
								    '''height''': None,
								    '''width''': None,
								    '''num_inference_steps''': 2,
								    '''guidance_scale''': 6.0,
								    '''output_type''': '''numpy''',
								}
								return inputs



					def 						snake_case			(			self				:    List[Any]	)->   int:
								lowerCamelCase__						:						str						      ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								lowerCamelCase__						:						int						      =self.get_dummy_components()
								lowerCamelCase__						:						Any						      =StableDiffusionModelEditingPipeline(**lowerCamelCase	)
								lowerCamelCase__						:						List[Any]						      =sd_pipe.to(lowerCamelCase	)
								sd_pipe.set_progress_bar_config(disable=lowerCamelCase	)

								lowerCamelCase__						:						Optional[int]						      =self.get_dummy_inputs(lowerCamelCase	)
								lowerCamelCase__						:						Dict						      =sd_pipe(**lowerCamelCase	).images
								lowerCamelCase__						:						List[Any]						      =image[0, -3:, -3:, -1]
								assert image.shape == (1, 64, 64, 3)

								lowerCamelCase__						:						Optional[Any]						      =np.array([0.4_755, 0.5_132, 0.4_976, 0.3_904, 0.3_554, 0.4_765, 0.5_139, 0.5_158, 0.4_889]	)

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



					def 						snake_case			(			self				:    int	)->   List[Any]:
								lowerCamelCase__						:						Dict						      ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								lowerCamelCase__						:						List[Any]						      =self.get_dummy_components()
								lowerCamelCase__						:						int						      =StableDiffusionModelEditingPipeline(**lowerCamelCase	)
								lowerCamelCase__						:						Optional[int]						      =sd_pipe.to(lowerCamelCase	)
								sd_pipe.set_progress_bar_config(disable=lowerCamelCase	)

								lowerCamelCase__						:						Optional[int]						      =self.get_dummy_inputs(lowerCamelCase	)
								lowerCamelCase__						:						Union[str, Any]						      ='''french fries'''
								lowerCamelCase__						:						Optional[int]						      =sd_pipe(**lowerCamelCase,    negative_prompt=lowerCamelCase	)
								lowerCamelCase__						:						Optional[Any]						      =output.images
								lowerCamelCase__						:						str						      =image[0, -3:, -3:, -1]

								assert image.shape == (1, 64, 64, 3)

								lowerCamelCase__						:						Union[str, Any]						      =np.array([0.4_992, 0.5_101, 0.5_004, 0.3_949, 0.3_604, 0.4_735, 0.5_216, 0.5_204, 0.4_913]	)

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



					def 						snake_case			(			self				:    List[Any]	)->   Tuple:
								lowerCamelCase__						:						int						      ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								lowerCamelCase__						:						Dict						      =self.get_dummy_components()
								lowerCamelCase__						:						Any						      =EulerAncestralDiscreteScheduler(
								    beta_start=0.00_085,    beta_end=0.012,    beta_schedule='''scaled_linear'''	)
								lowerCamelCase__						:						List[Any]						      =StableDiffusionModelEditingPipeline(**lowerCamelCase	)
								lowerCamelCase__						:						Union[str, Any]						      =sd_pipe.to(lowerCamelCase	)
								sd_pipe.set_progress_bar_config(disable=lowerCamelCase	)

								lowerCamelCase__						:						Optional[Any]						      =self.get_dummy_inputs(lowerCamelCase	)
								lowerCamelCase__						:						int						      =sd_pipe(**lowerCamelCase	).images
								lowerCamelCase__						:						Tuple						      =image[0, -3:, -3:, -1]

								assert image.shape == (1, 64, 64, 3)

								lowerCamelCase__						:						List[str]						      =np.array([0.4_747, 0.5_372, 0.4_779, 0.4_982, 0.5_543, 0.4_816, 0.5_238, 0.4_904, 0.5_027]	)

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



					def 						snake_case			(			self				:    List[str]	)->   Dict:
								lowerCamelCase__						:						Any						      ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								lowerCamelCase__						:						Union[str, Any]						      =self.get_dummy_components()
								lowerCamelCase__						:						str						      =PNDMScheduler()
								lowerCamelCase__						:						Dict						      =StableDiffusionModelEditingPipeline(**lowerCamelCase	)
								lowerCamelCase__						:						Optional[Any]						      =sd_pipe.to(lowerCamelCase	)
								sd_pipe.set_progress_bar_config(disable=lowerCamelCase	)

								lowerCamelCase__						:						Tuple						      =self.get_dummy_inputs(lowerCamelCase	)
								# the pipeline does not expect pndm so test if it raises error.
								with self.assertRaises(lowerCamelCase	):
											lowerCamelCase__						:						Union[str, Any]						      =sd_pipe(**lowerCamelCase	).images



					def 						snake_case			(			self				:    Optional[int]	)->   Optional[Any]:
								super().test_inference_batch_single_identical(expected_max_diff=5E-3	)




					def 						snake_case			(			self				:    Any	)->   str:
								super().test_attention_slicing_forward_pass(expected_max_diff=5E-3	)





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


					'''simple docstring'''





					def 						snake_case			(			self				:    Any	)->   int:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()



					def 						snake_case			(			self				:    Union[str, Any],    lowerCamelCase				:    Dict=0	)->   List[Any]:
								lowerCamelCase__						:						Any						      =torch.manual_seed(lowerCamelCase	)
								lowerCamelCase__						:						List[Any]						      ={
								    '''prompt''': '''A field of roses''',
								    '''generator''': generator,
								    '''num_inference_steps''': 3,
								    '''guidance_scale''': 7.5,
								    '''output_type''': '''numpy''',
								}
								return inputs



					def 						snake_case			(			self				:    int	)->   Any:
								lowerCamelCase__						:						Tuple						      ='''CompVis/stable-diffusion-v1-4'''
								lowerCamelCase__						:						Tuple						      =StableDiffusionModelEditingPipeline.from_pretrained(lowerCamelCase,    safety_checker=lowerCamelCase	)
								pipe.to(lowerCamelCase	)
								pipe.set_progress_bar_config(disable=lowerCamelCase	)
								pipe.enable_attention_slicing()

								lowerCamelCase__						:						Tuple						      =self.get_inputs()
								lowerCamelCase__						:						List[str]						      =pipe(**lowerCamelCase	).images
								lowerCamelCase__						:						Optional[int]						      =image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 512, 512, 3)

								lowerCamelCase__						:						Tuple						      =np.array(
								    [0.6_749_496, 0.6_386_453, 0.51_443_267, 0.66_094_905, 0.61_921_215, 0.5_491_332, 0.5_744_417, 0.58_075_106, 0.5_174_658]	)

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

								# make sure image changes after editing
								pipe.edit_model('''A pack of roses''',    '''A pack of blue roses'''	)

								lowerCamelCase__						:						Any						      =pipe(**lowerCamelCase	).images
								lowerCamelCase__						:						List[str]						      =image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 512, 512, 3)

								assert np.abs(expected_slice - image_slice	).max() > 1E-1




					def 						snake_case			(			self				:    Union[str, Any]	)->   Tuple:
								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

								lowerCamelCase__						:						Any						      ='''CompVis/stable-diffusion-v1-4'''
								lowerCamelCase__						:						Tuple						      =DDIMScheduler.from_pretrained(lowerCamelCase,    subfolder='''scheduler'''	)
								lowerCamelCase__						:						Tuple						      =StableDiffusionModelEditingPipeline.from_pretrained(
								    lowerCamelCase,    scheduler=lowerCamelCase,    safety_checker=lowerCamelCase	)
								lowerCamelCase__						:						int						      =pipe.to(lowerCamelCase	)
								pipe.set_progress_bar_config(disable=lowerCamelCase	)
								pipe.enable_attention_slicing(1	)
								pipe.enable_sequential_cpu_offload()

								lowerCamelCase__						:						int						      =self.get_inputs()
								lowerCamelCase__						:						Optional[int]						      =pipe(**lowerCamelCase	)

								lowerCamelCase__						:						int						      =torch.cuda.max_memory_allocated()
								# make sure that less than 4.4 GB is allocated
								assert mem_bytes < 4.4 * 10**9










'''simple docstring'''







import inspect
import unittest

from transformers import is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import MODEL_MAPPING, PoolFormerConfig, PoolFormerForImageClassification, PoolFormerModel
   from transformers.models.poolformer.modeling_poolformer import POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
   from PIL import Image

   from transformers import PoolFormerImageProcessor


class      _a    (      __a					):
       def        A (      self		:      Tuple							):



           '''simple docstring'''

           UpperCAmelCase     			=						self.config_class(**self.inputs_dict							)
           self.parent.assertTrue(hasattr(lowercase  ,							'''hidden_sizes'''							)							)
           self.parent.assertTrue(hasattr(lowercase  ,							'''num_encoder_blocks'''							)							)






class      _a    :
       def __init__(      self		:      Any  ,							lowercase		:      Any  ,							lowercase		:      Optional[Any]=13  ,							lowercase		:      List[str]=64  ,							lowercase		:      Optional[Any]=3  ,							lowercase		:      List[str]=4  ,							lowercase		:      Optional[int]=[2, 2, 2, 2]  ,							lowercase		:      Tuple=[8, 4, 2, 1]  ,							lowercase		:      Optional[int]=[16, 32, 64, 128]  ,							lowercase		:      Tuple=[1, 4, 8, 16]  ,							lowercase		:      List[Any]=False  ,							lowercase		:      List[Any]=True  ,							lowercase		:      List[str]="gelu"  ,							lowercase		:      Any=0.1  ,							lowercase		:      Dict=0.02  ,							lowercase		:      str=3  ,							lowercase		:      Optional[Any]=None  ,							):



           '''simple docstring'''

           UpperCAmelCase     			=						parent
           UpperCAmelCase     			=						batch_size
           UpperCAmelCase     			=						image_size
           UpperCAmelCase     			=						num_channels
           UpperCAmelCase     			=						num_encoder_blocks
           UpperCAmelCase     			=						sr_ratios
           UpperCAmelCase     			=						depths
           UpperCAmelCase     			=						hidden_sizes
           UpperCAmelCase     			=						downsampling_rates
           UpperCAmelCase     			=						is_training
           UpperCAmelCase     			=						use_labels
           UpperCAmelCase     			=						hidden_act
           UpperCAmelCase     			=						hidden_dropout_prob
           UpperCAmelCase     			=						initializer_range
           UpperCAmelCase     			=						num_labels
           UpperCAmelCase     			=						scope
       def        A (      self		:      Union[str, Any]							):



           '''simple docstring'''

           UpperCAmelCase     			=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

           UpperCAmelCase     			=						None
           if self.use_labels:
               UpperCAmelCase     			=						ids_tensor([self.batch_size, self.image_size, self.image_size]  ,							self.num_labels							)

           UpperCAmelCase     			=						PoolFormerConfig(
               image_size=self.image_size  ,							num_channels=self.num_channels  ,							num_encoder_blocks=self.num_encoder_blocks  ,							depths=self.depths  ,							hidden_sizes=self.hidden_sizes  ,							hidden_act=self.hidden_act  ,							hidden_dropout_prob=self.hidden_dropout_prob  ,							initializer_range=self.initializer_range  ,							)

           return config, pixel_values, labels
       def        A (      self		:      Optional[Any]  ,							lowercase		:      List[Any]  ,							lowercase		:      Dict  ,							lowercase		:      Union[str, Any]							):



           '''simple docstring'''

           UpperCAmelCase     			=						PoolFormerModel(config=lowercase							)
           model.to(lowercase							)
           model.eval()
           UpperCAmelCase     			=						model(lowercase							)
           UpperCAmelCase     			=						UpperCAmelCase     			=						self.image_size // 32.0
           self.parent.assertEqual(
               result.last_hidden_state.shape  ,							(self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)							)







       def        A (      self		:      Union[str, Any]							):



           '''simple docstring'''

           UpperCAmelCase     			=						self.prepare_config_and_inputs()
           UpperCAmelCase			,	UpperCAmelCase			,	UpperCAmelCase     			=						config_and_inputs
           UpperCAmelCase     			=						{'''pixel_values''': pixel_values}
           return config, inputs_dict






@require_torch
class      _a    (      __a				,				__a				,				unittest.TestCase					):
       __a							:	Optional[int]  							=   (PoolFormerModel, PoolFormerForImageClassification) if is_torch_available() else ()
       __a							:	Dict  							=   (
           {"""feature-extraction""": PoolFormerModel, """image-classification""": PoolFormerForImageClassification}
           if is_torch_available()
           else {}
       )

       __a							:	str  							=   False
       __a							:	Optional[int]  							=   False
       __a							:	Tuple  							=   False
       __a							:	Union[str, Any]  							=   False
       __a							:	Dict  							=   False
       def        A (      self		:      Any							):



           '''simple docstring'''

           UpperCAmelCase     			=						PoolFormerModelTester(self							)
           UpperCAmelCase     			=						PoolFormerConfigTester(self  ,							config_class=lowercase							)
       def        A (      self		:      Any							):



           '''simple docstring'''

           self.config_tester.run_common_tests()
       def        A (      self		:      Optional[int]							):



           '''simple docstring'''

           UpperCAmelCase     			=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*lowercase							)
       @unittest.skip('''PoolFormer does not use inputs_embeds'''							)
       def        A (      self		:      Dict							):



           '''simple docstring'''

           pass
       @unittest.skip('''PoolFormer does not have get_input_embeddings method and get_output_embeddings methods'''							)
       def        A (      self		:      Union[str, Any]							):



           '''simple docstring'''

           pass
       def        A (      self		:      Dict							):



           '''simple docstring'''

           def check_hidden_states_output(lowercase		:      int  ,							lowercase		:      Any  ,							lowercase		:      Union[str, Any]							):
               UpperCAmelCase     			=						model_class(lowercase							)
               model.to(lowercase							)
               model.eval()

               with torch.no_grad():
                   UpperCAmelCase     			=						model(**self._prepare_for_class(lowercase  ,							lowercase							)							)

               UpperCAmelCase     			=						outputs.hidden_states

               UpperCAmelCase     			=						self.model_tester.num_encoder_blocks
               self.assertEqual(len(lowercase							)  ,							lowercase							)

               # verify the first hidden states (first block)
               self.assertListEqual(
                   list(hidden_states[0].shape[-3:]							)  ,							[
                       self.model_tester.hidden_sizes[0],
                       self.model_tester.image_size // 4,
                       self.model_tester.image_size // 4,
                   ]  ,							)

           UpperCAmelCase			,	UpperCAmelCase     			=						self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               UpperCAmelCase     			=						True
               check_hidden_states_output(lowercase  ,							lowercase  ,							lowercase							)

               # check that output_hidden_states also work using config
               del inputs_dict["output_hidden_states"]
               UpperCAmelCase     			=						True

               check_hidden_states_output(lowercase  ,							lowercase  ,							lowercase							)
       def        A (      self		:      Optional[Any]							):



           '''simple docstring'''

           if not self.model_tester.is_training:
               return

           UpperCAmelCase			,	UpperCAmelCase     			=						self.model_tester.prepare_config_and_inputs_for_common()
           UpperCAmelCase     			=						True

           for model_class in self.all_model_classes:
               if model_class in get_values(lowercase							):
                   continue
               UpperCAmelCase     			=						model_class(lowercase							)
               model.to(lowercase							)
               model.train()
               UpperCAmelCase     			=						self._prepare_for_class(lowercase  ,							lowercase  ,							return_labels=lowercase							)
               UpperCAmelCase     			=						model(**lowercase							).loss
               loss.backward()
       def        A (      self		:      str							):



           '''simple docstring'''

           UpperCAmelCase			,	UpperCAmelCase     			=						self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               UpperCAmelCase     			=						model_class(lowercase							)
               UpperCAmelCase     			=						inspect.signature(model.forward							)
               # signature.parameters is an OrderedDict => so arg_names order is deterministic
               UpperCAmelCase     			=						[*signature.parameters.keys()]

               UpperCAmelCase     			=						['''pixel_values''']
               self.assertListEqual(arg_names[:1]  ,							lowercase							)







       @slow
       def        A (      self		:      Union[str, Any]							):



           '''simple docstring'''

           for model_name in POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               UpperCAmelCase     			=						PoolFormerModel.from_pretrained(lowercase							)
               self.assertIsNotNone(lowercase							)






def 		snake_case_				():
    UpperCAmelCase     			=						Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''	)
    return image







@require_torch
class      _a    (      unittest.TestCase					):
       @slow
       def        A (      self		:      str							):



           '''simple docstring'''

           UpperCAmelCase     			=						PoolFormerImageProcessor()
           UpperCAmelCase     			=						PoolFormerForImageClassification.from_pretrained('''sail/poolformer_s12'''							).to(lowercase							)

           UpperCAmelCase     			=						image_processor(images=prepare_img()  ,							return_tensors='''pt'''							).to(lowercase							)

           # forward pass
           with torch.no_grad():
               UpperCAmelCase     			=						model(**lowercase							)

           # verify the logits
           UpperCAmelCase     			=						torch.Size((1, 1_000)							)
           self.assertEqual(outputs.logits.shape  ,							lowercase							)

           UpperCAmelCase     			=						torch.tensor([-0.6113, 0.1685, -0.0492]							).to(lowercase							)
           self.assertTrue(torch.allclose(outputs.logits[0, :3]  ,							lowercase  ,							atol=1E-4							)							)




