



import os
from pickle import UnpicklingError
from typing import Dict, Tuple

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict, unflatten_dict

import transformers

from .utils import logging


__UpperCAmelCase          =       logging.get_logger(__name__)
def 		UpperCamelCase	(   snake_case__       :    Union[str, Any]       ,     snake_case__       :    Optional[int]       ,     snake_case__       :    int       ,     snake_case__       :    Dict=False	)	->     Optional[Any]:
   try:
      import torch  # noqa: F401
   except ImportError:
      logger.error(
          'Loading a PyTorch model in Flax, requires both PyTorch and Flax to be installed. Please see'
          ' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'
          ' instructions.'	)
      raise

   if not is_sharded:
      UpperCamelCase    :  Dict							=	os.path.abspath(__lowerCamelCase	)
      logger.info(F"""Loading PyTorch weights from {pt_path}"""	)

      UpperCamelCase    :  List[str]							=	torch.load(__lowerCamelCase       ,     map_location='cpu'	)
      logger.info(F"""PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()	):,} parameters."""	)

      UpperCamelCase    :  List[Any]							=	convert_pytorch_state_dict_to_flax(__lowerCamelCase       ,     __lowerCamelCase	)
   else:
      # model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files
      UpperCamelCase    :  Dict							=	convert_pytorch_sharded_state_dict_to_flax(__lowerCamelCase       ,     __lowerCamelCase	)
   return flax_state_dict
def 		UpperCamelCase	(   snake_case__       :    List[str]       ,     snake_case__       :    List[Any]       ,     snake_case__       :    Optional[Any]       ,     snake_case__       :    Any       ,     )	->     (Tuple[str], np.ndarray):

   def is_key_or_prefix_key_in_dict(snake_case__       :    int	) -> bool:
      return len(set(__lowerCamelCase	) & {key, (model_prefix,) + key}	) > 0

   # layer norm
   UpperCamelCase    :  Union[str, Any]							=	pt_tuple_key[:-1] + ('scale',)
   if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(__lowerCamelCase	):
      return renamed_pt_tuple_key, pt_tensor

   # batch norm layer mean
   UpperCamelCase    :  List[Any]							=	pt_tuple_key[:-1] + ('mean',)
   if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(__lowerCamelCase	):
      return renamed_pt_tuple_key, pt_tensor

   # batch norm layer var
   UpperCamelCase    :  Optional[Any]							=	pt_tuple_key[:-1] + ('var',)
   if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(__lowerCamelCase	):
      return renamed_pt_tuple_key, pt_tensor

   # embedding
   UpperCamelCase    :  Optional[int]							=	pt_tuple_key[:-1] + ('embedding',)
   if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(__lowerCamelCase	):
      return renamed_pt_tuple_key, pt_tensor

   # conv layer
   UpperCamelCase    :  Optional[int]							=	pt_tuple_key[:-1] + ('kernel',)
   if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and not is_key_or_prefix_key_in_dict(__lowerCamelCase	):
      UpperCamelCase    :  Any							=	pt_tensor.transpose(2       ,     3       ,     1       ,     0	)
      return renamed_pt_tuple_key, pt_tensor

   # linear layer
   UpperCamelCase    :  Dict							=	pt_tuple_key[:-1] + ('kernel',)
   if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(__lowerCamelCase	):
      UpperCamelCase    :  Union[str, Any]							=	pt_tensor.T
      return renamed_pt_tuple_key, pt_tensor

   # old PyTorch layer norm weight
   UpperCamelCase    :  List[Any]							=	pt_tuple_key[:-1] + ('weight',)
   if pt_tuple_key[-1] == "gamma":
      return renamed_pt_tuple_key, pt_tensor

   # old PyTorch layer norm bias
   UpperCamelCase    :  List[str]							=	pt_tuple_key[:-1] + ('bias',)
   if pt_tuple_key[-1] == "beta":
      return renamed_pt_tuple_key, pt_tensor

   # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
   UpperCamelCase    :  str							=	None
   if pt_tuple_key[-3::2] == ("parametrizations", "original0"):
      UpperCamelCase    :  List[Any]							=	pt_tuple_key[-2] + '_g'
   elif pt_tuple_key[-3::2] == ("parametrizations", "original1"):
      UpperCamelCase    :  Dict							=	pt_tuple_key[-2] + '_v'
   if name is not None:
      UpperCamelCase    :  Optional[int]							=	pt_tuple_key[:-3] + (name,)
      return renamed_pt_tuple_key, pt_tensor

   return pt_tuple_key, pt_tensor
def 		UpperCamelCase	(   snake_case__       :    Dict       ,     snake_case__       :    Optional[int]	)	->     List[str]:
   # convert pytorch tensor to numpy
   UpperCamelCase    :  List[str]							=	{k: v.numpy() for k, v in pt_state_dict.items()}

   UpperCamelCase    :  str							=	flax_model.base_model_prefix

   # use params dict if the model contains batch norm layers
   if "params" in flax_model.params:
      UpperCamelCase    :  Dict							=	flax_model.params['params']
   else:
      UpperCamelCase    :  Any							=	flax_model.params
   UpperCamelCase    :  List[Any]							=	flatten_dict(__lowerCamelCase	)

   # add batch_stats keys,values to dict
   if "batch_stats" in flax_model.params:
      UpperCamelCase    :  int							=	flatten_dict(flax_model.params['batch_stats']	)
      random_flax_state_dict.update(__lowerCamelCase	)

   UpperCamelCase    :  List[str]							=	{}

   UpperCamelCase    :  Dict							=	(model_prefix not in flax_model_params) and (
       model_prefix in {k.split('.'	)[0] for k in pt_state_dict.keys()}
   )
   UpperCamelCase    :  Optional[Any]							=	(model_prefix in flax_model_params) and (
       model_prefix not in {k.split('.'	)[0] for k in pt_state_dict.keys()}
   )

   # Need to change some parameters name to match Flax names
   for pt_key, pt_tensor in pt_state_dict.items():
      UpperCamelCase    :  str							=	tuple(pt_key.split('.'	)	)

      # remove base model prefix if necessary
      UpperCamelCase    :  int							=	pt_tuple_key[0] == model_prefix
      if load_model_with_head_into_base_model and has_base_model_prefix:
         UpperCamelCase    :  str							=	pt_tuple_key[1:]

      # Correctly rename weight parameters
      UpperCamelCase	,     UpperCamelCase    :  Dict							=	rename_key_and_reshape_tensor(
          __lowerCamelCase       ,     __lowerCamelCase       ,     __lowerCamelCase       ,     __lowerCamelCase	)

      # add model prefix if necessary
      UpperCamelCase    :  List[str]							=	(model_prefix,) + flax_key in random_flax_state_dict
      if load_base_model_into_model_with_head and require_base_model_prefix:
         UpperCamelCase    :  int							=	(model_prefix,) + flax_key

      if flax_key in random_flax_state_dict:
         if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
            raise ValueError(
                F"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
                F"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""	)

        # add batch stats if the model contains batchnorm layers
      if "batch_stats" in flax_model.params:
         if "mean" in flax_key[-1] or "var" in flax_key[-1]:
            UpperCamelCase    :  Dict							=	jnp.asarray(__lowerCamelCase	)
            continue
         # remove num_batches_tracked key
         if "num_batches_tracked" in flax_key[-1]:
            flax_state_dict.pop(__lowerCamelCase       ,     __lowerCamelCase	)
            continue

         # also add unexpected weight so that warning is thrown
         UpperCamelCase    :  Tuple							=	jnp.asarray(__lowerCamelCase	)

      else:
         # also add unexpected weight so that warning is thrown
         UpperCamelCase    :  Union[str, Any]							=	jnp.asarray(__lowerCamelCase	)

   return unflatten_dict(__lowerCamelCase	)
def 		UpperCamelCase	(   snake_case__       :    Any       ,     snake_case__       :    int	)	->     List[str]:
   import torch

   # Load the index
   UpperCamelCase    :  Any							=	{}
   for shard_file in shard_filenames:
      # load using msgpack utils
      UpperCamelCase    :  Dict							=	torch.load(__lowerCamelCase	)
      UpperCamelCase    :  Any							=	{k: v.numpy() for k, v in pt_state_dict.items()}

      UpperCamelCase    :  Any							=	flax_model.base_model_prefix

      # use params dict if the model contains batch norm layers and then add batch_stats keys,values to dict
      if "batch_stats" in flax_model.params:
         UpperCamelCase    :  str							=	flax_model.params['params']

         UpperCamelCase    :  List[str]							=	flatten_dict(__lowerCamelCase	)
         random_flax_state_dict.update(flatten_dict(flax_model.params['batch_stats']	)	)
      else:
         UpperCamelCase    :  Any							=	flax_model.params
         UpperCamelCase    :  Union[str, Any]							=	flatten_dict(__lowerCamelCase	)

      UpperCamelCase    :  List[str]							=	(model_prefix not in flax_model_params) and (
          model_prefix in {k.split('.'	)[0] for k in pt_state_dict.keys()}
      )
      UpperCamelCase    :  Optional[int]							=	(model_prefix in flax_model_params) and (
          model_prefix not in {k.split('.'	)[0] for k in pt_state_dict.keys()}
      )
      # Need to change some parameters name to match Flax names
      for pt_key, pt_tensor in pt_state_dict.items():
         UpperCamelCase    :  List[str]							=	tuple(pt_key.split('.'	)	)

         # remove base model prefix if necessary
         UpperCamelCase    :  List[Any]							=	pt_tuple_key[0] == model_prefix
         if load_model_with_head_into_base_model and has_base_model_prefix:
            UpperCamelCase    :  Optional[int]							=	pt_tuple_key[1:]

         # Correctly rename weight parameters
         UpperCamelCase	,     UpperCamelCase    :  Tuple							=	rename_key_and_reshape_tensor(
             __lowerCamelCase       ,     __lowerCamelCase       ,     __lowerCamelCase       ,     __lowerCamelCase	)
         # add model prefix if necessary
         UpperCamelCase    :  Optional[Any]							=	(model_prefix,) + flax_key in random_flax_state_dict
         if load_base_model_into_model_with_head and require_base_model_prefix:
            UpperCamelCase    :  Dict							=	(model_prefix,) + flax_key

         if flax_key in random_flax_state_dict:
            if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
               raise ValueError(
                   F"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
                   F"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""	)

            # add batch stats if the model contains batchnorm layers
         if "batch_stats" in flax_model.params:
            if "mean" in flax_key[-1]:
               UpperCamelCase    :  List[str]							=	jnp.asarray(__lowerCamelCase	)
               continue
            if "var" in flax_key[-1]:
               UpperCamelCase    :  Dict							=	jnp.asarray(__lowerCamelCase	)
               continue
            # remove num_batches_tracked key
            if "num_batches_tracked" in flax_key[-1]:
               flax_state_dict.pop(__lowerCamelCase       ,     __lowerCamelCase	)
               continue

            # also add unexpected weight so that warning is thrown
            UpperCamelCase    :  int							=	jnp.asarray(__lowerCamelCase	)

         else:
            # also add unexpected weight so that warning is thrown
            UpperCamelCase    :  List[str]							=	jnp.asarray(__lowerCamelCase	)
   return unflatten_dict(__lowerCamelCase	)
def 		UpperCamelCase	(   snake_case__       :    Optional[int]       ,     snake_case__       :    str	)	->     Any:
   UpperCamelCase    :  str							=	os.path.abspath(__lowerCamelCase	)
   logger.info(F"""Loading Flax weights from {flax_checkpoint_path}"""	)

   # import correct flax class
   UpperCamelCase    :  Optional[Any]							=	getattr(__lowerCamelCase       ,     'Flax' + model.__class__.__name__	)

   # load flax weight dict
   with open(__lowerCamelCase       ,     'rb'	) as state_f:
      try:
         UpperCamelCase    :  List[str]							=	from_bytes(__lowerCamelCase       ,     state_f.read()	)
      except UnpicklingError:
         raise EnvironmentError(F"""Unable to convert {flax_checkpoint_path} to Flax deserializable object. """	)

   return load_flax_weights_in_pytorch_model(__lowerCamelCase       ,     __lowerCamelCase	)


def 		UpperCamelCase	(   snake_case__       :    Optional[Any]       ,     snake_case__       :    Tuple	)	->     Optional[int]:

   try:
      import torch  # noqa: F401
   except ImportError:
      logger.error(
          'Loading a Flax weights in PyTorch, requires both PyTorch and Flax to be installed. Please see'
          ' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'
          ' instructions.'	)
      raise

   # check if we have bf16 weights
   UpperCamelCase    :  Union[str, Any]							=	flatten_dict(jax.tree_util.tree_map(lambda snake_case__	: x.dtype == jnp.bfloataa       ,     __lowerCamelCase	)	).values()
   if any(__lowerCamelCase	):
      # convert all weights to fp32 if the are bf16 since torch.from_numpy can-not handle bf16
      # and bf16 is not fully supported in PT yet.
      logger.warning(
          'Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` '
          'before loading those in PyTorch model.'	)
      UpperCamelCase    :  str							=	jax.tree_util.tree_map(
          lambda snake_case__	: params.astype(np.floataa	) if params.dtype == jnp.bfloataa else params       ,     __lowerCamelCase	)

   UpperCamelCase    :  Optional[int]							=	flatten_dict(__lowerCamelCase	)
   UpperCamelCase    :  List[Any]							=	pt_model.state_dict()

   UpperCamelCase    :  Any							=	(pt_model.base_model_prefix in flax_state) and (
       pt_model.base_model_prefix not in {k.split('.'	)[0] for k in pt_model_dict.keys()}
   )
   UpperCamelCase    :  Optional[int]							=	(pt_model.base_model_prefix not in flax_state) and (
       pt_model.base_model_prefix in {k.split('.'	)[0] for k in pt_model_dict.keys()}
   )

   # keep track of unexpected & missing keys
   UpperCamelCase    :  Optional[int]							=	[]
   UpperCamelCase    :  str							=	set(pt_model_dict.keys()	)

   for flax_key_tuple, flax_tensor in flax_state_dict.items():
      UpperCamelCase    :  Any							=	flax_key_tuple[0] == pt_model.base_model_prefix
      UpperCamelCase    :  List[str]							=	'.'.join((pt_model.base_model_prefix,) + flax_key_tuple	) in pt_model_dict

      # adapt flax_key to prepare for loading from/to base model only
      if load_model_with_head_into_base_model and has_base_model_prefix:
         UpperCamelCase    :  int							=	flax_key_tuple[1:]
      elif load_base_model_into_model_with_head and require_base_model_prefix:
         UpperCamelCase    :  List[Any]							=	(pt_model.base_model_prefix,) + flax_key_tuple

      # rename flax weights to PyTorch format
      if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(__lowerCamelCase	) not in pt_model_dict:
         # conv layer
         UpperCamelCase    :  str							=	flax_key_tuple[:-1] + ('weight',)
         UpperCamelCase    :  List[Any]							=	jnp.transpose(__lowerCamelCase       ,     (3, 2, 0, 1)	)
      elif flax_key_tuple[-1] == "kernel" and ".".join(__lowerCamelCase	) not in pt_model_dict:
         # linear layer
         UpperCamelCase    :  List[Any]							=	flax_key_tuple[:-1] + ('weight',)
         UpperCamelCase    :  int							=	flax_tensor.T
      elif flax_key_tuple[-1] in ["scale", "embedding"]:
         UpperCamelCase    :  Optional[int]							=	flax_key_tuple[:-1] + ('weight',)

      # adding batch stats from flax batch norm to pt
      elif "mean" in flax_key_tuple[-1]:
         UpperCamelCase    :  str							=	flax_key_tuple[:-1] + ('running_mean',)
      elif "var" in flax_key_tuple[-1]:
         UpperCamelCase    :  Optional[int]							=	flax_key_tuple[:-1] + ('running_var',)

      if "batch_stats" in flax_state:
         UpperCamelCase    :  Any							=	'.'.join(flax_key_tuple[1:]	)  # Remove the params/batch_stats header
      else:
         UpperCamelCase    :  Union[str, Any]							=	'.'.join(__lowerCamelCase	)

      # We also need to look at `pt_model_dict` and see if there are keys requiring further transformation.
      UpperCamelCase    :  Tuple							=	{}
      # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
      for key in pt_model_dict:
         UpperCamelCase    :  Tuple							=	key.split('.'	)
         UpperCamelCase    :  List[str]							=	None
         if key_components[-3::2] == ["parametrizations", "original0"]:
            UpperCamelCase    :  Optional[Any]							=	key_components[-2] + '_g'
         elif key_components[-3::2] == ["parametrizations", "original1"]:
            UpperCamelCase    :  Any							=	key_components[-2] + '_v'
         if name is not None:
            UpperCamelCase    :  Tuple							=	key_components[:-3] + [name]
            UpperCamelCase    :  Union[str, Any]							=	'.'.join(__lowerCamelCase	)
            UpperCamelCase    :  List[str]							=	key

      if flax_key in special_pt_names:
         UpperCamelCase    :  Optional[Any]							=	special_pt_names[flax_key]

      if flax_key in pt_model_dict:
         if flax_tensor.shape != pt_model_dict[flax_key].shape:
            raise ValueError(
                F"""Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected """
                F"""to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}."""	)
         else:
            # add weight to pytorch dict
            UpperCamelCase    :  Dict							=	np.asarray(__lowerCamelCase	) if not isinstance(__lowerCamelCase       ,     np.ndarray	) else flax_tensor
            UpperCamelCase    :  Dict							=	torch.from_numpy(__lowerCamelCase	)
            # remove from missing keys
            missing_keys.remove(__lowerCamelCase	)
      else:
         # weight is not expected by PyTorch model
         unexpected_keys.append(__lowerCamelCase	)

   pt_model.load_state_dict(__lowerCamelCase	)

   # re-transform missing_keys to list
   UpperCamelCase    :  str							=	list(__lowerCamelCase	)

   if len(__lowerCamelCase	) > 0:
      logger.warning(
          'Some weights of the Flax model were not used when initializing the PyTorch model'
          F""" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing"""
          F""" {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture"""
          ' (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This'
          F""" IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect"""
          ' to be exactly identical (e.g. initializing a BertForSequenceClassification model from a'
          ' FlaxBertForSequenceClassification model).'	)
   else:
      logger.warning(F"""All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n"""	)
   if len(__lowerCamelCase	) > 0:
      logger.warning(
          F"""Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly"""
          F""" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to"""
          ' use it for predictions and inference.'	)
   else:
      logger.warning(
          F"""All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n"""
          'If your task is similar to the task the model of the checkpoint was trained on, '
          F"""you can already use {pt_model.__class__.__name__} for predictions without further training."""	)

   return pt_model









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


__UpperCamelCase					:       Optional[Any]    					=      False

@skip_mps
class 		__lowerCAmelCase	(	__magic_name__      ,       __magic_name__      ,       __magic_name__      ,       unittest.TestCase ):
							UpperCamelCase__					     =  StableDiffusionAttendAndExcitePipeline
							UpperCamelCase__					     =  False
							UpperCamelCase__					     =  TEXT_TO_IMAGE_PARAMS
							UpperCamelCase__					     =  TEXT_TO_IMAGE_BATCH_PARAMS.union({'''token_indices'''} )
							UpperCamelCase__					     =  TEXT_TO_IMAGE_IMAGE_PARAMS
							UpperCamelCase__					     =  TEXT_TO_IMAGE_IMAGE_PARAMS
							@classmethod
							def       lowerCamelCase__							(						cls						:Union[str, Any]				):







								'''simple docstring'''







								super().setUpClass()
								torch.use_deterministic_algorithms(__magic_name__				)
							@classmethod
							def       lowerCamelCase__							(						cls						:Union[str, Any]				):







								'''simple docstring'''







								super().tearDownClass()
								torch.use_deterministic_algorithms(__magic_name__				)
							def       lowerCamelCase__							(						self						:Tuple				):







								'''simple docstring'''







								torch.manual_seed(0				)
								a										=							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__		,      )
								a										=							DDIMScheduler(
								    beta_start=0.00085		,      beta_end=0.012		,      beta_schedule="""scaled_linear"""		,      clip_sample=__magic_name__		,      set_alpha_to_one=__magic_name__		,      )
								torch.manual_seed(0				)
								a										=							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				)
								a										=							CLIPTextConfig(
								    bos_token_id=0		,      eos_token_id=2		,      hidden_size=32		,      intermediate_size=37		,      layer_norm_eps=1E-0_5		,      num_attention_heads=4		,      num_hidden_layers=5		,      pad_token_id=1		,      vocab_size=1000		,      hidden_act="""gelu"""		,      projection_dim=512		,      )
								a										=							CLIPTextModel(__magic_name__				)
								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       lowerCamelCase__							(						self						:Dict		,      __magic_name__						:Optional[int]		,      __magic_name__						:Optional[int]=0				):







								'''simple docstring'''







								if str(__magic_name__				).startswith("""mps"""				):
									a										=							torch.manual_seed(__magic_name__				)
								else:
									a										=							torch.Generator(device=__magic_name__				).manual_seed(__magic_name__				)
								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       lowerCamelCase__							(						self						:str				):







								'''simple docstring'''







								a										=							"""cpu"""

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

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

								self.assertEqual(image.shape		,      (1, 64, 64, 3)				)
								a										=							np.array(
								    [0.63905364, 0.62897307, 0.48599017, 0.5133624, 0.5550048, 0.45769516, 0.50326973, 0.5023139, 0.45384496]				)
								a										=							np.abs(image_slice.flatten() - expected_slice				).max()
								self.assertLessEqual(__magic_name__		,      1E-3				)
							def       lowerCamelCase__							(						self						:Tuple				):







								'''simple docstring'''







								super().test_cpu_offload_forward_pass(expected_max_diff=5E-4				)
							def       lowerCamelCase__							(						self						:List[Any]				):







								'''simple docstring'''







								self._test_inference_batch_consistent(batch_sizes=[1, 2]				)
							def       lowerCamelCase__							(						self						:List[str]				):







								'''simple docstring'''







								self._test_inference_batch_single_identical(batch_size=2		,      expected_max_diff=7E-4				)
							def       lowerCamelCase__							(						self						:List[str]				):







								'''simple docstring'''







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







								'''simple docstring'''







								super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5E-4				)
							def       lowerCamelCase__							(						self						:List[str]				):







								'''simple docstring'''







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





							def       lowerCamelCase__							(						self						:str				):







								'''simple docstring'''







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







@require_torch_gpu
@slow
class 		__lowerCAmelCase	(	unittest.TestCase ):
							@classmethod
							def       lowerCamelCase__							(						cls						:int				):







								'''simple docstring'''







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







								'''simple docstring'''







								super().tearDownClass()
								torch.use_deterministic_algorithms(__magic_name__				)
							def       lowerCamelCase__							(						self						:Union[str, Any]				):







								'''simple docstring'''







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





							def       lowerCamelCase__							(						self						:int				):







								'''simple docstring'''







								a										=							torch.manual_seed(51				)

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

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

								a										=							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]

								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

