





import gc
import random
import unittest

import numpy as np
import torch

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
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 ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()


class 			lowercase				(							a_		,			unittest.TestCase     ):



						"""simple docstring"""

						_UpperCamelCase  :	Optional[Any]								=	KandinskyVaaControlnetPipeline
						_UpperCamelCase  :	List[Any]								=	["image_embeds", "negative_image_embeds", "hint"]
						_UpperCamelCase  :	int								=	["image_embeds", "negative_image_embeds", "hint"]
						_UpperCamelCase  :	int								=	[
						    "generator",
						    "height",
						    "width",
						    "latents",
						    "guidance_scale",
						    "num_inference_steps",
						    "return_dict",
						    "guidance_scale",
						    "num_images_per_prompt",
						    "output_type",
						    "return_dict",
						]
						_UpperCamelCase  :	Any								=	False







						@property
						def 		__UpperCAmelCase       (    self			:		Any					):





										'''simple docstring'''







										return 32







						@property
						def 		__UpperCAmelCase       (    self			:		Union[str, Any]					):





										'''simple docstring'''







										return 32







						@property
						def 		__UpperCAmelCase       (    self			:		List[str]					):





										'''simple docstring'''







										return self.time_input_dim







						@property
						def 		__UpperCAmelCase       (    self			:		List[Any]					):





										'''simple docstring'''







										return self.time_input_dim * 4







						@property
						def 		__UpperCAmelCase       (    self			:		Dict					):





										'''simple docstring'''







										return 1_00







						@property
						def 		__UpperCAmelCase       (    self			:		Optional[Any]					):





										'''simple docstring'''







										torch.manual_seed(0					)

										_snake_case :			Optional[int]  						=   {
										    'in_channels': 8,
										    # Out channels is double in channels because predicts mean and variance
										    'out_channels': 8,
										    'addition_embed_type': 'image_hint',
										    'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
										    'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
										    'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
										    'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
										    'layers_per_block': 1,
										    'encoder_hid_dim': self.text_embedder_hidden_size,
										    'encoder_hid_dim_type': 'image_proj',
										    'cross_attention_dim': self.cross_attention_dim,
										    'attention_head_dim': 4,
										    'resnet_time_scale_shift': 'scale_shift',
										    'class_embed_type': None,
										}

										_snake_case :			str  						=   UNetaDConditionModel(**lowerCamelCase_					)
										return model







						@property
						def 		__UpperCAmelCase       (    self			:		int					):





										'''simple docstring'''







										return {
										    "block_out_channels": [32, 32, 64, 64],
										    "down_block_types": [
										        "DownEncoderBlock2D",
										        "DownEncoderBlock2D",
										        "DownEncoderBlock2D",
										        "AttnDownEncoderBlock2D",
										    ],
										    "in_channels": 3,
										    "latent_channels": 4,
										    "layers_per_block": 1,
										    "norm_num_groups": 8,
										    "norm_type": "spatial",
										    "num_vq_embeddings": 12,
										    "out_channels": 3,
										    "up_block_types": ["AttnUpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"],
										    "vq_embed_dim": 4,
										}







						@property
						def 		__UpperCAmelCase       (    self			:		Union[str, Any]					):





										'''simple docstring'''







										torch.manual_seed(0					)
										_snake_case :			Any  						=   VQModel(**self.dummy_movq_kwargs					)
										return model







						def 		__UpperCAmelCase       (    self			:		Dict					):





										'''simple docstring'''







										_snake_case :			Dict  						=   self.dummy_unet
										_snake_case :			int  						=   self.dummy_movq

										_snake_case :			Any  						=   DDIMScheduler(
										    num_train_timesteps=10_00 ,				beta_schedule='linear' ,				beta_start=0.0_0085 ,				beta_end=0.012 ,				clip_sample=lowerCamelCase_ ,				set_alpha_to_one=lowerCamelCase_ ,				steps_offset=1 ,				prediction_type='epsilon' ,				thresholding=lowerCamelCase_ ,				)

										_snake_case :			Optional[Any]  						=   {
										    'unet': unet,
										    'scheduler': scheduler,
										    'movq': movq,
										}
										return components







						def 		__UpperCAmelCase       (    self			:		Tuple ,				lowerCamelCase_			:		Any ,				lowerCamelCase_			:		Dict=0					):





										'''simple docstring'''







										_snake_case :			str  						=   floats_tensor((1, self.text_embedder_hidden_size) ,				rng=random.Random(lowerCamelCase_					)					).to(lowerCamelCase_					)
										_snake_case :			int  						=   floats_tensor((1, self.text_embedder_hidden_size) ,				rng=random.Random(seed + 1					)					).to(
										    lowerCamelCase_					)

										# create hint
										_snake_case :			Union[str, Any]  						=   floats_tensor((1, 3, 64, 64) ,				rng=random.Random(lowerCamelCase_					)					).to(lowerCamelCase_					)

										if str(lowerCamelCase_					).startswith('mps'					):
														_snake_case :			int  						=   torch.manual_seed(lowerCamelCase_					)
										else:
														_snake_case :			int  						=   torch.Generator(device=lowerCamelCase_					).manual_seed(lowerCamelCase_					)
										_snake_case :			str  						=   {
										    'image_embeds': image_embeds,
										    'negative_image_embeds': negative_image_embeds,
										    'hint': hint,
										    'generator': generator,
										    'height': 64,
										    'width': 64,
										    'guidance_scale': 4.0,
										    'num_inference_steps': 2,
										    'output_type': 'np',
										}
										return inputs







						def 		__UpperCAmelCase       (    self			:		Optional[int]					):





										'''simple docstring'''







										_snake_case :			Any  						=   'cpu'

										_snake_case :			Optional[int]  						=   self.get_dummy_components()

										_snake_case :			Union[str, Any]  						=   self.pipeline_class(**lowerCamelCase_					)
										_snake_case :			Optional[Any]  						=   pipe.to(lowerCamelCase_					)

										pipe.set_progress_bar_config(disable=lowerCamelCase_					)

										_snake_case :			Tuple  						=   pipe(**self.get_dummy_inputs(lowerCamelCase_					)					)
										_snake_case :			List[Any]  						=   output.images

										_snake_case :			str  						=   pipe(
										    **self.get_dummy_inputs(lowerCamelCase_					) ,				return_dict=lowerCamelCase_ ,				)[0]

										_snake_case :			Optional[Any]  						=   image[0, -3:, -3:, -1]
										_snake_case :			Optional[Any]  						=   image_from_tuple[0, -3:, -3:, -1]

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

										_snake_case :			Tuple  						=   np.array(
										    [0.695_9826, 0.86_8279, 0.755_8092, 0.6876_9467, 0.8580_5804, 0.6597_7496, 0.4488_5302, 0.595_9111, 0.425_1595]					)

										assert (
										    np.abs(image_slice.flatten() - expected_slice					).max() < 1e-2
										), f''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''

										assert (
										    np.abs(image_from_tuple_slice.flatten() - expected_slice					).max() < 1e-2
										), f''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''



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



						"""simple docstring"""

						def 		__UpperCAmelCase       (    self			:		Dict					):





										'''simple docstring'''







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







						def 		__UpperCAmelCase       (    self			:		Dict					):





										'''simple docstring'''







										_snake_case :			Dict  						=   load_numpy(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
										    '/kandinskyv22/kandinskyv22_controlnet_robotcat_fp16.npy'					)

										_snake_case :			Optional[Any]  						=   load_image(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
										    '/kandinskyv22/hint_image_cat.png'					)
										_snake_case :			List[Any]  						=   torch.from_numpy(np.array(lowerCamelCase_					)					).float() / 255.0
										_snake_case :			Optional[int]  						=   hint.permute(2 ,				0 ,				1					).unsqueeze(0					)

										_snake_case :			Optional[int]  						=   KandinskyVaaPriorPipeline.from_pretrained(
										    'kandinsky-community/kandinsky-2-2-prior' ,				torch_dtype=torch.floataa					)
										pipe_prior.to(lowerCamelCase_					)

										_snake_case :			str  						=   KandinskyVaaControlnetPipeline.from_pretrained(
										    'kandinsky-community/kandinsky-2-2-controlnet-depth' ,				torch_dtype=torch.floataa					)
										_snake_case :			Any  						=   pipeline.to(lowerCamelCase_					)
										pipeline.set_progress_bar_config(disable=lowerCamelCase_					)

										_snake_case :			str  						=   'A robot, 4k photo'

										_snake_case :			Optional[Any]  						=   torch.Generator(device='cuda'					).manual_seed(0					)
										_snake_case						,			_snake_case :			List[str]  						=   pipe_prior(
										    lowerCamelCase_ ,				generator=lowerCamelCase_ ,				num_inference_steps=5 ,				negative_prompt='' ,				).to_tuple()

										_snake_case :			Optional[Any]  						=   torch.Generator(device='cuda'					).manual_seed(0					)
										_snake_case :			Optional[Any]  						=   pipeline(
										    image_embeds=lowerCamelCase_ ,				negative_image_embeds=lowerCamelCase_ ,				hint=lowerCamelCase_ ,				generator=lowerCamelCase_ ,				num_inference_steps=1_00 ,				output_type='np' ,				)

										_snake_case :			List[Any]  						=   output.images[0]

										assert image.shape == (5_12, 5_12, 3)

										assert_mean_pixel_difference(lowerCamelCase_ ,				lowerCamelCase_					)













import itertools
import math

def       A__(		__lowerCAmelCase ):

				if 1 < number < 4:
								# 2 and 3 are primes
								return True
				elif number < 2 or number % 2 == 0 or number % 3 == 0:
								# Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
								return False

				# All primes number are in format of 6k +/- 1
				for i in range(5	,			int(math.sqrt(__lowerCAmelCase ) + 1 )	,			6 ):
								if number % i == 0 or number % (i + 2) == 0:
												return False
				return True

def       A__(		):
				_snake_case :			Optional[Any]  						=   2
				while True:
								if is_prime(__lowerCAmelCase ):
												yield num
								num += 1

def       A__(		__lowerCAmelCase = 1_00_01 ):
				return next(itertools.islice(prime_generator()	,			nth - 1	,			__lowerCAmelCase ) )


if __name__ == "__main__":
		print(F'''{solution() = }''')






