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_							=      False


@skip_mps
class 			UpperCAmelCase__    (		snake_case		,	snake_case		,	snake_case		,	unittest.TestCase					):



  """simple docstring"""






  lowerCAmelCase__		:	Optional[int]    		= StableDiffusionAttendAndExcitePipeline
  lowerCAmelCase__		:	Any    		= False
  lowerCAmelCase__		:	Optional[Any]    		= TEXT_TO_IMAGE_PARAMS
  lowerCAmelCase__		:	List[Any]    		= TEXT_TO_IMAGE_BATCH_PARAMS.union({'token_indices'}					)
  lowerCAmelCase__		:	int    		= TEXT_TO_IMAGE_IMAGE_PARAMS
  lowerCAmelCase__		:	str    		= TEXT_TO_IMAGE_IMAGE_PARAMS






  @classmethod
  def 	_UpperCAmelCase				(				cls:			Optional[Any]   ) -> List[str]:




    '''simple docstring'''




    super().setUpClass()
    torch.use_deterministic_algorithms(__lowerCAmelCase   )






  @classmethod
  def 	_UpperCAmelCase				(				cls:			Optional[int]   ) -> Optional[int]:




    '''simple docstring'''




    super().tearDownClass()
    torch.use_deterministic_algorithms(__lowerCAmelCase   )






  def 	_UpperCAmelCase				(				self:			int   ) -> List[str]:




    '''simple docstring'''




    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.00085				,				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-05				,				num_attention_heads=4				,				num_hidden_layers=5				,				pad_token_id=1				,				vocab_size=1_000				,				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 	_UpperCAmelCase				(				self:			str				,				__lowerCAmelCase:			Optional[Any]				,				__lowerCAmelCase:			List[str]=0   ) -> List[str]:




    '''simple docstring'''




    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 	_UpperCAmelCase				(				self:			Tuple   ) -> int:




    '''simple docstring'''




    __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.63905364, 0.62897307, 0.48599017, 0.5133624, 0.5550048, 0.45769516, 0.50326973, 0.5023139, 0.45384496]   )
    __UpperCAmelCase						    =       np.abs(image_slice.flatten() - expected_slice   ).max()
    self.assertLessEqual(__lowerCAmelCase				,				1E-3   )






  def 	_UpperCAmelCase				(				self:			Optional[Any]   ) -> List[Any]:




    '''simple docstring'''




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






  def 	_UpperCAmelCase				(				self:			int   ) -> str:




    '''simple docstring'''




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






  def 	_UpperCAmelCase				(				self:			Tuple   ) -> List[str]:




    '''simple docstring'''




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






  def 	_UpperCAmelCase				(				self:			List[Any]   ) -> Union[str, Any]:




    '''simple docstring'''




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






  def 	_UpperCAmelCase				(				self:			Tuple   ) -> str:




    '''simple docstring'''




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






  def 	_UpperCAmelCase				(				self:			Union[str, Any]   ) -> Union[str, Any]:




    '''simple docstring'''




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






  def 	_UpperCAmelCase				(				self:			Optional[Any]   ) -> str:




    '''simple docstring'''




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






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



  """simple docstring"""






  @classmethod
  def 	_UpperCAmelCase				(				cls:			Optional[Any]   ) -> Union[str, Any]:




    '''simple docstring'''




    super().setUpClass()
    torch.use_deterministic_algorithms(__lowerCAmelCase   )






  @classmethod
  def 	_UpperCAmelCase				(				cls:			Optional[Any]   ) -> Optional[Any]:




    '''simple docstring'''




    super().tearDownClass()
    torch.use_deterministic_algorithms(__lowerCAmelCase   )






  def 	_UpperCAmelCase				(				self:			Optional[Any]   ) -> List[Any]:




    '''simple docstring'''




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






  def 	_UpperCAmelCase				(				self:			Union[str, Any]   ) -> Dict:




    '''simple docstring'''




    __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




from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


a_							=      {
    """configuration_squeezebert""": [
        """SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """SqueezeBertConfig""",
        """SqueezeBertOnnxConfig""",
    ],
    """tokenization_squeezebert""": ["""SqueezeBertTokenizer"""],
}

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

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       a_							=      [
           """SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """SqueezeBertForMaskedLM""",
           """SqueezeBertForMultipleChoice""",
           """SqueezeBertForQuestionAnswering""",
           """SqueezeBertForSequenceClassification""",
           """SqueezeBertForTokenClassification""",
           """SqueezeBertModel""",
           """SqueezeBertModule""",
           """SqueezeBertPreTrainedModel""",
       ]


if TYPE_CHECKING:
       from .configuration_squeezebert import (
           SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
           SqueezeBertConfig,
           SqueezeBertOnnxConfig,
       )
       from .tokenization_squeezebert import SqueezeBertTokenizer

       try:
              if not is_tokenizers_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tokenization_squeezebert_fast import SqueezeBertTokenizerFast

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_squeezebert import (
                  SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                  SqueezeBertForMaskedLM,
                  SqueezeBertForMultipleChoice,
                  SqueezeBertForQuestionAnswering,
                  SqueezeBertForSequenceClassification,
                  SqueezeBertForTokenClassification,
                  SqueezeBertModel,
                  SqueezeBertModule,
                  SqueezeBertPreTrainedModel,
              )

else:
       import sys

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



