





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__      :     List[Any]         =		False





@skip_mps
class UpperCAmelCase_				( _SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE     ,		unittest.TestCase ):
  __lowerCAmelCase   :						str							=   StableDiffusionAttendAndExcitePipeline
  __lowerCAmelCase   :						Dict							=   False
  __lowerCAmelCase   :						Dict							=   TEXT_TO_IMAGE_PARAMS
  __lowerCAmelCase   :						Optional[Any]							=   TEXT_TO_IMAGE_BATCH_PARAMS.union({"""token_indices"""} )
  __lowerCAmelCase   :						Optional[int]							=   TEXT_TO_IMAGE_IMAGE_PARAMS
  __lowerCAmelCase   :						List[Any]							=   TEXT_TO_IMAGE_IMAGE_PARAMS


  @classmethod
  def 						__UpperCAmelCase     (							cls ):



       """simple docstring"""





       super().setUpClass()
       torch.use_deterministic_algorithms(__snake_case )


  @classmethod
  def 						__UpperCAmelCase     (							cls ):



       """simple docstring"""





       super().tearDownClass()
       torch.use_deterministic_algorithms(__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       torch.manual_seed(0 )
       A_          =       UNetaDConditionModel(
           block_out_channels=(3_2, 6_4) ,layers_per_block=1 ,sample_size=3_2 ,in_channels=4 ,out_channels=4 ,down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''') ,up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''') ,cross_attention_dim=3_2 ,attention_head_dim=(2, 4) ,use_linear_projection=__snake_case ,)
       A_          =       DDIMScheduler(
           beta_start=0.0_0085 ,beta_end=0.012 ,beta_schedule='''scaled_linear''' ,clip_sample=__snake_case ,set_alpha_to_one=__snake_case ,)
       torch.manual_seed(0 )
       A_          =       AutoencoderKL(
           block_out_channels=[3_2, 6_4] ,in_channels=3 ,out_channels=3 ,down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''] ,up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''] ,latent_channels=4 ,sample_size=1_2_8 ,)
       torch.manual_seed(0 )
       A_          =       CLIPTextConfig(
           bos_token_id=0 ,eos_token_id=2 ,hidden_size=3_2 ,intermediate_size=3_7 ,layer_norm_eps=1E-05 ,num_attention_heads=4 ,num_hidden_layers=5 ,pad_token_id=1 ,vocab_size=1_0_0_0 ,hidden_act='''gelu''' ,projection_dim=5_1_2 ,)
       A_          =       CLIPTextModel(__snake_case )
       A_          =       CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' )

       A_          =       {
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''safety_checker''': None,
           '''feature_extractor''': None,
       }

       return components


  def 						__UpperCAmelCase     (							self ,__snake_case ,__snake_case=0 ):



       """simple docstring"""





       if str(__snake_case ).startswith('''mps''' ):
            A_          =       torch.manual_seed(__snake_case )
       else:
            A_          =       torch.Generator(device=__snake_case ).manual_seed(__snake_case )
       A_          =       A_          =       {
           '''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 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       '''cpu'''

       A_          =       self.get_dummy_components()
       A_          =       self.pipeline_class(**__snake_case )
       pipe.to(__snake_case )
       pipe.set_progress_bar_config(disable=__snake_case )

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

       self.assertEqual(image.shape ,(1, 6_4, 6_4, 3) )
       A_          =       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] )
       A_          =       np.abs(image_slice.flatten() - expected_slice ).max()
       self.assertLessEqual(__snake_case ,1E-3 )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       super().test_cpu_offload_forward_pass(expected_max_diff=5E-4 )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       self._test_inference_batch_consistent(batch_sizes=[1, 2] )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       self._test_inference_batch_single_identical(batch_size=2 ,expected_max_diff=7E-4 )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3 )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5E-4 )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       super().test_save_load_local(expected_max_difference=5E-4 )




  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       super().test_save_load_optional_components(expected_max_difference=4E-4 )





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


  @classmethod
  def 						__UpperCAmelCase     (							cls ):



       """simple docstring"""





       super().setUpClass()
       torch.use_deterministic_algorithms(__snake_case )


  @classmethod
  def 						__UpperCAmelCase     (							cls ):



       """simple docstring"""





       super().tearDownClass()
       torch.use_deterministic_algorithms(__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()




  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       torch.manual_seed(5_1 )

       A_          =       StableDiffusionAttendAndExcitePipeline.from_pretrained(
           '''CompVis/stable-diffusion-v1-4''' ,safety_checker=__snake_case ,torch_dtype=torch.floataa )
       pipe.to('''cuda''' )

       A_          =       '''a painting of an elephant with glasses'''
       A_          =       [5, 7]

       A_          =       pipe(
           prompt=__snake_case ,token_indices=__snake_case ,guidance_scale=7.5 ,generator=__snake_case ,num_inference_steps=5 ,max_iter_to_alter=5 ,output_type='''numpy''' ,).images[0]

       A_          =       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_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


a__      :     str         =		{
    'configuration_roformer': ['ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP', 'RoFormerConfig', 'RoFormerOnnxConfig'],
    'tokenization_roformer': ['RoFormerTokenizer'],
}

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      a__      :     List[Any]         =		['RoFormerTokenizerFast']

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      a__      :     List[str]         =		[
          'ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
          'RoFormerForCausalLM',
          'RoFormerForMaskedLM',
          'RoFormerForMultipleChoice',
          'RoFormerForQuestionAnswering',
          'RoFormerForSequenceClassification',
          'RoFormerForTokenClassification',
          'RoFormerLayer',
          'RoFormerModel',
          'RoFormerPreTrainedModel',
          'load_tf_weights_in_roformer',
      ]


try:
      if not is_tf_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      a__      :     List[Any]         =		[
          'TF_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
          'TFRoFormerForCausalLM',
          'TFRoFormerForMaskedLM',
          'TFRoFormerForMultipleChoice',
          'TFRoFormerForQuestionAnswering',
          'TFRoFormerForSequenceClassification',
          'TFRoFormerForTokenClassification',
          'TFRoFormerLayer',
          'TFRoFormerModel',
          'TFRoFormerPreTrainedModel',
      ]


try:
      if not is_flax_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      a__      :     Dict         =		[
          'FLAX_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
          'FlaxRoFormerForMaskedLM',
          'FlaxRoFormerForMultipleChoice',
          'FlaxRoFormerForQuestionAnswering',
          'FlaxRoFormerForSequenceClassification',
          'FlaxRoFormerForTokenClassification',
          'FlaxRoFormerModel',
          'FlaxRoFormerPreTrainedModel',
      ]


if TYPE_CHECKING:
      from .configuration_roformer import ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, RoFormerConfig, RoFormerOnnxConfig
      from .tokenization_roformer import RoFormerTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_roformer_fast import RoFormerTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_roformer import (
                ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
                RoFormerForCausalLM,
                RoFormerForMaskedLM,
                RoFormerForMultipleChoice,
                RoFormerForQuestionAnswering,
                RoFormerForSequenceClassification,
                RoFormerForTokenClassification,
                RoFormerLayer,
                RoFormerModel,
                RoFormerPreTrainedModel,
                load_tf_weights_in_roformer,
            )

      try:
            if not is_tf_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_roformer import (
                TF_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
                TFRoFormerForCausalLM,
                TFRoFormerForMaskedLM,
                TFRoFormerForMultipleChoice,
                TFRoFormerForQuestionAnswering,
                TFRoFormerForSequenceClassification,
                TFRoFormerForTokenClassification,
                TFRoFormerLayer,
                TFRoFormerModel,
                TFRoFormerPreTrainedModel,
            )

      try:
            if not is_flax_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_flax_roformer import (
                FLAX_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
                FlaxRoFormerForMaskedLM,
                FlaxRoFormerForMultipleChoice,
                FlaxRoFormerForQuestionAnswering,
                FlaxRoFormerForSequenceClassification,
                FlaxRoFormerForTokenClassification,
                FlaxRoFormerModel,
                FlaxRoFormerPreTrainedModel,
            )


else:
      import sys

      a__      :     Union[str, Any]         =		_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



