

import math
def        UpperCAmelCase_     (		_A       ):





     '''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(math.sqrt(_A       ) + 1       )   ,	6       ):
          if number % i == 0 or number % (i + 2) == 0:
               return False
     return True
def        UpperCAmelCase_     (		_A = 1_00_01       ):





     '''simple docstring'''




     try:
          SCREAMING_SNAKE_CASE__					=					int(_A       )
     except (TypeError, ValueError):
          raise TypeError('''Parameter nth must be int or castable to int.'''       ) from None
     if nth <= 0:
          raise ValueError('''Parameter nth must be greater than or equal to one.'''       )
     SCREAMING_SNAKE_CASE__					=					[]
     SCREAMING_SNAKE_CASE__					=					2
     while len(_A       ) < nth:
          if is_prime(_A       ):
               primes.append(_A       )
               num += 1
          else:
               num += 1
     return primes[len(_A       ) - 1]


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





import gc
import unittest

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

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class  UpperCAmelCase__	( A__		,  A__		,  A__		,  unittest.TestCase ):






  """simple docstring"""


  a     =      AltDiffusionPipeline
  a     =      TEXT_TO_IMAGE_PARAMS
  a     =      TEXT_TO_IMAGE_BATCH_PARAMS
  a     =      TEXT_TO_IMAGE_IMAGE_PARAMS
  a     =      TEXT_TO_IMAGE_IMAGE_PARAMS






  def    lowercase_	(							self		:	Tuple      )	->							Tuple:
       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					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''')			,     cross_attention_dim=32			,     )
       SCREAMING_SNAKE_CASE__					=					DDIMScheduler(
           beta_start=0.00085			,     beta_end=0.012			,     beta_schedule='''scaled_linear'''			,     clip_sample=__lowerCamelCase			,     set_alpha_to_one=__lowerCamelCase			,     )
       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					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			,     )

       # TODO: address the non-deterministic text encoder (fails for save-load tests)
       # torch.manual_seed(0)
       # text_encoder_config = RobertaSeriesConfig(
       #     hidden_size=32,
       #     project_dim=32,
       #     intermediate_size=37,
       #     layer_norm_eps=1e-05,
       #     num_attention_heads=4,
       #     num_hidden_layers=5,
       #     vocab_size=5002,
       # )
       # text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					CLIPTextConfig(
           bos_token_id=0			,     eos_token_id=2			,     hidden_size=32			,     projection_dim=32			,     intermediate_size=37			,     layer_norm_eps=1e-05			,     num_attention_heads=4			,     num_hidden_layers=5			,     pad_token_id=1			,     vocab_size=5002			,     )
       SCREAMING_SNAKE_CASE__					=					CLIPTextModel(__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''      )
       SCREAMING_SNAKE_CASE__					=					77

       SCREAMING_SNAKE_CASE__					=					{
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''safety_checker''': None,
           '''feature_extractor''': None,
       }
       return components






  def    lowercase_	(							self		:	Any			,     __lowerCamelCase		:	Tuple			,     __lowerCamelCase		:	Optional[int]=0      )	->							Union[str, Any]:
       if str(__lowerCamelCase      ).startswith('''mps'''      ):
            SCREAMING_SNAKE_CASE__					=					torch.manual_seed(__lowerCamelCase      )
       else:
            SCREAMING_SNAKE_CASE__					=					torch.Generator(device=__lowerCamelCase      ).manual_seed(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					{
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''guidance_scale''': 6.0,
           '''output_type''': '''numpy''',
       }
       return inputs






  def    lowercase_	(							self		:	List[Any]      )	->							str:
       super().test_attention_slicing_forward_pass(expected_max_diff=3e-3      )






  def    lowercase_	(							self		:	List[Any]      )	->							int:
       super().test_inference_batch_single_identical(expected_max_diff=3e-3      )






  def    lowercase_	(							self		:	Dict      )	->							Optional[int]:
       SCREAMING_SNAKE_CASE__					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator

       SCREAMING_SNAKE_CASE__					=					self.get_dummy_components()
       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					RobertaSeriesConfig(
           hidden_size=32			,     project_dim=32			,     intermediate_size=37			,     layer_norm_eps=1e-05			,     num_attention_heads=4			,     num_hidden_layers=5			,     vocab_size=5002			,     )
       # TODO: remove after fixing the non-deterministic text encoder
       SCREAMING_SNAKE_CASE__					=					RobertaSeriesModelWithTransformation(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					text_encoder

       SCREAMING_SNAKE_CASE__					=					AltDiffusionPipeline(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe.to(__lowerCamelCase      )
       alt_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					self.get_dummy_inputs(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					'''A photo of an astronaut'''
       SCREAMING_SNAKE_CASE__					=					alt_pipe(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					output.images
       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)
       SCREAMING_SNAKE_CASE__					=					np.array(
           [0.5748162, 0.60447145, 0.48821217, 0.50100636, 0.5431185, 0.45763683, 0.49657696, 0.48132733, 0.47573093]      )

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







  def    lowercase_	(							self		:	List[Any]      )	->							Optional[Any]:
       SCREAMING_SNAKE_CASE__					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator

       SCREAMING_SNAKE_CASE__					=					self.get_dummy_components()
       SCREAMING_SNAKE_CASE__					=					PNDMScheduler(skip_prk_steps=__lowerCamelCase      )
       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					RobertaSeriesConfig(
           hidden_size=32			,     project_dim=32			,     intermediate_size=37			,     layer_norm_eps=1e-05			,     num_attention_heads=4			,     num_hidden_layers=5			,     vocab_size=5002			,     )
       # TODO: remove after fixing the non-deterministic text encoder
       SCREAMING_SNAKE_CASE__					=					RobertaSeriesModelWithTransformation(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					text_encoder
       SCREAMING_SNAKE_CASE__					=					AltDiffusionPipeline(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe.to(__lowerCamelCase      )
       alt_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					self.get_dummy_inputs(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					output.images
       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)
       SCREAMING_SNAKE_CASE__					=					np.array(
           [0.51605093, 0.5707241, 0.47365507, 0.50578886, 0.5633877, 0.4642503, 0.5182081, 0.48763484, 0.49084237]      )

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






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






  """simple docstring"""


  def    lowercase_	(							self		:	Tuple      )	->							str:
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






  def    lowercase_	(							self		:	str      )	->							Any:
       # make sure here that pndm scheduler skips prk
       SCREAMING_SNAKE_CASE__					=					AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''			,     safety_checker=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe.to(__lowerCamelCase      )
       alt_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					'''A painting of a squirrel eating a burger'''
       SCREAMING_SNAKE_CASE__					=					torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe([prompt]			,     generator=__lowerCamelCase			,     guidance_scale=6.0			,     num_inference_steps=20			,     output_type='''np'''      )

       SCREAMING_SNAKE_CASE__					=					output.images

       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       SCREAMING_SNAKE_CASE__					=					np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586]      )

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







  def    lowercase_	(							self		:	Optional[int]      )	->							List[str]:
       SCREAMING_SNAKE_CASE__					=					DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''			,     subfolder='''scheduler'''      )

       SCREAMING_SNAKE_CASE__					=					AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''			,     scheduler=__lowerCamelCase			,     safety_checker=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe.to(__lowerCamelCase      )
       alt_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					'''A painting of a squirrel eating a burger'''
       SCREAMING_SNAKE_CASE__					=					torch.manual_seed(0      )

       SCREAMING_SNAKE_CASE__					=					alt_pipe([prompt]			,     generator=__lowerCamelCase			,     num_inference_steps=2			,     output_type='''numpy'''      )
       SCREAMING_SNAKE_CASE__					=					output.images

       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       SCREAMING_SNAKE_CASE__					=					np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323]      )

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


