
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 						lowercase_				(				lowerCamelCase_							,		lowerCamelCase_							,		lowerCamelCase_							,		unittest.TestCase ):
 A__  : Union[str, Any]      					=			StableDiffusionAttendAndExcitePipeline
 A__  : List[str]      					=			False
 A__  : int      					=			TEXT_TO_IMAGE_PARAMS
 A__  : List[Any]      					=			TEXT_TO_IMAGE_BATCH_PARAMS.union({"""token_indices"""} )
 A__  : List[str]      					=			TEXT_TO_IMAGE_IMAGE_PARAMS
 A__  : str      					=			TEXT_TO_IMAGE_IMAGE_PARAMS

 @classmethod
 def lowerCamelCase_ (					cls   ):




        """simple docstring"""


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

 @classmethod
 def lowerCamelCase_ (					cls   ):




        """simple docstring"""


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

 def lowerCamelCase_ (					self   ):




        """simple docstring"""


        torch.manual_seed(0   )
        UpperCamelCase_											=				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   ,		)
        UpperCamelCase_											=				DDIMScheduler(
            beta_start=0.00_085   ,		beta_end=0.012   ,		beta_schedule="""scaled_linear"""   ,		clip_sample=__snake_case   ,		set_alpha_to_one=__snake_case   ,		)
        torch.manual_seed(0   )
        UpperCamelCase_											=				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   )
        UpperCamelCase_											=				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   ,		)
        UpperCamelCase_											=				CLIPTextModel(__snake_case   )
        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 lowerCamelCase_ (					self   ,		__UpperCamelCase   ,		__UpperCamelCase=0   ):




        """simple docstring"""


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




        """simple docstring"""


        UpperCamelCase_											=				"""cpu"""

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

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

        self.assertEqual(image.shape   ,		(1, 6_4, 6_4, 3)   )
        UpperCamelCase_											=				np.array(
            [0.63_905_364, 0.62_897_307, 0.48_599_017, 0.5_133_624, 0.5_550_048, 0.45_769_516, 0.50_326_973, 0.5_023_139, 0.45_384_496]   )
        UpperCamelCase_											=				np.abs(image_slice.flatten() - expected_slice   ).max()
        self.assertLessEqual(__snake_case   ,		1e-3   )

 def lowerCamelCase_ (					self   ):




        """simple docstring"""


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

 def lowerCamelCase_ (					self   ):




        """simple docstring"""


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

 def lowerCamelCase_ (					self   ):




        """simple docstring"""


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

 def lowerCamelCase_ (					self   ):




        """simple docstring"""


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

 def lowerCamelCase_ (					self   ):




        """simple docstring"""


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

 def lowerCamelCase_ (					self   ):




        """simple docstring"""


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

 def lowerCamelCase_ (					self   ):




        """simple docstring"""


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



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

 @classmethod
 def lowerCamelCase_ (					cls   ):




        """simple docstring"""


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

 @classmethod
 def lowerCamelCase_ (					cls   ):




        """simple docstring"""


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

 def lowerCamelCase_ (					self   ):




        """simple docstring"""


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

 def lowerCamelCase_ (					self   ):




        """simple docstring"""


        UpperCamelCase_											=				torch.manual_seed(5_1   )

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

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

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

        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_tf_available, is_torch_available


snake_case							:       str     			=      {
    '''configuration_rag''': ['''RagConfig'''],
    '''retrieval_rag''': ['''RagRetriever'''],
    '''tokenization_rag''': ['''RagTokenizer'''],
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						snake_case							:       Any     			=      [
						    '''RagModel''',
						    '''RagPreTrainedModel''',
						    '''RagSequenceForGeneration''',
						    '''RagTokenForGeneration''',
						]

try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						snake_case							:       Optional[int]     			=      [
						    '''TFRagModel''',
						    '''TFRagPreTrainedModel''',
						    '''TFRagSequenceForGeneration''',
						    '''TFRagTokenForGeneration''',
						]


if TYPE_CHECKING:
						from .configuration_rag import RagConfig
						from .retrieval_rag import RagRetriever
						from .tokenization_rag import RagTokenizer

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_rag import RagModel, RagPreTrainedModel, RagSequenceForGeneration, RagTokenForGeneration

						try:
												if not is_tf_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_tf_rag import (
												    TFRagModel,
												    TFRagPreTrainedModel,
												    TFRagSequenceForGeneration,
												    TFRagTokenForGeneration,
												)

else:
						import sys

						snake_case							:       int     			=      _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)




