


'''simple docstring'''

import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin
@skip_mps
class   __lowerCAmelCase						(__snake_case  ,     __snake_case  ,     unittest.TestCase     ):





      '''simple docstring'''






      lowerCAmelCase__					:     List[str]  					=							IFPipeline
      lowerCAmelCase__					:     Optional[Any]  					=							TEXT_TO_IMAGE_PARAMS - {"""width""", """height""", """latents"""}
      lowerCAmelCase__					:     List[Any]  					=							TEXT_TO_IMAGE_BATCH_PARAMS
      lowerCAmelCase__					:     Union[str, Any]  					=							PipelineTesterMixin.required_optional_params - {"""latents"""}



      def       UpperCamelCase__		(self    :    Optional[int] ):

             '''simple docstring'''







             return self._get_dummy_components()



      def       UpperCamelCase__		(self    :    Optional[Any]   ,						UpperCamelCase    :    List[Any]   ,						UpperCamelCase    :    Optional[int]=0 ):

             '''simple docstring'''







             if str(UpperCamelCase ).startswith('''mps''' ):
                    lowercase__      			= torch.manual_seed(UpperCamelCase )
             else:
                    lowercase__      			= torch.Generator(device=UpperCamelCase ).manual_seed(UpperCamelCase )

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

             return inputs



      def       UpperCamelCase__		(self    :    Tuple ):

             '''simple docstring'''







             self._test_save_load_optional_components()



      @unittest.skipIf(torch_device != '''cuda'''   ,						reason='''float16 requires CUDA''' )
      def       UpperCamelCase__		(self    :    Union[str, Any] ):

             '''simple docstring'''







             super().test_save_load_floataa(expected_max_diff=1E-1 )



      def       UpperCamelCase__		(self    :    Any ):

             '''simple docstring'''







             self._test_attention_slicing_forward_pass(expected_max_diff=1E-2 )



      def       UpperCamelCase__		(self    :    Any ):

             '''simple docstring'''







             self._test_save_load_local()



      def       UpperCamelCase__		(self    :    Dict ):

             '''simple docstring'''







             self._test_inference_batch_single_identical(
                 expected_max_diff=1E-2   ,						)






      @unittest.skipIf(
          torch_device != '''cuda''' or not is_xformers_available()   ,						reason='''XFormers attention is only available with CUDA and `xformers` installed'''   ,						)
      def       UpperCamelCase__		(self    :    Tuple ):

             '''simple docstring'''







             self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3 )






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





      '''simple docstring'''






      def       UpperCamelCase__		(self    :    List[str] ):

             '''simple docstring'''







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



      def       UpperCamelCase__		(self    :    Optional[Any] ):

             '''simple docstring'''







             lowercase__      			= IFPipeline.from_pretrained('''DeepFloyd/IF-I-XL-v1.0'''   ,						variant='''fp16'''   ,						torch_dtype=torch.floataa )

             lowercase__      			= IFSuperResolutionPipeline.from_pretrained(
                 '''DeepFloyd/IF-II-L-v1.0'''   ,						variant='''fp16'''   ,						torch_dtype=torch.floataa   ,						text_encoder=UpperCamelCase   ,						tokenizer=UpperCamelCase )

             # pre compute text embeddings and remove T5 to save memory

             pipe_a.text_encoder.to('''cuda''' )

             lowercase__      			= pipe_a.encode_prompt('''anime turtle'''   ,						device='''cuda''' )

             del pipe_a.tokenizer
             del pipe_a.text_encoder
             gc.collect()

             lowercase__      			= None
             lowercase__      			= None

             pipe_a.enable_model_cpu_offload()
             pipe_a.enable_model_cpu_offload()

             pipe_a.unet.set_attn_processor(AttnAddedKVProcessor() )
             pipe_a.unet.set_attn_processor(AttnAddedKVProcessor() )

             self._test_if(UpperCamelCase   ,						UpperCamelCase   ,						UpperCamelCase   ,						UpperCamelCase )

             pipe_a.remove_all_hooks()
             pipe_a.remove_all_hooks()

             # img2img

             lowercase__      			= IFImgaImgPipeline(**pipe_a.components )
             lowercase__      			= IFImgaImgSuperResolutionPipeline(**pipe_a.components )

             pipe_a.enable_model_cpu_offload()
             pipe_a.enable_model_cpu_offload()

             pipe_a.unet.set_attn_processor(AttnAddedKVProcessor() )
             pipe_a.unet.set_attn_processor(AttnAddedKVProcessor() )

             self._test_if_imgaimg(UpperCamelCase   ,						UpperCamelCase   ,						UpperCamelCase   ,						UpperCamelCase )

             pipe_a.remove_all_hooks()
             pipe_a.remove_all_hooks()

             # inpainting

             lowercase__      			= IFInpaintingPipeline(**pipe_a.components )
             lowercase__      			= IFInpaintingSuperResolutionPipeline(**pipe_a.components )

             pipe_a.enable_model_cpu_offload()
             pipe_a.enable_model_cpu_offload()

             pipe_a.unet.set_attn_processor(AttnAddedKVProcessor() )
             pipe_a.unet.set_attn_processor(AttnAddedKVProcessor() )

             self._test_if_inpainting(UpperCamelCase   ,						UpperCamelCase   ,						UpperCamelCase   ,						UpperCamelCase )



      def       UpperCamelCase__		(self    :    int   ,						UpperCamelCase    :    int   ,						UpperCamelCase    :    Tuple   ,						UpperCamelCase    :    List[str]   ,						UpperCamelCase    :    List[str] ):

             '''simple docstring'''







             _start_torch_memory_measurement()

             lowercase__      			= torch.Generator(device='''cpu''' ).manual_seed(0 )
             lowercase__      			= pipe_a(
                 prompt_embeds=UpperCamelCase   ,						negative_prompt_embeds=UpperCamelCase   ,						num_inference_steps=2   ,						generator=UpperCamelCase   ,						output_type='''np'''   ,						)

             lowercase__      			= output.images[0]

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

             lowercase__      			= torch.cuda.max_memory_allocated()
             assert mem_bytes < 13 * 10**9

             lowercase__      			= load_numpy(
                 '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy''' )
             assert_mean_pixel_difference(UpperCamelCase   ,						UpperCamelCase )

             # pipeline 2

             _start_torch_memory_measurement()

             lowercase__      			= torch.Generator(device='''cpu''' ).manual_seed(0 )

             lowercase__      			= floats_tensor((1, 3, 64, 64)   ,						rng=random.Random(0 ) ).to(UpperCamelCase )

             lowercase__      			= pipe_a(
                 prompt_embeds=UpperCamelCase   ,						negative_prompt_embeds=UpperCamelCase   ,						image=UpperCamelCase   ,						generator=UpperCamelCase   ,						num_inference_steps=2   ,						output_type='''np'''   ,						)

             lowercase__      			= output.images[0]

             assert image.shape == (256, 256, 3)

             lowercase__      			= torch.cuda.max_memory_allocated()
             assert mem_bytes < 4 * 10**9

             lowercase__      			= load_numpy(
                 '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy''' )
             assert_mean_pixel_difference(UpperCamelCase   ,						UpperCamelCase )



      def       UpperCamelCase__		(self    :    int   ,						UpperCamelCase    :    Any   ,						UpperCamelCase    :    Union[str, Any]   ,						UpperCamelCase    :    List[str]   ,						UpperCamelCase    :    Dict ):

             '''simple docstring'''







             _start_torch_memory_measurement()

             lowercase__      			= floats_tensor((1, 3, 64, 64)   ,						rng=random.Random(0 ) ).to(UpperCamelCase )

             lowercase__      			= torch.Generator(device='''cpu''' ).manual_seed(0 )

             lowercase__      			= pipe_a(
                 prompt_embeds=UpperCamelCase   ,						negative_prompt_embeds=UpperCamelCase   ,						image=UpperCamelCase   ,						num_inference_steps=2   ,						generator=UpperCamelCase   ,						output_type='''np'''   ,						)

             lowercase__      			= output.images[0]

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

             lowercase__      			= torch.cuda.max_memory_allocated()
             assert mem_bytes < 10 * 10**9

             lowercase__      			= load_numpy(
                 '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy''' )
             assert_mean_pixel_difference(UpperCamelCase   ,						UpperCamelCase )

             # pipeline 2

             _start_torch_memory_measurement()

             lowercase__      			= torch.Generator(device='''cpu''' ).manual_seed(0 )

             lowercase__      			= floats_tensor((1, 3, 256, 256)   ,						rng=random.Random(0 ) ).to(UpperCamelCase )
             lowercase__      			= floats_tensor((1, 3, 64, 64)   ,						rng=random.Random(0 ) ).to(UpperCamelCase )

             lowercase__      			= pipe_a(
                 prompt_embeds=UpperCamelCase   ,						negative_prompt_embeds=UpperCamelCase   ,						image=UpperCamelCase   ,						original_image=UpperCamelCase   ,						generator=UpperCamelCase   ,						num_inference_steps=2   ,						output_type='''np'''   ,						)

             lowercase__      			= output.images[0]

             assert image.shape == (256, 256, 3)

             lowercase__      			= torch.cuda.max_memory_allocated()
             assert mem_bytes < 4 * 10**9

             lowercase__      			= load_numpy(
                 '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy''' )
             assert_mean_pixel_difference(UpperCamelCase   ,						UpperCamelCase )






      def       UpperCamelCase__		(self    :    Optional[int]   ,						UpperCamelCase    :    Tuple   ,						UpperCamelCase    :    Tuple   ,						UpperCamelCase    :    List[Any]   ,						UpperCamelCase    :    List[str] ):

             '''simple docstring'''







             _start_torch_memory_measurement()

             lowercase__      			= floats_tensor((1, 3, 64, 64)   ,						rng=random.Random(0 ) ).to(UpperCamelCase )
             lowercase__      			= floats_tensor((1, 3, 64, 64)   ,						rng=random.Random(1 ) ).to(UpperCamelCase )

             lowercase__      			= torch.Generator(device='''cpu''' ).manual_seed(0 )
             lowercase__      			= pipe_a(
                 prompt_embeds=UpperCamelCase   ,						negative_prompt_embeds=UpperCamelCase   ,						image=UpperCamelCase   ,						mask_image=UpperCamelCase   ,						num_inference_steps=2   ,						generator=UpperCamelCase   ,						output_type='''np'''   ,						)

             lowercase__      			= output.images[0]

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

             lowercase__      			= torch.cuda.max_memory_allocated()
             assert mem_bytes < 10 * 10**9

             lowercase__      			= load_numpy(
                 '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy''' )
             assert_mean_pixel_difference(UpperCamelCase   ,						UpperCamelCase )

             # pipeline 2

             _start_torch_memory_measurement()

             lowercase__      			= torch.Generator(device='''cpu''' ).manual_seed(0 )

             lowercase__      			= floats_tensor((1, 3, 64, 64)   ,						rng=random.Random(0 ) ).to(UpperCamelCase )
             lowercase__      			= floats_tensor((1, 3, 256, 256)   ,						rng=random.Random(0 ) ).to(UpperCamelCase )
             lowercase__      			= floats_tensor((1, 3, 256, 256)   ,						rng=random.Random(1 ) ).to(UpperCamelCase )

             lowercase__      			= pipe_a(
                 prompt_embeds=UpperCamelCase   ,						negative_prompt_embeds=UpperCamelCase   ,						image=UpperCamelCase   ,						mask_image=UpperCamelCase   ,						original_image=UpperCamelCase   ,						generator=UpperCamelCase   ,						num_inference_steps=2   ,						output_type='''np'''   ,						)

             lowercase__      			= output.images[0]

             assert image.shape == (256, 256, 3)

             lowercase__      			= torch.cuda.max_memory_allocated()
             assert mem_bytes < 4 * 10**9

             lowercase__      			= load_numpy(
                 '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy''' )
             assert_mean_pixel_difference(UpperCamelCase   ,						UpperCamelCase )






def _SCREAMING_SNAKE_CASE      ()       ->       int:






       """simple docstring"""







       torch.cuda.empty_cache()
       torch.cuda.reset_max_memory_allocated()
       torch.cuda.reset_peak_memory_stats()






'''simple docstring'''




import unittest

from transformers import AutoConfig, AutoTokenizer, BertConfig, TensorType, is_flax_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, slow


if is_flax_available():
      import jax

      from transformers.models.auto.modeling_flax_auto import FlaxAutoModel
      from transformers.models.bert.modeling_flax_bert import FlaxBertModel
      from transformers.models.roberta.modeling_flax_roberta import FlaxRobertaModel





@require_flax
class  A				(							unittest.TestCase						):





     @slow
     def 				__lowerCAmelCase							(							self	) ->				List[Any]:






         """simple docstring"""






         for model_name in ["bert-base-cased", "bert-large-uncased"]:
             with self.subTest(SCREAMING_SNAKE_CASE	):
                 A       :   int					=						AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE	)
                 self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                 self.assertIsInstance(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	)

                 A       :   List[str]					=						FlaxAutoModel.from_pretrained(SCREAMING_SNAKE_CASE	)
                 self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                 self.assertIsInstance(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	)





     @slow
     def 				__lowerCAmelCase							(							self	) ->				int:






         """simple docstring"""






         for model_name in ["roberta-base", "roberta-large"]:
             with self.subTest(SCREAMING_SNAKE_CASE	):
                 A       :   Any					=						AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE	)
                 self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                 self.assertIsInstance(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	)

                 A       :   Any					=						FlaxAutoModel.from_pretrained(SCREAMING_SNAKE_CASE	)
                 self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                 self.assertIsInstance(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	)





     @slow
     def 				__lowerCAmelCase							(							self	) ->				Any:






         """simple docstring"""






         for model_name in ["bert-base-cased", "bert-large-uncased"]:
             A       :   Optional[int]					=						AutoTokenizer.from_pretrained(SCREAMING_SNAKE_CASE	)
             A       :   List[str]					=						FlaxBertModel.from_pretrained(SCREAMING_SNAKE_CASE	)
             A       :   Optional[Any]					=						tokenizer('''Do you support jax jitted function?'''						,			return_tensors=TensorType.JAX	)

             @jax.jit
             def eval(**SCREAMING_SNAKE_CASE	):
                 return model(**SCREAMING_SNAKE_CASE	)

             eval(**SCREAMING_SNAKE_CASE	).block_until_ready()





     @slow
     def 				__lowerCAmelCase							(							self	) ->				List[str]:






         """simple docstring"""






         for model_name in ["roberta-base", "roberta-large"]:
             A       :   List[str]					=						AutoTokenizer.from_pretrained(SCREAMING_SNAKE_CASE	)
             A       :   Union[str, Any]					=						FlaxRobertaModel.from_pretrained(SCREAMING_SNAKE_CASE	)
             A       :   int					=						tokenizer('''Do you support jax jitted function?'''						,			return_tensors=TensorType.JAX	)

             @jax.jit
             def eval(**SCREAMING_SNAKE_CASE	):
                 return model(**SCREAMING_SNAKE_CASE	)

             eval(**SCREAMING_SNAKE_CASE	).block_until_ready()





     def 				__lowerCAmelCase							(							self	) ->				List[str]:






         """simple docstring"""






         with self.assertRaisesRegex(
             SCREAMING_SNAKE_CASE						,			'''bert-base is not a local folder and is not a valid model identifier'''	):
             A       :   List[Any]					=						FlaxAutoModel.from_pretrained('''bert-base'''	)





     def 				__lowerCAmelCase							(							self	) ->				Tuple:






         """simple docstring"""






         with self.assertRaisesRegex(
             SCREAMING_SNAKE_CASE						,			R'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''	):
             A       :   Optional[int]					=						FlaxAutoModel.from_pretrained(SCREAMING_SNAKE_CASE						,			revision='''aaaaaa'''	)





     def 				__lowerCAmelCase							(							self	) ->				str:






         """simple docstring"""






         with self.assertRaisesRegex(
             SCREAMING_SNAKE_CASE						,			'''hf-internal-testing/config-no-model does not appear to have a file named flax_model.msgpack'''						,			):
             A       :   List[str]					=						FlaxAutoModel.from_pretrained('''hf-internal-testing/config-no-model'''	)





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






         """simple docstring"""






         with self.assertRaisesRegex(SCREAMING_SNAKE_CASE						,			'''Use `from_pt=True` to load this model'''	):
             A       :   Any					=						FlaxAutoModel.from_pretrained('''hf-internal-testing/tiny-bert-pt-only'''	)




