




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


lowerCAmelCase_					    = False






@skip_mps
class  _lowerCAmelCase     (  _lowercase    ,			_lowercase    ,			_lowercase    ,			unittest.TestCase	):
      A__		     =		StableDiffusionAttendAndExcitePipeline
      A__		     =		False
      A__		     =		TEXT_TO_IMAGE_PARAMS
      A__		     =		TEXT_TO_IMAGE_BATCH_PARAMS.union({'token_indices'}	)
      A__		     =		TEXT_TO_IMAGE_IMAGE_PARAMS
      A__		     =		TEXT_TO_IMAGE_IMAGE_PARAMS

      @classmethod
      def 	__magic_name__(   cls		):
        super().setUpClass()
        torch.use_deterministic_algorithms(__UpperCAmelCase		)

      @classmethod
      def 	__magic_name__(   cls		):
        super().tearDownClass()
        torch.use_deterministic_algorithms(__UpperCAmelCase		)

      def 	__magic_name__(   self		):
        torch.manual_seed(0		)
        lowerCAmelCase__				:  Optional[int]		       =     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=__UpperCAmelCase      ,				)
        lowerCAmelCase__				:  List[str]		       =     DDIMScheduler(
            beta_start=0.00085      ,				beta_end=0.012      ,				beta_schedule='''scaled_linear'''      ,				clip_sample=__UpperCAmelCase      ,				set_alpha_to_one=__UpperCAmelCase      ,				)
        torch.manual_seed(0		)
        lowerCAmelCase__				:  List[str]		       =     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		)
        lowerCAmelCase__				:  str		       =     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=1000      ,				hidden_act='''gelu'''      ,				projection_dim=512      ,				)
        lowerCAmelCase__				:  Optional[Any]		       =     CLIPTextModel(__UpperCAmelCase		)
        lowerCAmelCase__				:  Optional[Any]		       =     CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

        lowerCAmelCase__				:  List[str]		       =     {
            '''unet''': unet,
            '''scheduler''': scheduler,
            '''vae''': vae,
            '''text_encoder''': text_encoder,
            '''tokenizer''': tokenizer,
            '''safety_checker''': None,
            '''feature_extractor''': None,
        }

        return components

      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase=0		):
        if str(__UpperCAmelCase		).startswith('''mps'''		):
          lowerCAmelCase__				:  str		       =     torch.manual_seed(__UpperCAmelCase		)
        else:
          lowerCAmelCase__				:  Union[str, Any]		       =     torch.Generator(device=__UpperCAmelCase		).manual_seed(__UpperCAmelCase		)
        lowerCAmelCase__				:  Union[str, Any]		       =     {
            '''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 	__magic_name__(   self		):
        lowerCAmelCase__				:  Any		       =     '''cpu'''

        lowerCAmelCase__				:  Optional[int]		       =     self.get_dummy_components()
        lowerCAmelCase__				:  Tuple		       =     self.pipeline_class(**__UpperCAmelCase		)
        pipe.to(__UpperCAmelCase		)
        pipe.set_progress_bar_config(disable=__UpperCAmelCase		)

        lowerCAmelCase__				:  Optional[Any]		       =     self.get_dummy_inputs(__UpperCAmelCase		)
        lowerCAmelCase__				:  Any		       =     pipe(**__UpperCAmelCase		).images
        lowerCAmelCase__				:  int		       =     image[0, -3:, -3:, -1]

        self.assertEqual(image.shape      ,				(1, 64, 64, 3)		)
        lowerCAmelCase__				:  List[str]		       =     np.array(
            [0.63905364, 0.62897307, 0.48599017, 0.5133624, 0.5550048, 0.45769516, 0.50326973, 0.5023139, 0.45384496]		)
        lowerCAmelCase__				:  Tuple		       =     np.abs(image_slice.flatten() - expected_slice		).max()
        self.assertLessEqual(__UpperCAmelCase      ,				1e-3		)

      def 	__magic_name__(   self		):
        super().test_cpu_offload_forward_pass(expected_max_diff=5e-4		)

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

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

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

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

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







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






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

      @classmethod
      def 	__magic_name__(   cls		):
        super().setUpClass()
        torch.use_deterministic_algorithms(__UpperCAmelCase		)

      @classmethod
      def 	__magic_name__(   cls		):
        super().tearDownClass()
        torch.use_deterministic_algorithms(__UpperCAmelCase		)

      def 	__magic_name__(   self		):
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Union[str, Any]		       =     torch.manual_seed(51		)

        lowerCAmelCase__				:  Optional[Any]		       =     StableDiffusionAttendAndExcitePipeline.from_pretrained(
            '''CompVis/stable-diffusion-v1-4'''      ,				safety_checker=__UpperCAmelCase      ,				torch_dtype=torch.floataa		)
        pipe.to('''cuda'''		)

        lowerCAmelCase__				:  Any		       =     '''a painting of an elephant with glasses'''
        lowerCAmelCase__				:  Dict		       =     [5, 7]

        lowerCAmelCase__				:  Optional[Any]		       =     pipe(
            prompt=__UpperCAmelCase      ,				token_indices=__UpperCAmelCase      ,				guidance_scale=7.5      ,				generator=__UpperCAmelCase      ,				num_inference_steps=5      ,				max_iter_to_alter=5      ,				output_type='''numpy'''      ,				).images[0]

        lowerCAmelCase__				:  Tuple		       =     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







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


lowerCAmelCase_					    = {
    """configuration_longformer""": [
        """LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """LongformerConfig""",
        """LongformerOnnxConfig""",
    ],
    """tokenization_longformer""": ["""LongformerTokenizer"""],
}

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowerCAmelCase_					    = ["""LongformerTokenizerFast"""]

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowerCAmelCase_					    = [
         """LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """LongformerForMaskedLM""",
         """LongformerForMultipleChoice""",
         """LongformerForQuestionAnswering""",
         """LongformerForSequenceClassification""",
         """LongformerForTokenClassification""",
         """LongformerModel""",
         """LongformerPreTrainedModel""",
         """LongformerSelfAttention""",
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowerCAmelCase_					    = [
         """TF_LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """TFLongformerForMaskedLM""",
         """TFLongformerForMultipleChoice""",
         """TFLongformerForQuestionAnswering""",
         """TFLongformerForSequenceClassification""",
         """TFLongformerForTokenClassification""",
         """TFLongformerModel""",
         """TFLongformerPreTrainedModel""",
         """TFLongformerSelfAttention""",
     ]


if TYPE_CHECKING:
     from .configuration_longformer import (
         LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
         LongformerConfig,
         LongformerOnnxConfig,
     )
     from .tokenization_longformer import LongformerTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_longformer_fast import LongformerTokenizerFast

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_longformer import (
              LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
              LongformerForMaskedLM,
              LongformerForMultipleChoice,
              LongformerForQuestionAnswering,
              LongformerForSequenceClassification,
              LongformerForTokenClassification,
              LongformerModel,
              LongformerPreTrainedModel,
              LongformerSelfAttention,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_longformer import (
              TF_LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFLongformerForMaskedLM,
              TFLongformerForMultipleChoice,
              TFLongformerForQuestionAnswering,
              TFLongformerForSequenceClassification,
              TFLongformerForTokenClassification,
              TFLongformerModel,
              TFLongformerPreTrainedModel,
              TFLongformerSelfAttention,
          )

else:
     import sys

     lowerCAmelCase_					    = _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

