

'''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 							lowercase_						(		_a			,							_a			,							unittest.TestCase			):

  '''simple docstring'''





  __lowerCAmelCase       :			Any     	= IFPipeline
  __lowerCAmelCase       :			str     	= TEXT_TO_IMAGE_PARAMS - {'''width''', '''height''', '''latents'''}
  __lowerCAmelCase       :			Dict     	= TEXT_TO_IMAGE_BATCH_PARAMS
  __lowerCAmelCase       :			Tuple     	= PipelineTesterMixin.required_optional_params - {'''latents'''}
  def     snake_case_      (  self  )				->					str:


    """simple docstring"""
    return self._get_dummy_components()
  def     snake_case_      (  self ,			a_ ,			a_=0  )				->					str:


    """simple docstring"""
    if str(__UpperCamelCase  ).startswith('mps'  ):
      UpperCAmelCase    				=						torch.manual_seed(__UpperCamelCase  )
    else:
      UpperCAmelCase    				=						torch.Generator(device=__UpperCamelCase  ).manual_seed(__UpperCamelCase  )

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

    return inputs
  def     snake_case_      (  self  )				->					int:


    """simple docstring"""
    self._test_save_load_optional_components()
  @unittest.skipIf(torch_device != 'cuda' ,			reason='float16 requires CUDA'  )
  def     snake_case_      (  self  )				->					List[Any]:


    """simple docstring"""
    # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
    super().test_save_load_floataa(expected_max_diff=1E-1  )
  def     snake_case_      (  self  )				->					Tuple:


    """simple docstring"""
    self._test_attention_slicing_forward_pass(expected_max_diff=1E-2  )
  def     snake_case_      (  self  )				->					Union[str, Any]:


    """simple docstring"""
    self._test_save_load_local()
  def     snake_case_      (  self  )				->					List[Any]:


    """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     snake_case_      (  self  )				->					str:


    """simple docstring"""
    self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3  )




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

  '''simple docstring'''





  def     snake_case_      (  self  )				->					Dict:


    """simple docstring"""
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()
  def     snake_case_      (  self  )				->					List[Any]:


    """simple docstring"""
    # if

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

    UpperCAmelCase    				=						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'  )

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

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

    UpperCAmelCase    				=						None
    UpperCAmelCase    				=						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

    UpperCAmelCase    				=						IFImgaImgPipeline(**pipe_a.components  )
    UpperCAmelCase    				=						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

    UpperCAmelCase    				=						IFInpaintingPipeline(**pipe_a.components  )
    UpperCAmelCase    				=						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     snake_case_      (  self ,			a_ ,			a_ ,			a_ ,			a_  )				->					Tuple:


    """simple docstring"""
    # pipeline 1

    _start_torch_memory_measurement()

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

    UpperCAmelCase    				=						output.images[0]

    assert image.shape == (6_4, 6_4, 3)

    UpperCAmelCase    				=						torch.cuda.max_memory_allocated()
    assert mem_bytes < 1_3 * 1_0**9

    UpperCAmelCase    				=						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()

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

    UpperCAmelCase    				=						floats_tensor((1, 3, 6_4, 6_4) ,			rng=random.Random(0  )  ).to(__UpperCamelCase  )

    UpperCAmelCase    				=						pipe_a(
        prompt_embeds=__UpperCamelCase ,			negative_prompt_embeds=__UpperCamelCase ,			image=__UpperCamelCase ,			generator=__UpperCamelCase ,			num_inference_steps=2 ,			output_type='np' ,			)

    UpperCAmelCase    				=						output.images[0]

    assert image.shape == (2_5_6, 2_5_6, 3)

    UpperCAmelCase    				=						torch.cuda.max_memory_allocated()
    assert mem_bytes < 4 * 1_0**9

    UpperCAmelCase    				=						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     snake_case_      (  self ,			a_ ,			a_ ,			a_ ,			a_  )				->					Dict:


    """simple docstring"""
    # pipeline 1

    _start_torch_memory_measurement()

    UpperCAmelCase    				=						floats_tensor((1, 3, 6_4, 6_4) ,			rng=random.Random(0  )  ).to(__UpperCamelCase  )

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

    UpperCAmelCase    				=						pipe_a(
        prompt_embeds=__UpperCamelCase ,			negative_prompt_embeds=__UpperCamelCase ,			image=__UpperCamelCase ,			num_inference_steps=2 ,			generator=__UpperCamelCase ,			output_type='np' ,			)

    UpperCAmelCase    				=						output.images[0]

    assert image.shape == (6_4, 6_4, 3)

    UpperCAmelCase    				=						torch.cuda.max_memory_allocated()
    assert mem_bytes < 1_0 * 1_0**9

    UpperCAmelCase    				=						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()

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

    UpperCAmelCase    				=						floats_tensor((1, 3, 2_5_6, 2_5_6) ,			rng=random.Random(0  )  ).to(__UpperCamelCase  )
    UpperCAmelCase    				=						floats_tensor((1, 3, 6_4, 6_4) ,			rng=random.Random(0  )  ).to(__UpperCamelCase  )

    UpperCAmelCase    				=						pipe_a(
        prompt_embeds=__UpperCamelCase ,			negative_prompt_embeds=__UpperCamelCase ,			image=__UpperCamelCase ,			original_image=__UpperCamelCase ,			generator=__UpperCamelCase ,			num_inference_steps=2 ,			output_type='np' ,			)

    UpperCAmelCase    				=						output.images[0]

    assert image.shape == (2_5_6, 2_5_6, 3)

    UpperCAmelCase    				=						torch.cuda.max_memory_allocated()
    assert mem_bytes < 4 * 1_0**9

    UpperCAmelCase    				=						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     snake_case_      (  self ,			a_ ,			a_ ,			a_ ,			a_  )				->					int:


    """simple docstring"""
    # pipeline 1

    _start_torch_memory_measurement()

    UpperCAmelCase    				=						floats_tensor((1, 3, 6_4, 6_4) ,			rng=random.Random(0  )  ).to(__UpperCamelCase  )
    UpperCAmelCase    				=						floats_tensor((1, 3, 6_4, 6_4) ,			rng=random.Random(1  )  ).to(__UpperCamelCase  )

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

    UpperCAmelCase    				=						output.images[0]

    assert image.shape == (6_4, 6_4, 3)

    UpperCAmelCase    				=						torch.cuda.max_memory_allocated()
    assert mem_bytes < 1_0 * 1_0**9

    UpperCAmelCase    				=						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()

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

    UpperCAmelCase    				=						floats_tensor((1, 3, 6_4, 6_4) ,			rng=random.Random(0  )  ).to(__UpperCamelCase  )
    UpperCAmelCase    				=						floats_tensor((1, 3, 2_5_6, 2_5_6) ,			rng=random.Random(0  )  ).to(__UpperCamelCase  )
    UpperCAmelCase    				=						floats_tensor((1, 3, 2_5_6, 2_5_6) ,			rng=random.Random(1  )  ).to(__UpperCamelCase  )

    UpperCAmelCase    				=						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' ,			)

    UpperCAmelCase    				=						output.images[0]

    assert image.shape == (2_5_6, 2_5_6, 3)

    UpperCAmelCase    				=						torch.cuda.max_memory_allocated()
    assert mem_bytes < 4 * 1_0**9

    UpperCAmelCase    				=						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    lowerCamelCase__		(    ):
  torch.cuda.empty_cache()
  torch.cuda.reset_max_memory_allocated()
  torch.cuda.reset_peak_memory_stats()




import datasets


A		: Optional[int]								=     '''\
@InProceedings{conneau2018xnli,
  author = "Conneau, Alexis
                 and Rinott, Ruty
                 and Lample, Guillaume
                 and Williams, Adina
                 and Bowman, Samuel R.
                 and Schwenk, Holger
                 and Stoyanov, Veselin",
  title = "XNLI: Evaluating Cross-lingual Sentence Representations",
  booktitle = "Proceedings of the 2018 Conference on Empirical Methods
               in Natural Language Processing",
  year = "2018",
  publisher = "Association for Computational Linguistics",
  location = "Brussels, Belgium",
}
'''

A		: Optional[int]								=     '''\
XNLI is a subset of a few thousand examples from MNLI which has been translated
into a 14 different languages (some low-ish resource). As with MNLI, the goal is
to predict textual entailment (does sentence A imply/contradict/neither sentence
B) and is a classification task (given two sentences, predict one of three
labels).
'''

A		: str								=     '''
Computes XNLI score which is just simple accuracy.
Args:
    predictions: Predicted labels.
    references: Ground truth labels.
Returns:
    \'accuracy\': accuracy
Examples:

    >>> predictions = [0, 1]
    >>> references = [0, 1]
    >>> xnli_metric = datasets.load_metric("xnli")
    >>> results = xnli_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'accuracy\': 1.0}
'''




def 	UpperCamelCase__  (     SCREAMING_SNAKE_CASE_	:	List[str]		,   SCREAMING_SNAKE_CASE_	:	List[Any]       )  ->     str:
						return (preds == labels).mean()



@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION					,			_KWARGS_DESCRIPTION    )
class    a_ (			datasets.Metric    ):






			def    UpperCamelCase_				(      self  ):
									return datasets.MetricInfo(
									    description=_DESCRIPTION				,			citation=_CITATION				,			inputs_description=_KWARGS_DESCRIPTION				,			features=datasets.Features(
									        {
									            """predictions""": datasets.Value("""int64""" if self.config_name != """sts-b""" else """float32"""  ),
									            """references""": datasets.Value("""int64""" if self.config_name != """sts-b""" else """float32"""  ),
									        }  )				,			codebase_urls=[]				,			reference_urls=[]				,			format="""numpy"""				,			)






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase  ):
									return {"accuracy": simple_accuracy(__UpperCamelCase				,			__UpperCamelCase  )}