





import inspect
from typing import Callable, List, Optional, Union

import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import DiffusionPipeline
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import logging


_lowerCAmelCase				:   str						       =						logging.get_logger(__name__)  # pylint: disable=invalid-name


class        __magic_name__	(				lowerCamelCase__				):





     """simple docstring"""



     def __init__(				self							:int      ,   snake_case							:AutoencoderKL      ,   snake_case							:CLIPTextModel      ,   snake_case							:CLIPTokenizer      ,   snake_case							:UNetaDConditionModel      ,   snake_case							:Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]      ,   snake_case							:StableDiffusionSafetyChecker      ,   snake_case							:CLIPImageProcessor      ,   ):





        '''simple docstring'''




        super().__init__()
        self.register_modules(
            vae=snake_case      ,   text_encoder=snake_case      ,   tokenizer=snake_case      ,   unet=snake_case      ,   scheduler=snake_case      ,   safety_checker=snake_case      ,   feature_extractor=snake_case      ,   )



     def     SCREAMING_SNAKE_CASE		(				self							:str      ,   snake_case							:Optional[Union[str, int]] = "auto"	):





        '''simple docstring'''




        if slice_size == "auto":
           # half the attention head size is usually a good trade-off between
           # speed and memory
           A_       :					List[str]        =      self.unet.config.attention_head_dim // 2
        self.unet.set_attention_slice(snake_case	)



     def     SCREAMING_SNAKE_CASE		(				self							:Optional[int]	):





        '''simple docstring'''




        self.enable_attention_slicing(snake_case	)





     @torch.no_grad()
     def __call__(				self							:int      ,   snake_case							:Union[str, List[str]]      ,   snake_case							:int = 512      ,   snake_case							:int = 512      ,   snake_case							:int = 50      ,   snake_case							:float = 7.5      ,   snake_case							:Optional[Union[str, List[str]]] = None      ,   snake_case							:Optional[int] = 1      ,   snake_case							:float = 0.0      ,   snake_case							:Optional[torch.Generator] = None      ,   snake_case							:Optional[torch.FloatTensor] = None      ,   snake_case							:Optional[str] = "pil"      ,   snake_case							:bool = True      ,   snake_case							:Optional[Callable[[int, int, torch.FloatTensor], None]] = None      ,   snake_case							:int = 1      ,   snake_case							:Optional[torch.FloatTensor] = None      ,   **snake_case							:Union[str, Any]      ,   ):





        '''simple docstring'''




        if isinstance(snake_case      ,   snake_case	):
           A_       :					Union[str, Any]        =      1
        elif isinstance(snake_case      ,   snake_case	):
           A_       :					Optional[Any]        =      len(snake_case	)
        else:
           raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(snake_case	)}"	)

        if height % 8 != 0 or width % 8 != 0:
           raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}."	)

        if (callback_steps is None) or (
            callback_steps is not None and (not isinstance(snake_case      ,   snake_case	) or callback_steps <= 0)
        ):
           raise ValueError(
               f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
               f" {type(snake_case	)}."	)

        # get prompt text embeddings
        A_       :					Any        =      self.tokenizer(
            snake_case      ,   padding="max_length"      ,   max_length=self.tokenizer.model_max_length      ,   return_tensors="pt"      ,   )
        A_       :					Optional[Any]        =      text_inputs.input_ids

        if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
           A_       :					str        =      self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]	)
           logger.warning(
               "The following part of your input was truncated because CLIP can only handle sequences up to"
               f" {self.tokenizer.model_max_length} tokens: {removed_text}"	)
           A_       :					str        =      text_input_ids[:, : self.tokenizer.model_max_length]

        if text_embeddings is None:
           A_       :					int        =      self.text_encoder(text_input_ids.to(self.device	)	)[0]

        # duplicate text embeddings for each generation per prompt, using mps friendly method
        A_   ,    A_   ,    A_       :					Optional[Any]        =      text_embeddings.shape
        A_       :					Any        =      text_embeddings.repeat(1      ,   snake_case      ,   1	)
        A_       :					int        =      text_embeddings.view(bs_embed * num_images_per_prompt      ,   snake_case      ,   -1	)

        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
        # corresponds to doing no classifier free guidance.
        A_       :					int        =      guidance_scale > 1.0
        # get unconditional embeddings for classifier free guidance
        if do_classifier_free_guidance:
           A_       :					List[str]
           if negative_prompt is None:
              A_       :					Any        =      [""]
           elif type(snake_case	) is not type(snake_case	):
              raise TypeError(
                  f"`negative_prompt` should be the same type to `prompt`, but got {type(snake_case	)} !="
                  f" {type(snake_case	)}."	)
           elif isinstance(snake_case      ,   snake_case	):
              A_       :					Union[str, Any]        =      [negative_prompt]
           elif batch_size != len(snake_case	):
              raise ValueError(
                  f"`negative_prompt`: {negative_prompt} has batch size {len(snake_case	)}, but `prompt`:"
                  f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
                  " the batch size of `prompt`."	)
           else:
              A_       :					Tuple        =      negative_prompt

           A_       :					List[Any]        =      text_input_ids.shape[-1]
           A_       :					Union[str, Any]        =      self.tokenizer(
               snake_case      ,   padding="max_length"      ,   max_length=snake_case      ,   truncation=snake_case      ,   return_tensors="pt"      ,   )
           A_       :					Optional[int]        =      self.text_encoder(uncond_input.input_ids.to(self.device	)	)[0]

           # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
           A_       :					List[Any]        =      uncond_embeddings.shape[1]
           A_       :					Union[str, Any]        =      uncond_embeddings.repeat(snake_case      ,   snake_case      ,   1	)
           A_       :					List[str]        =      uncond_embeddings.view(batch_size * num_images_per_prompt      ,   snake_case      ,   -1	)

           # For classifier free guidance, we need to do two forward passes.
           # Here we concatenate the unconditional and text embeddings into a single batch
           # to avoid doing two forward passes
           A_       :					str        =      torch.cat([uncond_embeddings, text_embeddings]	)

        # get the initial random noise unless the user supplied it

        # Unlike in other pipelines, latents need to be generated in the target device
        # for 1-to-1 results reproducibility with the CompVis implementation.
        # However this currently doesn't work in `mps`.
        A_       :					List[Any]        =      (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
        A_       :					Optional[int]        =      (batch_size * num_images_per_prompt, self.unet.config.in_channels, 64, 64)
        A_       :					Tuple        =      text_embeddings.dtype
        if latents is None:
           if self.device.type == "mps":
              # randn does not exist on mps
              A_       :					Optional[int]        =      torch.randn(
                  snake_case      ,   generator=snake_case      ,   device="cpu"      ,   dtype=snake_case	).to(self.device	)
              A_       :					Dict        =      torch.randn(snake_case      ,   generator=snake_case      ,   device="cpu"      ,   dtype=snake_case	).to(
                  self.device	)
           else:
              A_       :					int        =      torch.randn(
                  snake_case      ,   generator=snake_case      ,   device=self.device      ,   dtype=snake_case	)
              A_       :					Any        =      torch.randn(snake_case      ,   generator=snake_case      ,   device=self.device      ,   dtype=snake_case	)
        else:
           if latents_reference.shape != latents_shape:
              raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}"	)
           A_       :					str        =      latents_reference.to(self.device	)
           A_       :					int        =      latents.to(self.device	)

        # This is the key part of the pipeline where we
        # try to ensure that the generated images w/ the same seed
        # but different sizes actually result in similar images
        A_       :					str        =      (latents_shape[3] - latents_shape_reference[3]) // 2
        A_       :					int        =      (latents_shape[2] - latents_shape_reference[2]) // 2
        A_       :					Union[str, Any]        =      latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
        A_       :					Any        =      latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
        A_       :					int        =      0 if dx < 0 else dx
        A_       :					Dict        =      0 if dy < 0 else dy
        A_       :					str        =      max(-dx      ,   0	)
        A_       :					Dict        =      max(-dy      ,   0	)
        # import pdb
        # pdb.set_trace()
        A_       :					Optional[Any]        =      latents_reference[:, :, dy : dy + h, dx : dx + w]

        # set timesteps
        self.scheduler.set_timesteps(snake_case	)

        # Some schedulers like PNDM have timesteps as arrays
        # It's more optimized to move all timesteps to correct device beforehand
        A_       :					List[str]        =      self.scheduler.timesteps.to(self.device	)

        # scale the initial noise by the standard deviation required by the scheduler
        A_       :					List[Any]        =      latents * self.scheduler.init_noise_sigma

        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
        # and should be between [0, 1]
        A_       :					Any        =      "eta" in set(inspect.signature(self.scheduler.step	).parameters.keys()	)
        A_       :					int        =      {}
        if accepts_eta:
           A_       :					Dict        =      eta

        for i, t in enumerate(self.progress_bar(snake_case	)	):
           # expand the latents if we are doing classifier free guidance
           A_       :					int        =      torch.cat([latents] * 2	) if do_classifier_free_guidance else latents
           A_       :					List[Any]        =      self.scheduler.scale_model_input(snake_case      ,   snake_case	)

           # predict the noise residual
           A_       :					Tuple        =      self.unet(snake_case      ,   snake_case      ,   encoder_hidden_states=snake_case	).sample

           # perform guidance
           if do_classifier_free_guidance:
              A_   ,    A_       :					Optional[int]        =      noise_pred.chunk(2	)
              A_       :					List[Any]        =      noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

           # compute the previous noisy sample x_t -> x_t-1
           A_       :					str        =      self.scheduler.step(snake_case      ,   snake_case      ,   snake_case      ,   **snake_case	).prev_sample

           # call the callback, if provided
           if callback is not None and i % callback_steps == 0:
              callback(snake_case      ,   snake_case      ,   snake_case	)

        A_       :					Optional[int]        =      1 / 0.18215 * latents
        A_       :					str        =      self.vae.decode(snake_case	).sample

        A_       :					List[Any]        =      (image / 2 + 0.5).clamp(0      ,   1	)

        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
        A_       :					List[str]        =      image.cpu().permute(0      ,   2      ,   3      ,   1	).float().numpy()

        if self.safety_checker is not None:
           A_       :					Dict        =      self.feature_extractor(self.numpy_to_pil(snake_case	)      ,   return_tensors="pt"	).to(
               self.device	)
           A_   ,    A_       :					int        =      self.safety_checker(
               images=snake_case      ,   clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype	)	)
        else:
           A_       :					Optional[Any]        =      None

        if output_type == "pil":
           A_       :					Optional[int]        =      self.numpy_to_pil(snake_case	)

        if not return_dict:
           return (image, has_nsfw_concept)

        return StableDiffusionPipelineOutput(images=snake_case      ,   nsfw_content_detected=snake_case	)







from math import sqrt

def   __snake_case    (       _lowerCAmelCase			:			int						)			->    bool:

   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			:			int = 10001						)			->    int:
   A_       :					Any        =      0
   A_       :					Tuple        =      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() = }''')
