



"""simple docstring"""







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


lowercase__		:				List[str]          =	False







@skip_mps
class       UpperCamelCase__   (	_UpperCAmelCase,					_UpperCAmelCase,					_UpperCAmelCase,					unittest.TestCase  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     StableDiffusionAttendAndExcitePipeline
      _SCREAMING_SNAKE_CASE  				=     False
      _SCREAMING_SNAKE_CASE  				=     TEXT_TO_IMAGE_PARAMS
      _SCREAMING_SNAKE_CASE  				=     TEXT_TO_IMAGE_BATCH_PARAMS.union({"""token_indices"""}  )
      _SCREAMING_SNAKE_CASE  				=     TEXT_TO_IMAGE_IMAGE_PARAMS
      _SCREAMING_SNAKE_CASE  				=     TEXT_TO_IMAGE_IMAGE_PARAMS



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Tuple       ):
        super().setUpClass()
        torch.use_deterministic_algorithms(SCREAMING_SNAKE_CASE_       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Tuple       ):
        super().tearDownClass()
        torch.use_deterministic_algorithms(SCREAMING_SNAKE_CASE_       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      Dict       ):
        torch.manual_seed(0       )
        lowerCAmelCase_   :				Dict								= 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=SCREAMING_SNAKE_CASE_    ,			)
        lowerCAmelCase_   :				Optional[Any]								= DDIMScheduler(
            beta_start=0.0_00_85    ,			beta_end=0.0_12    ,			beta_schedule='scaled_linear'    ,			clip_sample=SCREAMING_SNAKE_CASE_    ,			set_alpha_to_one=SCREAMING_SNAKE_CASE_    ,			)
        torch.manual_seed(0       )
        lowerCAmelCase_   :				Tuple								= 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       )
        lowerCAmelCase_   :				Union[str, Any]								= 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    ,			)
        lowerCAmelCase_   :				Optional[Any]								= CLIPTextModel(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Dict								= 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      SCREAMING_SNAKE_CASE__ (      self				:      Any    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]    ,			SCREAMING_SNAKE_CASE_				:      List[Any]=0       ):
        if str(SCREAMING_SNAKE_CASE_       ).startswith('mps'       ):
          lowerCAmelCase_   :				Any								= torch.manual_seed(SCREAMING_SNAKE_CASE_       )
        else:
          lowerCAmelCase_   :				Tuple								= torch.Generator(device=SCREAMING_SNAKE_CASE_       ).manual_seed(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				str								= {
            '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      SCREAMING_SNAKE_CASE__ (      self				:      int       ):
        lowerCAmelCase_   :				Optional[int]								= 'cpu'

        lowerCAmelCase_   :				str								= self.get_dummy_components()
        lowerCAmelCase_   :				Optional[int]								= self.pipeline_class(**SCREAMING_SNAKE_CASE_       )
        pipe.to(SCREAMING_SNAKE_CASE_       )
        pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_       )

        lowerCAmelCase_   :				Union[str, Any]								= self.get_dummy_inputs(SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Any								= pipe(**SCREAMING_SNAKE_CASE_       ).images
        lowerCAmelCase_   :				int								= image[0, -3:, -3:, -1]

        self.assertEqual(image.shape    ,			(1, 6_4, 6_4, 3)       )
        lowerCAmelCase_   :				int								= np.array(
            [0.63_90_53_64, 0.62_89_73_07, 0.48_59_90_17, 0.5_13_36_24, 0.5_55_00_48, 0.45_76_95_16, 0.50_32_69_73, 0.5_02_31_39, 0.45_38_44_96]       )
        lowerCAmelCase_   :				Optional[int]								= np.abs(image_slice.flatten() - expected_slice       ).max()
        self.assertLessEqual(SCREAMING_SNAKE_CASE_    ,			1E-3       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      Tuple       ):
        super().test_cpu_offload_forward_pass(expected_max_diff=5E-4       )



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



      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[Any]       ):
        self._test_inference_batch_single_identical(batch_size=2    ,			expected_max_diff=7E-4       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      Dict       ):
        super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      Dict       ):
        super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5E-4       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      Tuple       ):
        super().test_save_load_local(expected_max_difference=5E-4       )





      def      SCREAMING_SNAKE_CASE__ (      self				:      int       ):
        super().test_save_load_optional_components(expected_max_difference=4E-4       )







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




      """simple docstring"""




      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      Optional[Any]       ):
        super().setUpClass()
        torch.use_deterministic_algorithms(SCREAMING_SNAKE_CASE_       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      str       ):
        super().tearDownClass()
        torch.use_deterministic_algorithms(SCREAMING_SNAKE_CASE_       )



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





      def      SCREAMING_SNAKE_CASE__ (      self				:      str       ):
        lowerCAmelCase_   :				List[Any]								= torch.manual_seed(5_1       )

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

        lowerCAmelCase_   :				str								= 'a painting of an elephant with glasses'
        lowerCAmelCase_   :				int								= [5, 7]

        lowerCAmelCase_   :				Dict								= pipe(
            prompt=SCREAMING_SNAKE_CASE_    ,			token_indices=SCREAMING_SNAKE_CASE_    ,			guidance_scale=7.5    ,			generator=SCREAMING_SNAKE_CASE_    ,			num_inference_steps=5    ,			max_iter_to_alter=5    ,			output_type='numpy'    ,			).images[0]

        lowerCAmelCase_   :				Dict								= 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












"""simple docstring"""





from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Mapping, Optional

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


if TYPE_CHECKING:
       from ... import FeatureExtractionMixin, TensorType

SCREAMING_SNAKE_CASE__					=   logging.get_logger(__name__)

SCREAMING_SNAKE_CASE__					=   {
    "openai/imagegpt-small": "",
    "openai/imagegpt-medium": "",
    "openai/imagegpt-large": "",
}

class 		lowercase					(					_UpperCAmelCase					):
      _SCREAMING_SNAKE_CASE					=      'imagegpt'
      _SCREAMING_SNAKE_CASE					=      ['past_key_values']
      _SCREAMING_SNAKE_CASE					=      {
          'hidden_size': 'n_embd',
          'max_position_embeddings': 'n_positions',
          'num_attention_heads': 'n_head',
          'num_hidden_layers': 'n_layer',
      }





      def __init__(       self      ,						lowercase=512 + 1      ,						lowercase=32 * 32      ,						lowercase=512      ,						lowercase=24      ,						lowercase=8      ,						lowercase=None      ,						lowercase="quick_gelu"      ,						lowercase=0.1      ,						lowercase=0.1      ,						lowercase=0.1      ,						lowercase=1e-5      ,						lowercase=0.02      ,						lowercase=True      ,						lowercase=True      ,						lowercase=False      ,						lowercase=False      ,						lowercase=False      ,						**lowercase      ,						)		->     Any:
            lowerCAmelCase  						=       vocab_size
            lowerCAmelCase  						=       n_positions
            lowerCAmelCase  						=       n_embd
            lowerCAmelCase  						=       n_layer
            lowerCAmelCase  						=       n_head
            lowerCAmelCase  						=       n_inner
            lowerCAmelCase  						=       activation_function
            lowerCAmelCase  						=       resid_pdrop
            lowerCAmelCase  						=       embd_pdrop
            lowerCAmelCase  						=       attn_pdrop
            lowerCAmelCase  						=       layer_norm_epsilon
            lowerCAmelCase  						=       initializer_range
            lowerCAmelCase  						=       scale_attn_weights
            lowerCAmelCase  						=       use_cache
            lowerCAmelCase  						=       scale_attn_by_inverse_layer_idx
            lowerCAmelCase  						=       reorder_and_upcast_attn
            lowerCAmelCase  						=       tie_word_embeddings

            super().__init__(tie_word_embeddings=lowercase      ,						**lowercase )







class 		lowercase					(					_UpperCAmelCase					):





      @property
      def     _snake_case			(       self )		->     Mapping[str, Mapping[int, str]]:
            return OrderedDict(
                [
                    ("""input_ids""", {0: """batch""", 1: """sequence"""}),
                ] )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase = 1      ,						lowercase = -1      ,						lowercase = False      ,						lowercase = None      ,						lowercase = 3      ,						lowercase = 32      ,						lowercase = 32      ,						)		->     Mapping[str, Any]:
            lowerCAmelCase  						=       self._generate_dummy_images(lowercase      ,						lowercase      ,						lowercase      ,						lowercase )
            lowerCAmelCase  						=       dict(preprocessor(images=lowercase      ,						return_tensors=lowercase ) )

            return inputs






