

"""simple docstring"""







import gc
import unittest

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

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


__SCREAMING_SNAKE_CASE								=False







class  UpperCamelCase	(    unittest.TestCase						):
   def     _UpperCAmelCase     (    self    )			->   Any:


        '''simple docstring'''


        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()
   @property
   def     _UpperCAmelCase     (    self    )			->   Dict:


        '''simple docstring'''


        return 12
   @property
   def     _UpperCAmelCase     (    self    )			->   List[Any]:


        '''simple docstring'''


        return 12
   @property
   def     _UpperCAmelCase     (    self    )			->   Dict:


        '''simple docstring'''


        return 32
   @property
   def     _UpperCAmelCase     (    self    )			->   List[str]:


        '''simple docstring'''


        torch.manual_seed(0    )
        lowercase_			:       Dict									=			VQModel(
            block_out_channels=[32, 64]  ,in_channels=3  ,out_channels=3  ,down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']  ,up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']  ,latent_channels=3  ,num_vq_embeddings=self.num_embed  ,vq_embed_dim=3  ,)
        return model
   @property
   def     _UpperCAmelCase     (    self    )			->   Tuple:


        '''simple docstring'''


        lowercase_			:       str									=			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'    )
        return tokenizer
   @property
   def     _UpperCAmelCase     (    self    )			->   List[Any]:


        '''simple docstring'''


        torch.manual_seed(0    )
        lowercase_			:       int									=			CLIPTextConfig(
            bos_token_id=0  ,eos_token_id=2  ,hidden_size=self.text_embedder_hidden_size  ,intermediate_size=37  ,layer_norm_eps=1e-05  ,num_attention_heads=4  ,num_hidden_layers=5  ,pad_token_id=1  ,vocab_size=1000  ,)
        return CLIPTextModel(__UpperCamelCase    )
   @property
   def     _UpperCAmelCase     (    self    )			->   List[str]:


        '''simple docstring'''


        torch.manual_seed(0    )

        lowercase_			:       List[Any]									=			12
        lowercase_			:       Tuple									=			12

        lowercase_			:       Optional[int]									=			{
            'attention_bias': True,
            'cross_attention_dim': 32,
            'attention_head_dim': height * width,
            'num_attention_heads': 1,
            'num_vector_embeds': self.num_embed,
            'num_embeds_ada_norm': self.num_embeds_ada_norm,
            'norm_num_groups': 32,
            'sample_size': width,
            'activation_fn': 'geglu-approximate',
        }

        lowercase_			:       Optional[int]									=			TransformeraDModel(**__UpperCamelCase    )
        return model
   def     _UpperCAmelCase     (    self    )			->   str:


        '''simple docstring'''


        lowercase_			:       List[str]									=			'cpu'

        lowercase_			:       Any									=			self.dummy_vqvae
        lowercase_			:       Dict									=			self.dummy_text_encoder
        lowercase_			:       int									=			self.dummy_tokenizer
        lowercase_			:       Dict									=			self.dummy_transformer
        lowercase_			:       Tuple									=			VQDiffusionScheduler(self.num_embed    )
        lowercase_			:       Tuple									=			LearnedClassifierFreeSamplingEmbeddings(learnable=__UpperCamelCase    )

        lowercase_			:       Union[str, Any]									=			VQDiffusionPipeline(
            vqvae=__UpperCamelCase  ,text_encoder=__UpperCamelCase  ,tokenizer=__UpperCamelCase  ,transformer=__UpperCamelCase  ,scheduler=__UpperCamelCase  ,learned_classifier_free_sampling_embeddings=__UpperCamelCase  ,)
        lowercase_			:       List[str]									=			pipe.to(__UpperCamelCase    )
        pipe.set_progress_bar_config(disable=__UpperCamelCase    )

        lowercase_			:       int									=			'teddy bear playing in the pool'

        lowercase_			:       int									=			torch.Generator(device=__UpperCamelCase    ).manual_seed(0    )
        lowercase_			:       Dict									=			pipe([prompt]  ,generator=__UpperCamelCase  ,num_inference_steps=2  ,output_type='np'    )
        lowercase_			:       Union[str, Any]									=			output.images

        lowercase_			:       List[Any]									=			torch.Generator(device=__UpperCamelCase    ).manual_seed(0    )
        lowercase_			:       Union[str, Any]									=			pipe(
            [prompt]  ,generator=__UpperCamelCase  ,output_type='np'  ,return_dict=__UpperCamelCase  ,num_inference_steps=2    )[0]

        lowercase_			:       Union[str, Any]									=			image[0, -3:, -3:, -1]
        lowercase_			:       Dict									=			image_from_tuple[0, -3:, -3:, -1]

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

        lowercase_			:       Optional[int]									=			np.array([0.6551, 0.6168, 0.5008, 0.5676, 0.5659, 0.4295, 0.6073, 0.5599, 0.4992]    )

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







   def     _UpperCAmelCase     (    self    )			->   str:


        '''simple docstring'''


        lowercase_			:       Optional[int]									=			'cpu'

        lowercase_			:       int									=			self.dummy_vqvae
        lowercase_			:       str									=			self.dummy_text_encoder
        lowercase_			:       Union[str, Any]									=			self.dummy_tokenizer
        lowercase_			:       int									=			self.dummy_transformer
        lowercase_			:       Tuple									=			VQDiffusionScheduler(self.num_embed    )
        lowercase_			:       Optional[Any]									=			LearnedClassifierFreeSamplingEmbeddings(
            learnable=__UpperCamelCase  ,hidden_size=self.text_embedder_hidden_size  ,length=tokenizer.model_max_length    )

        lowercase_			:       int									=			VQDiffusionPipeline(
            vqvae=__UpperCamelCase  ,text_encoder=__UpperCamelCase  ,tokenizer=__UpperCamelCase  ,transformer=__UpperCamelCase  ,scheduler=__UpperCamelCase  ,learned_classifier_free_sampling_embeddings=__UpperCamelCase  ,)
        lowercase_			:       List[str]									=			pipe.to(__UpperCamelCase    )
        pipe.set_progress_bar_config(disable=__UpperCamelCase    )

        lowercase_			:       Tuple									=			'teddy bear playing in the pool'

        lowercase_			:       int									=			torch.Generator(device=__UpperCamelCase    ).manual_seed(0    )
        lowercase_			:       List[Any]									=			pipe([prompt]  ,generator=__UpperCamelCase  ,num_inference_steps=2  ,output_type='np'    )
        lowercase_			:       Dict									=			output.images

        lowercase_			:       Optional[int]									=			torch.Generator(device=__UpperCamelCase    ).manual_seed(0    )
        lowercase_			:       int									=			pipe(
            [prompt]  ,generator=__UpperCamelCase  ,output_type='np'  ,return_dict=__UpperCamelCase  ,num_inference_steps=2    )[0]

        lowercase_			:       List[str]									=			image[0, -3:, -3:, -1]
        lowercase_			:       int									=			image_from_tuple[0, -3:, -3:, -1]

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

        lowercase_			:       int									=			np.array([0.6693, 0.6075, 0.4959, 0.5701, 0.5583, 0.4333, 0.6171, 0.5684, 0.4988]    )

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







@slow
@require_torch_gpu
class  UpperCamelCase	(    unittest.TestCase						):
   def     _UpperCAmelCase     (    self    )			->   Optional[Any]:


        '''simple docstring'''


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







   def     _UpperCAmelCase     (    self    )			->   Dict:


        '''simple docstring'''


        lowercase_			:       Tuple									=			load_numpy(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'    )

        lowercase_			:       Any									=			VQDiffusionPipeline.from_pretrained('microsoft/vq-diffusion-ithq'    )
        lowercase_			:       List[Any]									=			pipeline.to(__UpperCamelCase    )
        pipeline.set_progress_bar_config(disable=__UpperCamelCase    )

        # requires GPU generator for gumbel softmax
        # don't use GPU generator in tests though
        lowercase_			:       Tuple									=			torch.Generator(device=__UpperCamelCase    ).manual_seed(0    )
        lowercase_			:       Tuple									=			pipeline(
            'teddy bear playing in the pool'  ,num_images_per_prompt=1  ,generator=__UpperCamelCase  ,output_type='np'  ,)

        lowercase_			:       Union[str, Any]									=			output.images[0]

        assert image.shape == (256, 256, 3)
        assert np.abs(expected_image - image    ).max() < 2.0


"""simple docstring"""







from math import sqrt



def     lowercase__( __SCREAMING_SNAKE_CASE							:  int   ):

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



def     lowercase__( __SCREAMING_SNAKE_CASE							:  int = 1_00_01   ):
     lowercase_			:       str									=			0
     lowercase_			:       Optional[Any]									=			1
     while count != nth and number < 3:
          number += 1
          if is_prime(__SCREAMING_SNAKE_CASE   ):
               count += 1
     while count != nth:
          number += 2
          if is_prime(__SCREAMING_SNAKE_CASE   ):
               count += 1
     return number


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

