




import gc
import random
import unittest

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class    _a     (				lowercase__					,							lowercase__					,							unittest.TestCase							):






							"""simple docstring"""

							snake_case_							 =						StableDiffusionXLImgaImgPipeline
							snake_case_							 =						TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
							snake_case_							 =						PipelineTesterMixin.required_optional_params - {"latents"}
							snake_case_							 =						TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
							snake_case_							 =						IMAGE_TO_IMAGE_IMAGE_PARAMS
							snake_case_							 =						IMAGE_TO_IMAGE_IMAGE_PARAMS

							def      A_  (      self				:     Optional[Any]    )	->Union[str, Any]:
													torch.manual_seed(0    )
													SCREAMING_SNAKE_CASE__     :				Dict   =					UNetaDConditionModel(
													    block_out_channels=(32, 64)							,	layers_per_block=2							,	sample_size=32							,	in_channels=4							,	out_channels=4							,	down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")							,	up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")							,	attention_head_dim=(2, 4)							,	use_linear_projection=__lowerCAmelCase							,	addition_embed_type="text_time"							,	addition_time_embed_dim=8							,	transformer_layers_per_block=(1, 2)							,	projection_class_embeddings_input_dim=80							,	cross_attention_dim=64							,	)
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					EulerDiscreteScheduler(
													    beta_start=0.0_0085							,	beta_end=0.012							,	steps_offset=1							,	beta_schedule="scaled_linear"							,	timestep_spacing="leading"							,	)
													torch.manual_seed(0    )
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					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=1_28							,	)
													torch.manual_seed(0    )
													SCREAMING_SNAKE_CASE__     :				Tuple   =					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=10_00							,	hidden_act="gelu"							,	projection_dim=32							,	)
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					CLIPTextModel(__lowerCAmelCase    )
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"							,	local_files_only=__lowerCAmelCase    )

													SCREAMING_SNAKE_CASE__     :				int   =					CLIPTextModelWithProjection(__lowerCAmelCase    )
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"							,	local_files_only=__lowerCAmelCase    )

													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					{
													    "unet": unet,
													    "scheduler": scheduler,
													    "vae": vae,
													    "text_encoder": text_encoder,
													    "tokenizer": tokenizer,
													    "text_encoder_2": text_encoder_a,
													    "tokenizer_2": tokenizer_a,
													    # "safety_checker": None,
													    # "feature_extractor": None,
													}
													return components

							def      A_  (      self				:     Any							,	a				:     List[Any]							,	a				:     str=0    )	->Optional[Any]:
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					floats_tensor((1, 3, 32, 32)							,	rng=random.Random(__lowerCAmelCase    )    ).to(__lowerCAmelCase    )
													SCREAMING_SNAKE_CASE__     :				Dict   =					image / 2 + 0.5
													if str(__lowerCAmelCase    ).startswith("mps"    ):
																			SCREAMING_SNAKE_CASE__     :				int   =					torch.manual_seed(__lowerCAmelCase    )
													else:
																			SCREAMING_SNAKE_CASE__     :				List[Any]   =					torch.Generator(device=__lowerCAmelCase    ).manual_seed(__lowerCAmelCase    )
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					{
													    "prompt": "A painting of a squirrel eating a burger",
													    "image": image,
													    "generator": generator,
													    "num_inference_steps": 2,
													    "guidance_scale": 5.0,
													    "output_type": "numpy",
													    "strength": 0.75,
													}
													return inputs

							def      A_  (      self				:     str    )	->str:
													SCREAMING_SNAKE_CASE__     :				Tuple   =					"cpu"  # ensure determinism for the device-dependent torch.Generator
													SCREAMING_SNAKE_CASE__     :				List[str]   =					self.get_dummy_components()
													SCREAMING_SNAKE_CASE__     :				str   =					StableDiffusionXLImgaImgPipeline(**__lowerCAmelCase    )
													SCREAMING_SNAKE_CASE__     :				List[str]   =					sd_pipe.to(__lowerCAmelCase    )
													sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase    )

													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					self.get_dummy_inputs(__lowerCAmelCase    )
													SCREAMING_SNAKE_CASE__     :				Dict   =					sd_pipe(**__lowerCAmelCase    ).images
													SCREAMING_SNAKE_CASE__     :				Tuple   =					image[0, -3:, -3:, -1]

													assert image.shape == (1, 32, 32, 3)

													SCREAMING_SNAKE_CASE__     :				Dict   =					np.array([0.4656, 0.4840, 0.4439, 0.6698, 0.5574, 0.4524, 0.5799, 0.5943, 0.5165]    )

													assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2

							def      A_  (      self				:     Tuple    )	->Tuple:
													super().test_attention_slicing_forward_pass(expected_max_diff=3E-3    )

							def      A_  (      self				:     Tuple    )	->Any:
													super().test_inference_batch_single_identical(expected_max_diff=3E-3    )

							def      A_  (      self				:     Optional[Any]    )	->Any:
													pass






							def      A_  (      self				:     Tuple    )	->Optional[Any]:
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					self.get_dummy_components()
													SCREAMING_SNAKE_CASE__     :				Tuple   =					StableDiffusionXLImgaImgPipeline(**__lowerCAmelCase    )
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					sd_pipe.to(__lowerCAmelCase    )
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					sd_pipe.to(__lowerCAmelCase    )
													sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase    )

													# forward without prompt embeds
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					self.get_dummy_inputs(__lowerCAmelCase    )
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					3 * ["this is a negative prompt"]
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					negative_prompt
													SCREAMING_SNAKE_CASE__     :				str   =					3 * [inputs["prompt"]]

													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					sd_pipe(**__lowerCAmelCase    )
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					output.images[0, -3:, -3:, -1]

													# forward with prompt embeds
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					self.get_dummy_inputs(__lowerCAmelCase    )
													SCREAMING_SNAKE_CASE__     :				Dict   =					3 * ["this is a negative prompt"]
													SCREAMING_SNAKE_CASE__     :				Dict   =					3 * [inputs.pop("prompt"    )]

													(
													    (
													    SCREAMING_SNAKE_CASE__			
													),							(
													    SCREAMING_SNAKE_CASE__			
													),							(
													    SCREAMING_SNAKE_CASE__			
													),							(
													    SCREAMING_SNAKE_CASE__			
													),										
													)     :				Dict   =					sd_pipe.encode_prompt(__lowerCAmelCase							,	negative_prompt=__lowerCAmelCase    )

													SCREAMING_SNAKE_CASE__     :				Tuple   =					sd_pipe(
													    **__lowerCAmelCase							,	prompt_embeds=__lowerCAmelCase							,	negative_prompt_embeds=__lowerCAmelCase							,	pooled_prompt_embeds=__lowerCAmelCase							,	negative_pooled_prompt_embeds=__lowerCAmelCase							,	)
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					output.images[0, -3:, -3:, -1]

													# make sure that it's equal
													assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()    ).max() < 1E-4







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






							"""simple docstring"""

							def      A_  (      self				:     Optional[int]    )	->Dict:
													super().tearDown()
													gc.collect()
													torch.cuda.empty_cache()

							def      A_  (      self				:     List[str]							,	a				:     Optional[int]							,	a				:     Any="cpu"							,	a				:     int=torch.floataa							,	a				:     List[Any]=0    )	->Dict:
													SCREAMING_SNAKE_CASE__     :				int   =					torch.Generator(device=__lowerCAmelCase    ).manual_seed(__lowerCAmelCase    )
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					np.random.RandomState(__lowerCAmelCase    ).standard_normal((1, 4, 64, 64)    )
													SCREAMING_SNAKE_CASE__     :				Dict   =					torch.from_numpy(__lowerCAmelCase    ).to(device=__lowerCAmelCase							,	dtype=__lowerCAmelCase    )
													SCREAMING_SNAKE_CASE__     :				List[str]   =					{
													    "prompt": "a photograph of an astronaut riding a horse",
													    "latents": latents,
													    "generator": generator,
													    "num_inference_steps": 3,
													    "guidance_scale": 7.5,
													    "output_type": "numpy",
													}
													return inputs






							def      A_  (      self				:     str    )	->Tuple:
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-base"    )
													pipe.to(__lowerCAmelCase    )
													pipe.set_progress_bar_config(disable=__lowerCAmelCase    )

													SCREAMING_SNAKE_CASE__     :				List[Any]   =					self.get_inputs(__lowerCAmelCase    )
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					pipe(**__lowerCAmelCase    ).images
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					image[0, -3:, -3:, -1].flatten()

													assert image.shape == (1, 5_12, 5_12, 3)
													SCREAMING_SNAKE_CASE__     :				int   =					np.array([0.4_9493, 0.4_7896, 0.4_0798, 0.5_4214, 0.5_3212, 0.4_8202, 0.4_7656, 0.4_6329, 0.4_8506]    )
													assert np.abs(image_slice - expected_slice    ).max() < 7E-3






def       UpperCAmelCase							(     _lowerCamelCase      :  int			,   _lowerCamelCase      :  bool = False     ):





						'''simple docstring'''


						if n == 2:
												return True
						if not n % 2 or n < 2:
												return False
						if n > 5 and n % 10 not in (1, 3, 7, 9):  # can quickly check last digit
												return False
						if n > 3_317_044_064_679_887_385_961_981 and not allow_probable:
												raise ValueError(
												    "Warning: upper bound of deterministic test is exceeded. "
												    "Pass allow_probable=True to allow probabilistic test. "
												    "A return value of True indicates a probable prime."     )
						# array bounds provided by analysis
						SCREAMING_SNAKE_CASE__     :				List[str]   =					[
						    2_047,
						    1_373_653,
						    25_326_001,
						    3_215_031_751,
						    2_152_302_898_747,
						    3_474_749_660_383,
						    341_550_071_728_321,
						    1,
						    3_825_123_056_546_413_051,
						    1,
						    1,
						    318_665_857_834_031_151_167_461,
						    3_317_044_064_679_887_385_961_981,
						]

						SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41]
						for idx, _p in enumerate(_lowerCamelCase			,   1     ):
												if n < _p:
																		# then we have our last prime to check
																		SCREAMING_SNAKE_CASE__     :				Dict   =					primes[:idx]
																		break
						SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				Dict   =					n - 1, 0
						# break up n -1 into a power of 2 (s) and
						# remaining odd component
						# essentially, solve for d * 2 ** s == n - 1
						while d % 2 == 0:
												d //= 2
												s += 1
						for prime in plist:
												SCREAMING_SNAKE_CASE__     :				str   =					False
												for r in range(_lowerCamelCase     ):
																		SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					pow(_lowerCamelCase			,   d * 2**r			,   _lowerCamelCase     )
																		# see article for analysis explanation for m
																		if (r == 0 and m == 1) or ((m + 1) % n == 0):
																								SCREAMING_SNAKE_CASE__     :				str   =					True
																								# this loop will not determine compositeness
																								break
												if pr:
																		continue
												# if pr is False, then the above loop never evaluated to true,
												# and the n MUST be composite
												return False
						return True







def       UpperCAmelCase							(     ):





						'''simple docstring'''


						assert not miller_rabin(561     )
						assert miller_rabin(563     )
						# 2047

						assert not miller_rabin(838_201     )
						assert miller_rabin(838_207     )
						# 1_373_653

						assert not miller_rabin(17_316_001     )
						assert miller_rabin(17_316_017     )
						# 25_326_001

						assert not miller_rabin(3_078_386_641     )
						assert miller_rabin(3_078_386_653     )
						# 3_215_031_751

						assert not miller_rabin(1_713_045_574_801     )
						assert miller_rabin(1_713_045_574_819     )
						# 2_152_302_898_747

						assert not miller_rabin(2_779_799_728_307     )
						assert miller_rabin(2_779_799_728_327     )
						# 3_474_749_660_383

						assert not miller_rabin(113_850_023_909_441     )
						assert miller_rabin(113_850_023_909_527     )
						# 341_550_071_728_321

						assert not miller_rabin(1_275_041_018_848_804_351     )
						assert miller_rabin(1_275_041_018_848_804_391     )
						# 3_825_123_056_546_413_051

						assert not miller_rabin(79_666_464_458_507_787_791_867     )
						assert miller_rabin(79_666_464_458_507_787_791_951     )
						# 318_665_857_834_031_151_167_461

						assert not miller_rabin(552_840_677_446_647_897_660_333     )
						assert miller_rabin(552_840_677_446_647_897_660_359     )
						# 3_317_044_064_679_887_385_961_981
						# upper limit for probabilistic test


if __name__ == "__main__":
				test_miller_rabin()