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


A     :  int					     =       False


@skip_mps
class   lowerCamelCase					(SCREAMING_SNAKE_CASE__  ,      SCREAMING_SNAKE_CASE__  ,      SCREAMING_SNAKE_CASE__  ,      unittest.TestCase					):
  """simple docstring"""







  lowerCamelCase__     =						StableDiffusionAttendAndExcitePipeline
  lowerCamelCase__     =						False
  lowerCamelCase__     =						TEXT_TO_IMAGE_PARAMS
  lowerCamelCase__     =						TEXT_TO_IMAGE_BATCH_PARAMS.union({'''token_indices'''}					)
  lowerCamelCase__     =						TEXT_TO_IMAGE_IMAGE_PARAMS
  lowerCamelCase__     =						TEXT_TO_IMAGE_IMAGE_PARAMS







  @classmethod
  def    __A     (							cls :       Any   )  ->	List[Any]:
        super().setUpClass()
        torch.use_deterministic_algorithms(__magic_name__   )







  @classmethod
  def    __A     (							cls :       List[Any]   )  ->	Dict:
        super().tearDownClass()
        torch.use_deterministic_algorithms(__magic_name__   )







  def    __A     (							self :       int   )  ->	int:
        torch.manual_seed(0   )
        SCREAMING_SNAKE_CASE_          =  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=__magic_name__  ,     )
        SCREAMING_SNAKE_CASE_          =  DDIMScheduler(
            beta_start=0.0_0085  ,     beta_end=0.012  ,     beta_schedule="scaled_linear"  ,     clip_sample=__magic_name__  ,     set_alpha_to_one=__magic_name__  ,     )
        torch.manual_seed(0   )
        SCREAMING_SNAKE_CASE_          =  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   )
        SCREAMING_SNAKE_CASE_          =  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=1_000  ,     hidden_act="gelu"  ,     projection_dim=512  ,     )
        SCREAMING_SNAKE_CASE_          =  CLIPTextModel(__magic_name__   )
        SCREAMING_SNAKE_CASE_          =  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"   )

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

        return components







  def    __A     (							self :       List[str]  ,     __magic_name__ :       Optional[Any]  ,     __magic_name__ :       str=0   )  ->	List[Any]:
        if str(__magic_name__   ).startswith("mps"   ):
              SCREAMING_SNAKE_CASE_          =  torch.manual_seed(__magic_name__   )
        else:
              SCREAMING_SNAKE_CASE_          =  torch.Generator(device=__magic_name__   ).manual_seed(__magic_name__   )
        SCREAMING_SNAKE_CASE_          =  SCREAMING_SNAKE_CASE_          =  {
            "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    __A     (							self :       Any   )  ->	Any:
        SCREAMING_SNAKE_CASE_          =  "cpu"

        SCREAMING_SNAKE_CASE_          =  self.get_dummy_components()
        SCREAMING_SNAKE_CASE_          =  self.pipeline_class(**__magic_name__   )
        pipe.to(__magic_name__   )
        pipe.set_progress_bar_config(disable=__magic_name__   )

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

        self.assertEqual(image.shape  ,     (1, 64, 64, 3)   )
        SCREAMING_SNAKE_CASE_          =  np.array(
            [0.6390_5364, 0.6289_7307, 0.4859_9017, 0.513_3624, 0.555_0048, 0.4576_9516, 0.5032_6973, 0.502_3139, 0.4538_4496]   )
        SCREAMING_SNAKE_CASE_          =  np.abs(image_slice.flatten() - expected_slice   ).max()
        self.assertLessEqual(__magic_name__  ,     1e-3   )







  def    __A     (							self :       Optional[int]   )  ->	Any:
        super().test_cpu_offload_forward_pass(expected_max_diff=5e-4   )







  def    __A     (							self :       Dict   )  ->	List[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    __A     (							self :       List[str]   )  ->	str:
        self._test_inference_batch_single_identical(batch_size=2  ,     expected_max_diff=7e-4   )







  def    __A     (							self :       Tuple   )  ->	int:
        super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3   )







  def    __A     (							self :       int   )  ->	str:
        super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5e-4   )







  def    __A     (							self :       int   )  ->	Union[str, Any]:
        super().test_save_load_local(expected_max_difference=5e-4   )







  def    __A     (							self :       Any   )  ->	str:
        super().test_save_load_optional_components(expected_max_difference=4e-4   )



@require_torch_gpu
@slow
class   lowerCamelCase					(unittest.TestCase					):
  """simple docstring"""







  @classmethod
  def    __A     (							cls :       List[Any]   )  ->	Any:
        super().setUpClass()
        torch.use_deterministic_algorithms(__magic_name__   )







  @classmethod
  def    __A     (							cls :       Optional[int]   )  ->	List[Any]:
        super().tearDownClass()
        torch.use_deterministic_algorithms(__magic_name__   )







  def    __A     (							self :       Optional[Any]   )  ->	List[Any]:
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







  def    __A     (							self :       Union[str, Any]   )  ->	Any:
        SCREAMING_SNAKE_CASE_          =  torch.manual_seed(51   )

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

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

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

        SCREAMING_SNAKE_CASE_          =  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



import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow


if is_torch_available():
      import torch

      from transformers import XLMRobertaModel


@require_sentencepiece
@require_tokenizers
@require_torch
class   lowerCamelCase					(unittest.TestCase					):
  """simple docstring"""







  @slow
  def    __A     (							self :       Dict   )  ->	Any:
        SCREAMING_SNAKE_CASE_          =  XLMRobertaModel.from_pretrained("xlm-roberta-base"   )
        SCREAMING_SNAKE_CASE_          =  torch.tensor([[0, 581, 10_269, 83, 99_942, 136, 60_742, 23, 70, 80_583, 18_276, 2]]   )
        # The dog is cute and lives in the garden house

        SCREAMING_SNAKE_CASE_          =  torch.Size((1, 12, 768)   )  # batch_size, sequence_length, embedding_vector_dim
        SCREAMING_SNAKE_CASE_          =  torch.tensor(
            [[-0.0101, 0.1218, -0.0803, 0.0801, 0.1327, 0.0776, -0.1215, 0.2383, 0.3338, 0.3106, 0.0300, 0.0252]]   )
        #  xlmr = torch.hub.load('pytorch/fairseq', 'xlmr.base')
        #  xlmr.eval()
        #  expected_output_values_last_dim = xlmr.extract_features(input_ids[0])[:, :, -1]
        with torch.no_grad():
              SCREAMING_SNAKE_CASE_          =  model(__magic_name__   )["last_hidden_state"].detach()
        self.assertEqual(output.shape  ,     __magic_name__   )
        # compare the actual values for a slice of last dim
        self.assertTrue(torch.allclose(output[:, :, -1]  ,     __magic_name__  ,     atol=1e-3   )   )







  @slow
  def    __A     (							self :       Tuple   )  ->	List[str]:
        SCREAMING_SNAKE_CASE_          =  XLMRobertaModel.from_pretrained("xlm-roberta-large"   )
        SCREAMING_SNAKE_CASE_          =  torch.tensor([[0, 581, 10_269, 83, 99_942, 136, 60_742, 23, 70, 80_583, 18_276, 2]]   )
        # The dog is cute and lives in the garden house

        SCREAMING_SNAKE_CASE_          =  torch.Size((1, 12, 1_024)   )  # batch_size, sequence_length, embedding_vector_dim
        SCREAMING_SNAKE_CASE_          =  torch.tensor(
            [[-0.0699, -0.0318, 0.0705, -0.1241, 0.0999, -0.0520, 0.1004, -0.1838, -0.4704, 0.1437, 0.0821, 0.0126]]   )
        #  xlmr = torch.hub.load('pytorch/fairseq', 'xlmr.large')
        #  xlmr.eval()
        #  expected_output_values_last_dim = xlmr.extract_features(input_ids[0])[:, :, -1]
        with torch.no_grad():
              SCREAMING_SNAKE_CASE_          =  model(__magic_name__   )["last_hidden_state"].detach()
        self.assertEqual(output.shape  ,     __magic_name__   )
        # compare the actual values for a slice of last dim
        self.assertTrue(torch.allclose(output[:, :, -1]  ,     __magic_name__  ,     atol=1e-3   )   )


