






import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


class        __UpperCamelCase    (     __snake_case   ,			unittest.TestCase   ):
	__snake_case	:Union[str, Any]   =       ShapEImgaImgPipeline
	__snake_case	:Tuple   =       ['image']
	__snake_case	:Optional[Any]   =       ['image']
	__snake_case	:Any   =       [
	    'num_images_per_prompt',
	    'num_inference_steps',
	    'generator',
	    'latents',
	    'guidance_scale',
	    'frame_size',
	    'output_type',
	    'return_dict',
	]
	__snake_case	:int   =       False







	@property
	def  _a   (  self      :							Tuple			)		-> Optional[int]:
					"""simple docstring"""







					return 32







	@property
	def  _a   (  self      :							Optional[Any]			)		-> Union[str, Any]:
					"""simple docstring"""







					return 32







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







					return self.time_input_dim * 4







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







					return 8







	@property
	def  _a   (  self      :							str			)		-> Dict:
					"""simple docstring"""







					torch.manual_seed(0			)
					__lowercase             =  CLIPVisionConfig(
					    hidden_size=self.text_embedder_hidden_size       ,				image_size=64       ,				projection_dim=self.text_embedder_hidden_size       ,				intermediate_size=37       ,				num_attention_heads=4       ,				num_channels=3       ,				num_hidden_layers=5       ,				patch_size=1       ,				)

					__lowercase             =  CLIPVisionModel(__UpperCamelCase			)
					return model







	@property
	def  _a   (  self      :							Tuple			)		-> str:
					"""simple docstring"""







					__lowercase             =  CLIPImageProcessor(
					    crop_size=224       ,				do_center_crop=__UpperCamelCase       ,				do_normalize=__UpperCamelCase       ,				do_resize=__UpperCamelCase       ,				image_mean=[0.48_145_466, 0.4_578_275, 0.40_821_073]       ,				image_std=[0.26_862_954, 0.26_130_258, 0.27_577_711]       ,				resample=3       ,				size=224       ,				)

					return image_processor







	@property
	def  _a   (  self      :							Dict			)		-> int:
					"""simple docstring"""







					torch.manual_seed(0			)

					__lowercase             =  {
					    """num_attention_heads""": 2,
					    """attention_head_dim""": 16,
					    """embedding_dim""": self.time_input_dim,
					    """num_embeddings""": 32,
					    """embedding_proj_dim""": self.text_embedder_hidden_size,
					    """time_embed_dim""": self.time_embed_dim,
					    """num_layers""": 1,
					    """clip_embed_dim""": self.time_input_dim * 2,
					    """additional_embeddings""": 0,
					    """time_embed_act_fn""": """gelu""",
					    """norm_in_type""": """layer""",
					    """embedding_proj_norm_type""": """layer""",
					    """encoder_hid_proj_type""": None,
					    """added_emb_type""": None,
					}

					__lowercase             =  PriorTransformer(**__UpperCamelCase			)
					return model







	@property
	def  _a   (  self      :							List[str]			)		-> int:
					"""simple docstring"""







					torch.manual_seed(0			)

					__lowercase             =  {
					    """param_shapes""": (
					        (self.renderer_dim, 93),
					        (self.renderer_dim, 8),
					        (self.renderer_dim, 8),
					        (self.renderer_dim, 8),
					    ),
					    """d_latent""": self.time_input_dim,
					    """d_hidden""": self.renderer_dim,
					    """n_output""": 12,
					    """background""": (
					        0.1,
					        0.1,
					        0.1,
					    ),
					}
					__lowercase             =  ShapERenderer(**__UpperCamelCase			)
					return model







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







					__lowercase             =  self.dummy_prior
					__lowercase             =  self.dummy_image_encoder
					__lowercase             =  self.dummy_image_processor
					__lowercase             =  self.dummy_renderer

					__lowercase             =  HeunDiscreteScheduler(
					    beta_schedule="""exp"""       ,				num_train_timesteps=1024       ,				prediction_type="""sample"""       ,				use_karras_sigmas=__UpperCamelCase       ,				clip_sample=__UpperCamelCase       ,				clip_sample_range=1.0       ,				)
					__lowercase             =  {
					    """prior""": prior,
					    """image_encoder""": image_encoder,
					    """image_processor""": image_processor,
					    """renderer""": renderer,
					    """scheduler""": scheduler,
					}

					return components







	def  _a   (  self      :							Union[str, Any]       ,				_lowerCAmelCase      :							int       ,				_lowerCAmelCase      :							Optional[Any]=0			)		-> Union[str, Any]:
					"""simple docstring"""







					__lowercase             =  floats_tensor((1, 3, 64, 64)       ,				rng=random.Random(__UpperCamelCase			)			).to(__UpperCamelCase			)

					if str(__UpperCamelCase			).startswith("""mps"""			):
									__lowercase             =  torch.manual_seed(__UpperCamelCase			)
					else:
									__lowercase             =  torch.Generator(device=__UpperCamelCase			).manual_seed(__UpperCamelCase			)
					__lowercase             =  {
					    """image""": input_image,
					    """generator""": generator,
					    """num_inference_steps""": 1,
					    """frame_size""": 32,
					    """output_type""": """np""",
					}
					return inputs







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







					__lowercase             =  """cpu"""

					__lowercase             =  self.get_dummy_components()

					__lowercase             =  self.pipeline_class(**__UpperCamelCase			)
					__lowercase             =  pipe.to(__UpperCamelCase			)

					pipe.set_progress_bar_config(disable=__UpperCamelCase			)

					__lowercase             =  pipe(**self.get_dummy_inputs(__UpperCamelCase			)			)
					__lowercase             =  output.images[0]
					__lowercase             =  image[0, -3:, -3:, -1]

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

					__lowercase             =  np.array(
					    [
					        0.00_039_216,
					        0.00_039_216,
					        0.00_039_216,
					        0.00_039_216,
					        0.00_039_216,
					        0.00_039_216,
					        0.00_039_216,
					        0.00_039_216,
					        0.00_039_216,
					    ]			)

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







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







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







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







					__lowercase             =  torch_device == """cpu"""
					__lowercase             =  True
					self._test_inference_batch_single_identical(
					    batch_size=2       ,				test_max_difference=__UpperCamelCase       ,				relax_max_difference=__UpperCamelCase       ,				)







	def  _a   (  self      :							int			)		-> int:
					"""simple docstring"""







					__lowercase             =  self.get_dummy_components()
					__lowercase             =  self.pipeline_class(**__UpperCamelCase			)
					__lowercase             =  pipe.to(__UpperCamelCase			)
					pipe.set_progress_bar_config(disable=__UpperCamelCase			)

					__lowercase             =  1
					__lowercase             =  2

					__lowercase             =  self.get_dummy_inputs(__UpperCamelCase			)

					for key in inputs.keys():
									if key in self.batch_params:
													__lowercase             =  batch_size * [inputs[key]]

					__lowercase             =  pipe(**__UpperCamelCase       ,				num_images_per_prompt=__UpperCamelCase			)[0]

					assert images.shape[0] == batch_size * num_images_per_prompt


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







	def  _a   (  self      :							int			)		-> int:
					"""simple docstring"""







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







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







					__lowercase             =  load_image(
					    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main""" """/shap_e/corgi.png"""			)
					__lowercase             =  load_numpy(
					    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
					    """/shap_e/test_shap_e_img2img_out.npy"""			)
					__lowercase             =  ShapEImgaImgPipeline.from_pretrained("""openai/shap-e-img2img"""			)
					__lowercase             =  pipe.to(__UpperCamelCase			)
					pipe.set_progress_bar_config(disable=__UpperCamelCase			)

					__lowercase             =  torch.Generator(device=__UpperCamelCase			).manual_seed(0			)

					__lowercase             =  pipe(
					    __UpperCamelCase       ,				generator=__UpperCamelCase       ,				guidance_scale=3.0       ,				num_inference_steps=64       ,				frame_size=64       ,				output_type="""np"""       ,				).images[0]

					assert images.shape == (20, 64, 64, 3)

					assert_mean_pixel_difference(__UpperCamelCase       ,				__UpperCamelCase			)









from math import sqrt





def 				snake_case  (     lowerCamelCase					):




	'''simple docstring'''







	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(sqrt(lowerCamelCase					) + 1					)				,  6					):
		if number % i == 0 or number % (i + 2) == 0:
			return False
	return True





def 				snake_case  (     lowerCamelCase = 10_001					):




	'''simple docstring'''







	__lowercase             =  0
	__lowercase             =  1
	while count != nth and number < 3:
		number += 1
		if is_prime(lowerCamelCase					):
			count += 1
	while count != nth:
		number += 2
		if is_prime(lowerCamelCase					):
			count += 1
	return number


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

