






"""simple docstring"""


import types
from typing import List, Optional, Tuple, Union

import torch
from transformers import CLIPTextModelWithProjection, CLIPTokenizer
from transformers.models.clip.modeling_clip import CLIPTextModelOutput

from diffusers.models import PriorTransformer
from diffusers.pipelines import DiffusionPipeline, StableDiffusionImageVariationPipeline
from diffusers.schedulers import UnCLIPScheduler
from diffusers.utils import logging, randn_tensor


_UpperCamelCase							:       Optional[Any]       					=		logging.get_logger(__name__)  # pylint: disable=invalid-name




def      snake_case    (self :Optional[Any]      ,     A_ :Any      ,     A_ :List[str]      ,     A_ :Optional[int]      ,     A_ :Optional[Any] ):







  '''simple docstring'''

  a	:     Tuple							=  image.to(device=A_ )
  a	:     int							=  image  # take image as image_embeddings
  a	:     Optional[Any]							=  image_embeddings.unsqueeze(1 )

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

  if do_classifier_free_guidance:
    a	:     Optional[int]							=  torch.zeros_like(A_ )

    # 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	:     Dict							=  torch.cat([uncond_embeddings, image_embeddings] )

  return image_embeddings






class 		snake_case    (    UpperCAmelCase			):





    def __init__(				self      :   Any							,						A      :   PriorTransformer							,						A      :   CLIPTokenizer							,						A      :   CLIPTextModelWithProjection							,						A      :   UnCLIPScheduler							,						A      :   Optional[dict] = None							,						):


      '''simple docstring'''




      super().__init__()

      a	:     str							=  {'image_encoder': None} if decoder_pipe_kwargs is None else decoder_pipe_kwargs

      a	:     List[Any]							=  decoder_pipe_kwargs.get('torch_dtype'							,						A     ) or prior.dtype

      a	:     Dict							=  StableDiffusionImageVariationPipeline.from_pretrained(
          'lambdalabs/sd-image-variations-diffusers'							,						**A     )

      # replace `_encode_image` method
      a	:     str							=  types.MethodType(_encode_image							,						self.decoder_pipe     )

      self.register_modules(
          prior=A							,						tokenizer=A							,						text_encoder=A							,						prior_scheduler=A							,						)





    def 							lowerCamelCase__ (				self      :   Dict							,						A      :   int							,						A      :   Dict							,						A      :   str							,						A      :   List[Any]							,						A      :   Optional[Union[CLIPTextModelOutput, Tuple]] = None							,						A      :   Optional[torch.Tensor] = None							,						):


      '''simple docstring'''




      if text_model_output is None:
        a	:     Any							=  len(A     ) if isinstance(A							,						A     ) else 1
        # get prompt text embeddings
        a	:     int							=  self.tokenizer(
            A							,						padding='max_length'							,						max_length=self.tokenizer.model_max_length							,						return_tensors='pt'							,						)
        a	:     Optional[Any]							=  text_inputs.input_ids
        a	:     Optional[Any]							=  text_inputs.attention_mask.bool().to(A     )

        if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
          a	:     Optional[Any]							=  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	:     Any							=  text_input_ids[:, : self.tokenizer.model_max_length]

        a	:     Dict							=  self.text_encoder(text_input_ids.to(A     )     )

        a	:     Optional[int]							=  text_encoder_output.text_embeds
        a	:     List[str]							=  text_encoder_output.last_hidden_state

      else:
        a	:     int							=  text_model_output[0].shape[0]
        a,	a	:     List[Any]							=  text_model_output[0], text_model_output[1]
        a	:     str							=  text_attention_mask

      a	:     List[Any]							=  text_embeddings.repeat_interleave(A							,						dim=0     )
      a	:     Dict							=  text_encoder_hidden_states.repeat_interleave(A							,						dim=0     )
      a	:     Any							=  text_mask.repeat_interleave(A							,						dim=0     )

      if do_classifier_free_guidance:
        a	:     Union[str, Any]							=  [''] * batch_size

        a	:     str							=  self.tokenizer(
            A							,						padding='max_length'							,						max_length=self.tokenizer.model_max_length							,						truncation=A							,						return_tensors='pt'							,						)
        a	:     Dict							=  uncond_input.attention_mask.bool().to(A     )
        a	:     Dict							=  self.text_encoder(uncond_input.input_ids.to(A     )     )

        a	:     Union[str, Any]							=  uncond_embeddings_text_encoder_output.text_embeds
        a	:     str							=  uncond_embeddings_text_encoder_output.last_hidden_state

        # duplicate unconditional embeddings for each generation per prompt, using mps friendly method

        a	:     int							=  uncond_embeddings.shape[1]
        a	:     int							=  uncond_embeddings.repeat(1							,						A     )
        a	:     int							=  uncond_embeddings.view(batch_size * num_images_per_prompt							,						A     )

        a	:     Dict							=  uncond_text_encoder_hidden_states.shape[1]
        a	:     Tuple							=  uncond_text_encoder_hidden_states.repeat(1							,						A							,						1     )
        a	:     Dict							=  uncond_text_encoder_hidden_states.view(
            batch_size * num_images_per_prompt							,						A							,						-1     )
        a	:     List[str]							=  uncond_text_mask.repeat_interleave(A							,						dim=0     )

        # done duplicates

        # 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]     )
        a	:     Tuple							=  torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]     )

        a	:     List[str]							=  torch.cat([uncond_text_mask, text_mask]     )

      return text_embeddings, text_encoder_hidden_states, text_mask





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


      '''simple docstring'''




      if self.device != torch.device('meta'     ) or not hasattr(self.prior							,						'_hf_hook'     ):
        return self.device
      for module in self.prior.modules():
        if (
            hasattr(A							,						'_hf_hook'     )
            and hasattr(module._hf_hook							,						'execution_device'     )
            and module._hf_hook.execution_device is not None
        ):
          return torch.device(module._hf_hook.execution_device     )
      return self.device





    def 							lowerCamelCase__ (				self      :   List[str]							,						A      :   Union[str, Any]							,						A      :   Optional[Any]							,						A      :   List[Any]							,						A      :   Optional[int]							,						A      :   Tuple							,						A      :   Tuple     ):


      '''simple docstring'''




      if latents is None:
        a	:     str							=  randn_tensor(A							,						generator=A							,						device=A							,						dtype=A     )
      else:
        if latents.shape != shape:
          raise ValueError(F'''Unexpected latents shape, got {latents.shape}, expected {shape}'''     )
        a	:     int							=  latents.to(A     )

      a	:     Optional[int]							=  latents * scheduler.init_noise_sigma
      return latents





    def 							lowerCamelCase__ (				self      :   int							,						A      :   Optional[Union[str, torch.device]] = None     ):


      '''simple docstring'''




      self.decoder_pipe.to(A     )
      super().to(A     )





    @torch.no_grad()
    def __call__(				self      :   List[Any]							,						A      :   Optional[Union[str, List[str]]] = None							,						A      :   Optional[int] = None							,						A      :   Optional[int] = None							,						A      :   int = 1							,						A      :   int = 2_5							,						A      :   Optional[torch.Generator] = None							,						A      :   Optional[torch.FloatTensor] = None							,						A      :   Optional[Union[CLIPTextModelOutput, Tuple]] = None							,						A      :   Optional[torch.Tensor] = None							,						A      :   float = 4.0							,						A      :   float = 8.0							,						A      :   int = 5_0							,						A      :   Optional[int] = 1							,						A      :   float = 0.0							,						A      :   Optional[str] = "pil"							,						A      :   bool = True							,						):


      '''simple docstring'''




      if prompt is not None:
        if isinstance(A							,						A     ):
          a	:     Optional[Any]							=  1
        elif isinstance(A							,						A     ):
          a	:     Tuple							=  len(A     )
        else:
          raise ValueError(F'''`prompt` has to be of type `str` or `list` but is {type(A     )}'''     )
      else:
        a	:     str							=  text_model_output[0].shape[0]

      a	:     Optional[int]							=  self._execution_device

      a	:     Optional[Any]							=  batch_size * num_images_per_prompt

      a	:     Any							=  prior_guidance_scale > 1.0 or decoder_guidance_scale > 1.0

      a,	a,	a	:     Tuple							=  self._encode_prompt(
          A							,						A							,						A							,						A							,						A							,						A     )

      # prior

      self.prior_scheduler.set_timesteps(A							,						device=A     )
      a	:     str							=  self.prior_scheduler.timesteps

      a	:     str							=  self.prior.config.embedding_dim

      a	:     Optional[int]							=  self.prepare_latents(
          (batch_size, embedding_dim)							,						text_embeddings.dtype							,						A							,						A							,						A							,						self.prior_scheduler							,						)

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

        a	:     Any							=  self.prior(
            A							,						timestep=A							,						proj_embedding=A							,						encoder_hidden_states=A							,						attention_mask=A							,						).predicted_image_embedding

        if do_classifier_free_guidance:
          a,	a	:     Any							=  predicted_image_embedding.chunk(2     )
          a	:     str							=  predicted_image_embedding_uncond + prior_guidance_scale * (
              predicted_image_embedding_text - predicted_image_embedding_uncond
          )

        if i + 1 == prior_timesteps_tensor.shape[0]:
          a	:     List[str]							=  None
        else:
          a	:     Tuple							=  prior_timesteps_tensor[i + 1]

        a	:     int							=  self.prior_scheduler.step(
            A							,						timestep=A							,						sample=A							,						generator=A							,						prev_timestep=A							,						).prev_sample

      a	:     List[str]							=  self.prior.post_process_latents(A     )

      a	:     int							=  prior_latents

      a	:     List[str]							=  self.decoder_pipe(
          image=A							,						height=A							,						width=A							,						num_inference_steps=A							,						guidance_scale=A							,						generator=A							,						output_type=A							,						return_dict=A							,						num_images_per_prompt=A							,						eta=A							,						)
      return output












"""simple docstring"""


_UpperCamelCase							:       str       					=		10**8




def      snake_case    (A_ :int ):







  '''simple docstring'''

  if n % 1_0 == 0:
    return False
  a	:     Optional[int]							=  str(A_ )
  return s == s[::-1]




def      snake_case    ():







  '''simple docstring'''

  a	:     Tuple							=  set()
  a	:     Optional[Any]							=  1
  a	:     List[str]							=  5
  while sum_squares < LIMIT:
    a	:     int							=  first_square + 1
    while sum_squares < LIMIT:
      if is_palindrome(A_ ):
        answer.add(A_ )
      last_square += 1
      sum_squares += last_square**2
    first_square += 1
    a	:     List[str]							=  first_square**2 + (first_square + 1) ** 2

  return sum(A_ )


if __name__ == "__main__":
   print(solution())




