"""simple docstring"""


import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionAttendAndExcitePipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_numpy, skip_mps, slow
from diffusers.utils.testing_utils import require_torch_gpu

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


UpperCAmelCase__  		=						False


@skip_mps
class        a				(		lowerCAmelCase_	,							lowerCAmelCase_	,							lowerCAmelCase_	,							unittest.TestCase   ):
    _snake_case     :  int     							=							StableDiffusionAttendAndExcitePipeline
    _snake_case     :  Optional[Any]     							=							False
    _snake_case     :  Union[str, Any]     							=							TEXT_TO_IMAGE_PARAMS
    _snake_case     :  Optional[Any]     							=							TEXT_TO_IMAGE_BATCH_PARAMS.union({'token_indices'}   )
    _snake_case     :  Optional[int]     							=							TEXT_TO_IMAGE_IMAGE_PARAMS
    _snake_case     :  Dict     							=							TEXT_TO_IMAGE_IMAGE_PARAMS


    @classmethod
    def      lowerCAmelCase_  (					cls	:  Any					):
      super().setUpClass()
      torch.use_deterministic_algorithms(__lowerCAmelCase					)


    @classmethod
    def      lowerCAmelCase_  (					cls	:  Optional[Any]					):
      super().tearDownClass()
      torch.use_deterministic_algorithms(__lowerCAmelCase					)


    def      lowerCAmelCase_  (					self	:  Optional[int]					):
      torch.manual_seed(0					)
      _UpperCAmelCase						 =      UNetaDConditionModel(
          block_out_channels=(32, 64)   ,				layers_per_block=1   ,				sample_size=32   ,				in_channels=4   ,				out_channels=4   ,				down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")   ,				up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")   ,				cross_attention_dim=32   ,				attention_head_dim=(2, 4)   ,				use_linear_projection=__lowerCAmelCase   ,				)
      _UpperCAmelCase						 =      DDIMScheduler(
          beta_start=0.00_085   ,				beta_end=0.012   ,				beta_schedule="""scaled_linear"""   ,				clip_sample=__lowerCAmelCase   ,				set_alpha_to_one=__lowerCAmelCase   ,				)
      torch.manual_seed(0					)
      _UpperCAmelCase						 =      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   ,				sample_size=128   ,				)
      torch.manual_seed(0					)
      _UpperCAmelCase						 =      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   ,				hidden_act="""gelu"""   ,				projection_dim=512   ,				)
      _UpperCAmelCase						 =      CLIPTextModel(__lowerCAmelCase					)
      _UpperCAmelCase						 =      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""					)

      _UpperCAmelCase						 =      {
          """unet""": unet,
          """scheduler""": scheduler,
          """vae""": vae,
          """text_encoder""": text_encoder,
          """tokenizer""": tokenizer,
          """safety_checker""": None,
          """feature_extractor""": None,
      }

      return components


    def      lowerCAmelCase_  (					self	:  int   ,				__lowerCAmelCase	:  Union[str, Any]   ,				__lowerCAmelCase	:  Tuple=0					):
      if str(__lowerCAmelCase					).startswith("""mps"""					):
        _UpperCAmelCase						 =      torch.manual_seed(__lowerCAmelCase					)
      else:
        _UpperCAmelCase						 =      torch.Generator(device=__lowerCAmelCase					).manual_seed(__lowerCAmelCase					)
      _UpperCAmelCase						 =      _UpperCAmelCase						 =      {
          """prompt""": """a cat and a frog""",
          """token_indices""": [2, 5],
          """generator""": generator,
          """num_inference_steps""": 1,
          """guidance_scale""": 6.0,
          """output_type""": """numpy""",
          """max_iter_to_alter""": 2,
          """thresholds""": {0: 0.7},
      }
      return inputs


    def      lowerCAmelCase_  (					self	:  Dict					):
      _UpperCAmelCase						 =      """cpu"""

      _UpperCAmelCase						 =      self.get_dummy_components()
      _UpperCAmelCase						 =      self.pipeline_class(**__lowerCAmelCase					)
      pipe.to(__lowerCAmelCase					)
      pipe.set_progress_bar_config(disable=__lowerCAmelCase					)

      _UpperCAmelCase						 =      self.get_dummy_inputs(__lowerCAmelCase					)
      _UpperCAmelCase						 =      pipe(**__lowerCAmelCase					).images
      _UpperCAmelCase						 =      image[0, -3:, -3:, -1]

      self.assertEqual(image.shape   ,				(1, 64, 64, 3)					)
      _UpperCAmelCase						 =      np.array(
          [0.63_905_364, 0.62_897_307, 0.48_599_017, 0.5_133_624, 0.5_550_048, 0.45_769_516, 0.50_326_973, 0.5_023_139, 0.45_384_496]					)
      _UpperCAmelCase						 =      np.abs(image_slice.flatten() - expected_slice					).max()
      self.assertLessEqual(__lowerCAmelCase   ,				1e-3					)


    def      lowerCAmelCase_  (					self	:  Union[str, Any]					):
      super().test_cpu_offload_forward_pass(expected_max_diff=5e-4					)


    def      lowerCAmelCase_  (					self	:  str					):
      # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
      self._test_inference_batch_consistent(batch_sizes=[1, 2]					)


    def      lowerCAmelCase_  (					self	:  Tuple					):
      self._test_inference_batch_single_identical(batch_size=2   ,				expected_max_diff=7e-4					)


    def      lowerCAmelCase_  (					self	:  Tuple					):
      super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3					)


    def      lowerCAmelCase_  (					self	:  Dict					):
      super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5e-4					)


    def      lowerCAmelCase_  (					self	:  int					):
      super().test_save_load_local(expected_max_difference=5e-4					)


    def      lowerCAmelCase_  (					self	:  int					):
      super().test_save_load_optional_components(expected_max_difference=4e-4					)



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


    @classmethod
    def      lowerCAmelCase_  (					cls	:  List[Any]					):
      super().setUpClass()
      torch.use_deterministic_algorithms(__lowerCAmelCase					)


    @classmethod
    def      lowerCAmelCase_  (					cls	:  Optional[int]					):
      super().tearDownClass()
      torch.use_deterministic_algorithms(__lowerCAmelCase					)


    def      lowerCAmelCase_  (					self	:  List[Any]					):
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()


    def      lowerCAmelCase_  (					self	:  Dict					):
      _UpperCAmelCase						 =      torch.manual_seed(51					)

      _UpperCAmelCase						 =      StableDiffusionAttendAndExcitePipeline.from_pretrained(
          """CompVis/stable-diffusion-v1-4"""   ,				safety_checker=__lowerCAmelCase   ,				torch_dtype=torch.floataa					)
      pipe.to("""cuda"""					)

      _UpperCAmelCase						 =      """a painting of an elephant with glasses"""
      _UpperCAmelCase						 =      [5, 7]

      _UpperCAmelCase						 =      pipe(
          prompt=__lowerCAmelCase   ,				token_indices=__lowerCAmelCase   ,				guidance_scale=7.5   ,				generator=__lowerCAmelCase   ,				num_inference_steps=5   ,				max_iter_to_alter=5   ,				output_type="""numpy"""   ,				).images[0]

      _UpperCAmelCase						 =      load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/attend-and-excite/elephant_glasses.npy"""					)
      assert np.abs((expected_image - image).max()					) < 5e-1

"""simple docstring"""


import datasets


UpperCAmelCase__  		=						"""\
@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\",
}
"""

UpperCAmelCase__  		=						"""\
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).
"""

UpperCAmelCase__  		=						"""
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					(			lowercase   ,lowercase		):







  """simple docstring"""
  return (preds == labels).mean()


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


    def      lowerCAmelCase_  (					self	:  Union[str, Any]					):
      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      lowerCAmelCase_  (					self	:  List[str]   ,				__lowerCAmelCase	:  int   ,				__lowerCAmelCase	:  Any					):
      return {"accuracy": simple_accuracy(__lowerCAmelCase   ,				__lowerCAmelCase					)}
