


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__   (	__magic_name__							,				__magic_name__							,				__magic_name__							,				unittest.TestCase   ):
					__SCREAMING_SNAKE_CASE       :  int						=   StableDiffusionAttendAndExcitePipeline
					__SCREAMING_SNAKE_CASE       :  Optional[Any]						=   False
					__SCREAMING_SNAKE_CASE       :  List[str]						=   TEXT_TO_IMAGE_PARAMS
					__SCREAMING_SNAKE_CASE       :  Tuple						=   TEXT_TO_IMAGE_BATCH_PARAMS.union({'token_indices'}   )
					__SCREAMING_SNAKE_CASE       :  Tuple						=   TEXT_TO_IMAGE_IMAGE_PARAMS
					__SCREAMING_SNAKE_CASE       :  Any						=   TEXT_TO_IMAGE_IMAGE_PARAMS

					@classmethod
					def   UpperCAmelCase__       (     cls       :   List[Any]       ):






										'''simple docstring'''






										super().setUpClass()
										torch.use_deterministic_algorithms(UpperCamelCase__       )

					@classmethod
					def   UpperCAmelCase__       (     cls       :   Optional[Any]       ):






										'''simple docstring'''






										super().tearDownClass()
										torch.use_deterministic_algorithms(UpperCamelCase__       )

					def   UpperCAmelCase__       (     self       :   Optional[int]       ):






										'''simple docstring'''






										torch.manual_seed(0       )
										lowercase_  							=			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=UpperCamelCase__							,					)
										lowercase_  							=			DDIMScheduler(
										    beta_start=0.00_085							,					beta_end=0.012							,					beta_schedule="""scaled_linear"""							,					clip_sample=UpperCamelCase__							,					set_alpha_to_one=UpperCamelCase__							,					)
										torch.manual_seed(0       )
										lowercase_  							=			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       )
										lowercase_  							=			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							,					)
										lowercase_  							=			CLIPTextModel(UpperCamelCase__       )
										lowercase_  							=			CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       )

										lowercase_  							=			{
										    """unet""": unet,
										    """scheduler""": scheduler,
										    """vae""": vae,
										    """text_encoder""": text_encoder,
										    """tokenizer""": tokenizer,
										    """safety_checker""": None,
										    """feature_extractor""": None,
										}

										return components

					def   UpperCAmelCase__       (     self       :   List[Any]							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   List[Any]=0       ):






										'''simple docstring'''






										if str(UpperCamelCase__       ).startswith("""mps"""       ):
															lowercase_  							=			torch.manual_seed(UpperCamelCase__       )
										else:
															lowercase_  							=			torch.Generator(device=UpperCamelCase__       ).manual_seed(UpperCamelCase__       )
										lowercase_  							=			lowercase_  							=			{
										    """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       :   List[Any]       ):






										'''simple docstring'''






										lowercase_  							=			"""cpu"""

										lowercase_  							=			self.get_dummy_components()
										lowercase_  							=			self.pipeline_class(**UpperCamelCase__       )
										pipe.to(UpperCamelCase__       )
										pipe.set_progress_bar_config(disable=UpperCamelCase__       )

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

										self.assertEqual(image.shape							,					(1, 64, 64, 3)       )
										lowercase_  							=			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]       )
										lowercase_  							=			np.abs(image_slice.flatten() - expected_slice       ).max()
										self.assertLessEqual(UpperCamelCase__							,					1e-3       )

					def   UpperCAmelCase__       (     self       :   List[Any]       ):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   int       ):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   Dict       ):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   Optional[int]       ):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   int       ):






										'''simple docstring'''






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







					def   UpperCAmelCase__       (     self       :   Optional[int]       ):






										'''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       :   Optional[Any]       ):






										'''simple docstring'''






										super().setUpClass()
										torch.use_deterministic_algorithms(UpperCamelCase__       )

					@classmethod
					def   UpperCAmelCase__       (     cls       :   Dict       ):






										'''simple docstring'''






										super().tearDownClass()
										torch.use_deterministic_algorithms(UpperCamelCase__       )

					def   UpperCAmelCase__       (     self       :   Tuple       ):






										'''simple docstring'''






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







					def   UpperCAmelCase__       (     self       :   List[Any]       ):






										'''simple docstring'''






										lowercase_  							=			torch.manual_seed(51       )

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

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

										lowercase_  							=			pipe(
										    prompt=UpperCamelCase__							,					token_indices=UpperCamelCase__							,					guidance_scale=7.5							,					generator=UpperCamelCase__							,					num_inference_steps=5							,					max_iter_to_alter=5							,					output_type="""numpy"""							,					).images[0]

										lowercase_  							=			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 ...utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_torch_available,
    is_transformers_available,
)


try:
							if not (is_transformers_available() and is_torch_available()):
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
							from .multicontrolnet import MultiControlNetModel
							from .pipeline_controlnet import StableDiffusionControlNetPipeline
							from .pipeline_controlnet_imgaimg import StableDiffusionControlNetImgaImgPipeline
							from .pipeline_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline


if is_transformers_available() and is_flax_available():
							from .pipeline_flax_controlnet import FlaxStableDiffusionControlNetPipeline


