



import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

import diffusers
from diffusers import (
    AutoencoderKL,
    EulerDiscreteScheduler,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






def  UpperCAmelCase__    (    lowercase__   ) ->      Union[str, Any]:
						__lowercase					=				[tensor.shape for tensor in tensor_list]
						return all(shape == shapes[0] for shape in shapes[1:]   )
class 			_lowerCAmelCase   (							_UpperCAmelCase    ,					_UpperCAmelCase    ,					_UpperCAmelCase    ,					unittest.TestCase			):






			"""simple docstring"""





			lowercase__      :  Union[str, Any]      =			StableDiffusionLatentUpscalePipeline
			lowercase__      :  Optional[int]      =			TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
			    "height",
			    "width",
			    "cross_attention_kwargs",
			    "negative_prompt_embeds",
			    "prompt_embeds",
			}
			lowercase__      :  List[str]      =			PipelineTesterMixin.required_optional_params - {"num_images_per_prompt"}
			lowercase__      :  Optional[Any]      =			TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
			lowercase__      :  Optional[Any]      =			frozenset(
			    []			)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
			lowercase__      :  Any      =			frozenset([]			)

			lowercase__      :  Optional[int]      =			True





			@property
			def 		snake_case__			(				self	:       List[Any]   )						->					Tuple:
									"""simple docstring"""







									__lowercase					=				1
									__lowercase					=				4
									__lowercase					=				(16, 16)

									__lowercase					=				floats_tensor((batch_size, num_channels) + sizes			,     rng=random.Random(0   )   ).to(UpperCamelCase_   )
									return image





			def 		snake_case__			(				self	:       Dict   )						->					str:
									"""simple docstring"""







									torch.manual_seed(0   )
									__lowercase					=				UNetaDConditionModel(
									    act_fn="""gelu"""			,     attention_head_dim=8			,     norm_num_groups=UpperCamelCase_			,     block_out_channels=[32, 32, 64, 64]			,     time_cond_proj_dim=160			,     conv_in_kernel=1			,     conv_out_kernel=1			,     cross_attention_dim=32			,     down_block_types=(
									        """KDownBlock2D""",
									        """KCrossAttnDownBlock2D""",
									        """KCrossAttnDownBlock2D""",
									        """KCrossAttnDownBlock2D""",
									    )			,     in_channels=8			,     mid_block_type=UpperCamelCase_			,     only_cross_attention=UpperCamelCase_			,     out_channels=5			,     resnet_time_scale_shift="""scale_shift"""			,     time_embedding_type="""fourier"""			,     timestep_post_act="""gelu"""			,     up_block_types=("""KCrossAttnUpBlock2D""", """KCrossAttnUpBlock2D""", """KCrossAttnUpBlock2D""", """KUpBlock2D""")			,     )
									__lowercase					=				AutoencoderKL(
									    block_out_channels=[32, 32, 64, 64]			,     in_channels=3			,     out_channels=3			,     down_block_types=[
									        """DownEncoderBlock2D""",
									        """DownEncoderBlock2D""",
									        """DownEncoderBlock2D""",
									        """DownEncoderBlock2D""",
									    ]			,     up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D""", """UpDecoderBlock2D""", """UpDecoderBlock2D"""]			,     latent_channels=4			,     )
									__lowercase					=				EulerDiscreteScheduler(prediction_type="""sample"""   )
									__lowercase					=				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=1_000			,     hidden_act="""quick_gelu"""			,     projection_dim=512			,     )
									__lowercase					=				CLIPTextModel(UpperCamelCase_   )
									__lowercase					=				CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""   )

									__lowercase					=				{
									    "unet": model.eval(),
									    "vae": vae.eval(),
									    "scheduler": scheduler,
									    "text_encoder": text_encoder,
									    "tokenizer": tokenizer,
									}

									return components





			def 		snake_case__			(				self	:       List[Any]			,     lowercase	:       List[Any]			,     lowercase	:       int=0   )						->					Optional[int]:
									"""simple docstring"""







									if str(UpperCamelCase_   ).startswith("""mps"""   ):
															__lowercase					=				torch.manual_seed(UpperCamelCase_   )
									else:
															__lowercase					=				torch.Generator(device=UpperCamelCase_   ).manual_seed(UpperCamelCase_   )
									__lowercase					=				{
									    "prompt": "A painting of a squirrel eating a burger",
									    "image": self.dummy_image.cpu(),
									    "generator": generator,
									    "num_inference_steps": 2,
									    "output_type": "numpy",
									}
									return inputs





			def 		snake_case__			(				self	:       str   )						->					Optional[int]:
									"""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, 256, 256, 3)   )
									__lowercase					=				np.array(
									    [0.4722_2412, 0.4192_1633, 0.4471_7434, 0.4687_4192, 0.4258_8258, 0.4615_0726, 0.467_7534, 0.4558_3832, 0.4857_9055]   )
									__lowercase					=				np.abs(image_slice.flatten() - expected_slice   ).max()
									self.assertLessEqual(UpperCamelCase_			,     1E-3   )





			def 		snake_case__			(				self	:       str   )						->					Union[str, Any]:
									"""simple docstring"""







									super().test_attention_slicing_forward_pass(expected_max_diff=7E-3   )





			def 		snake_case__			(				self	:       Union[str, Any]   )						->					List[str]:
									"""simple docstring"""







									super().test_cpu_offload_forward_pass(expected_max_diff=3E-3   )





			def 		snake_case__			(				self	:       Optional[int]   )						->					Optional[Any]:
									"""simple docstring"""







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





			def 		snake_case__			(				self	:       int   )						->					List[str]:
									"""simple docstring"""







									super().test_inference_batch_single_identical(expected_max_diff=7E-3   )





			def 		snake_case__			(				self	:       Dict   )						->					Union[str, Any]:
									"""simple docstring"""







									super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3E-3   )





			def 		snake_case__			(				self	:       Tuple   )						->					Optional[Any]:
									"""simple docstring"""







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





			def 		snake_case__			(				self	:       Dict   )						->					List[Any]:
									"""simple docstring"""







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







			def 		snake_case__			(				self	:       Dict   )						->					Any:
									"""simple docstring"""







									__lowercase					=				[
									    "DDIMScheduler",
									    "DDPMScheduler",
									    "PNDMScheduler",
									    "HeunDiscreteScheduler",
									    "EulerAncestralDiscreteScheduler",
									    "KDPM2DiscreteScheduler",
									    "KDPM2AncestralDiscreteScheduler",
									    "DPMSolverSDEScheduler",
									]
									__lowercase					=				self.get_dummy_components()
									__lowercase					=				self.pipeline_class(**UpperCamelCase_   )

									# make sure that PNDM does not need warm-up
									pipe.scheduler.register_to_config(skip_prk_steps=UpperCamelCase_   )

									pipe.to(UpperCamelCase_   )
									pipe.set_progress_bar_config(disable=UpperCamelCase_   )
									__lowercase					=				self.get_dummy_inputs(UpperCamelCase_   )
									__lowercase					=				2

									__lowercase					=				[]
									for scheduler_enum in KarrasDiffusionSchedulers:
															if scheduler_enum.name in skip_schedulers:
																					# no sigma schedulers are not supported
																					# no schedulers
																					continue

															__lowercase					=				getattr(UpperCamelCase_			,     scheduler_enum.name   )
															__lowercase					=				scheduler_cls.from_config(pipe.scheduler.config   )
															__lowercase					=				pipe(**UpperCamelCase_   )[0]
															outputs.append(UpperCamelCase_   )

									assert check_same_shape(UpperCamelCase_   )





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






			"""simple docstring"""





			def 		snake_case__			(				self	:       Optional[Any]   )						->					str:
									"""simple docstring"""







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





			def 		snake_case__			(				self	:       str   )						->					Tuple:
									"""simple docstring"""







									__lowercase					=				torch.manual_seed(33   )

									__lowercase					=				StableDiffusionPipeline.from_pretrained("""CompVis/stable-diffusion-v1-4"""			,     torch_dtype=torch.floataa   )
									pipe.to("""cuda"""   )

									__lowercase					=				StableDiffusionLatentUpscalePipeline.from_pretrained(
									    """stabilityai/sd-x2-latent-upscaler"""			,     torch_dtype=torch.floataa   )
									upscaler.to("""cuda"""   )

									__lowercase					=				"a photo of an astronaut high resolution, unreal engine, ultra realistic"

									__lowercase					=				pipe(UpperCamelCase_			,     generator=UpperCamelCase_			,     output_type="""latent"""   ).images

									__lowercase					=				upscaler(
									    prompt=UpperCamelCase_			,     image=UpperCamelCase_			,     num_inference_steps=20			,     guidance_scale=0			,     generator=UpperCamelCase_			,     output_type="""np"""			,     ).images[0]

									__lowercase					=				load_numpy(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/astronaut_1024.npy"""   )
									assert np.abs((expected_image - image).mean()   ) < 5E-2







			def 		snake_case__			(				self	:       List[str]   )						->					str:
									"""simple docstring"""







									__lowercase					=				torch.manual_seed(33   )

									__lowercase					=				StableDiffusionLatentUpscalePipeline.from_pretrained(
									    """stabilityai/sd-x2-latent-upscaler"""			,     torch_dtype=torch.floataa   )
									upscaler.to("""cuda"""   )

									__lowercase					=				"the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas"

									__lowercase					=				load_image(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_512.png"""   )

									__lowercase					=				upscaler(
									    prompt=UpperCamelCase_			,     image=UpperCamelCase_			,     num_inference_steps=20			,     guidance_scale=0			,     generator=UpperCamelCase_			,     output_type="""np"""			,     ).images[0]

									__lowercase					=				load_numpy(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_1024.npy"""   )
									assert np.abs((expected_image - image).max()   ) < 5E-2










import random






def  UpperCAmelCase__    (    lowercase__						,       lowercase__						,       lowercase__ = False   ) ->      dict:
						__lowercase					=				{i: [] for i in range(lowercase__   )}

						# if probability is greater or equal than 1, then generate a complete graph
						if probability >= 1:
												return complete_graph(lowercase__   )
						# if probability is lower or equal than 0, then return a graph without edges
						if probability <= 0:
												return graph

						# for each couple of nodes, add an edge from u to v
						# if the number randomly generated is greater than probability probability
						for i in range(lowercase__   ):
												for j in range(i + 1						,       lowercase__   ):
																		if random.random() < probability:
																								graph[i].append(lowercase__   )
																								if not directed:
																														# if the graph is undirected, add an edge in from j to i, either
																														graph[j].append(lowercase__   )
						return graph






def  UpperCAmelCase__    (    lowercase__   ) ->      dict:
						return {
						    i: [j for j in range(lowercase__   ) if i != j] for i in range(lowercase__   )
						}


if __name__ == "__main__":
			import doctest

			doctest.testmod()





