"""simple docstring"""


import gc
import unittest

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

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


class        a				(		lowerCAmelCase_	,							unittest.TestCase   ):
    _snake_case     :  Dict     							=							ShapEPipeline
    _snake_case     :  Tuple     							=							['prompt']
    _snake_case     :  Union[str, Any]     							=							['prompt']
    _snake_case     :  Tuple     							=							[
        'num_images_per_prompt',
        'num_inference_steps',
        'generator',
        'latents',
        'guidance_scale',
        'frame_size',
        'output_type',
        'return_dict',
    ]
    _snake_case     :  Tuple     							=							False


    @property
    def      lowerCAmelCase_  (					self	:  int					):
      return 32


    @property
    def      lowerCAmelCase_  (					self	:  Dict					):
      return 32


    @property
    def      lowerCAmelCase_  (					self	:  str					):
      return self.time_input_dim * 4


    @property
    def      lowerCAmelCase_  (					self	:  List[Any]					):
      return 8


    @property
    def      lowerCAmelCase_  (					self	:  Optional[Any]					):
      _UpperCAmelCase						 =      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""					)
      return tokenizer


    @property
    def      lowerCAmelCase_  (					self	:  Union[str, Any]					):
      torch.manual_seed(0					)
      _UpperCAmelCase						 =      CLIPTextConfig(
          bos_token_id=0   ,				eos_token_id=2   ,				hidden_size=self.text_embedder_hidden_size   ,				projection_dim=self.text_embedder_hidden_size   ,				intermediate_size=37   ,				layer_norm_eps=1e-0_5   ,				num_attention_heads=4   ,				num_hidden_layers=5   ,				pad_token_id=1   ,				vocab_size=1000   ,				)
      return CLIPTextModelWithProjection(__lowerCAmelCase					)


    @property
    def      lowerCAmelCase_  (					self	:  Union[str, Any]					):
      torch.manual_seed(0					)

      _UpperCAmelCase						 =      {
          """num_attention_heads""": 2,
          """attention_head_dim""": 16,
          """embedding_dim""": self.time_input_dim,
          """num_embeddings""": 32,
          """embedding_proj_dim""": self.text_embedder_hidden_size,
          """time_embed_dim""": self.time_embed_dim,
          """num_layers""": 1,
          """clip_embed_dim""": self.time_input_dim * 2,
          """additional_embeddings""": 0,
          """time_embed_act_fn""": """gelu""",
          """norm_in_type""": """layer""",
          """encoder_hid_proj_type""": None,
          """added_emb_type""": None,
      }

      _UpperCAmelCase						 =      PriorTransformer(**__lowerCAmelCase					)
      return model


    @property
    def      lowerCAmelCase_  (					self	:  List[Any]					):
      torch.manual_seed(0					)

      _UpperCAmelCase						 =      {
          """param_shapes""": (
              (self.renderer_dim, 93),
              (self.renderer_dim, 8),
              (self.renderer_dim, 8),
              (self.renderer_dim, 8),
          ),
          """d_latent""": self.time_input_dim,
          """d_hidden""": self.renderer_dim,
          """n_output""": 12,
          """background""": (
              0.1,
              0.1,
              0.1,
          ),
      }
      _UpperCAmelCase						 =      ShapERenderer(**__lowerCAmelCase					)
      return model


    def      lowerCAmelCase_  (					self	:  Union[str, Any]					):
      _UpperCAmelCase						 =      self.dummy_prior
      _UpperCAmelCase						 =      self.dummy_text_encoder
      _UpperCAmelCase						 =      self.dummy_tokenizer
      _UpperCAmelCase						 =      self.dummy_renderer

      _UpperCAmelCase						 =      HeunDiscreteScheduler(
          beta_schedule="""exp"""   ,				num_train_timesteps=1024   ,				prediction_type="""sample"""   ,				use_karras_sigmas=__lowerCAmelCase   ,				clip_sample=__lowerCAmelCase   ,				clip_sample_range=1.0   ,				)
      _UpperCAmelCase						 =      {
          """prior""": prior,
          """text_encoder""": text_encoder,
          """tokenizer""": tokenizer,
          """renderer""": renderer,
          """scheduler""": scheduler,
      }

      return components


    def      lowerCAmelCase_  (					self	:  List[Any]   ,				__lowerCAmelCase	:  Optional[int]   ,				__lowerCAmelCase	:  Tuple=0					):
      if str(__lowerCAmelCase					).startswith("""mps"""					):
        _UpperCAmelCase						 =      torch.manual_seed(__lowerCAmelCase					)
      else:
        _UpperCAmelCase						 =      torch.Generator(device=__lowerCAmelCase					).manual_seed(__lowerCAmelCase					)
      _UpperCAmelCase						 =      {
          """prompt""": """horse""",
          """generator""": generator,
          """num_inference_steps""": 1,
          """frame_size""": 32,
          """output_type""": """np""",
      }
      return inputs


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

      _UpperCAmelCase						 =      self.get_dummy_components()

      _UpperCAmelCase						 =      self.pipeline_class(**__lowerCAmelCase					)
      _UpperCAmelCase						 =      pipe.to(__lowerCAmelCase					)

      pipe.set_progress_bar_config(disable=__lowerCAmelCase					)

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

      assert image.shape == (20, 32, 32, 3)

      _UpperCAmelCase						 =      np.array(
          [
              0.00_039_216,
              0.00_039_216,
              0.00_039_216,
              0.00_039_216,
              0.00_039_216,
              0.00_039_216,
              0.00_039_216,
              0.00_039_216,
              0.00_039_216,
          ]					)

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


    def      lowerCAmelCase_  (					self	:  Any					):
      # 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	:  Optional[Any]					):
      _UpperCAmelCase						 =      torch_device == """cpu"""
      _UpperCAmelCase						 =      True

      self._test_inference_batch_single_identical(
          batch_size=2   ,				test_max_difference=__lowerCAmelCase   ,				relax_max_difference=__lowerCAmelCase   ,				)


    def      lowerCAmelCase_  (					self	:  Optional[Any]					):
      _UpperCAmelCase						 =      self.get_dummy_components()
      _UpperCAmelCase						 =      self.pipeline_class(**__lowerCAmelCase					)
      _UpperCAmelCase						 =      pipe.to(__lowerCAmelCase					)
      pipe.set_progress_bar_config(disable=__lowerCAmelCase					)

      _UpperCAmelCase						 =      1
      _UpperCAmelCase						 =      2

      _UpperCAmelCase						 =      self.get_dummy_inputs(__lowerCAmelCase					)

      for key in inputs.keys():
        if key in self.batch_params:
          _UpperCAmelCase						 =      batch_size * [inputs[key]]

      _UpperCAmelCase						 =      pipe(**__lowerCAmelCase   ,				num_images_per_prompt=__lowerCAmelCase					)[0]

      assert images.shape[0] == batch_size * num_images_per_prompt



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


    def      lowerCAmelCase_  (					self	:  int					):
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()


    def      lowerCAmelCase_  (					self	:  Optional[Any]					):
      _UpperCAmelCase						 =      load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
          """/shap_e/test_shap_e_np_out.npy"""					)
      _UpperCAmelCase						 =      ShapEPipeline.from_pretrained("""openai/shap-e"""					)
      _UpperCAmelCase						 =      pipe.to(__lowerCAmelCase					)
      pipe.set_progress_bar_config(disable=__lowerCAmelCase					)

      _UpperCAmelCase						 =      torch.Generator(device=__lowerCAmelCase					).manual_seed(0					)

      _UpperCAmelCase						 =      pipe(
          """a shark"""   ,				generator=__lowerCAmelCase   ,				guidance_scale=15.0   ,				num_inference_steps=64   ,				frame_size=64   ,				output_type="""np"""   ,				).images[0]

      assert images.shape == (20, 64, 64, 3)

      assert_mean_pixel_difference(__lowerCAmelCase   ,				__lowerCAmelCase					)

"""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					)}
