'''simple docstring'''
import inspect
from typing import List, Optional, Union

import torch
from torch import nn
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPImageProcessor, CLIPModel, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    DPMSolverMultistepScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput




class _lowerCAmelCase (  nn.Module				):







      def __init__(     self     ,						lowercase     ,						lowercase=1.0    ):
        super().__init__()

        lowercase							:       Tuple											=      cut_size
        lowercase							:       Tuple											=      cut_power







      def     lowerCamelCase__			(     self     ,						lowercase     ,						lowercase    ):
        lowercase							,			lowercase							:       Optional[int]											=      pixel_values.shape[2:4]
        lowercase							:       List[Any]											=      min(lowercase     ,						lowercase    )
        lowercase							:       Dict											=      min(lowercase     ,						lowercase     ,						self.cut_size    )
        lowercase							:       Optional[Any]											=      []
        for _ in range(lowercase    ):
          lowercase							:       Union[str, Any]											=      int(torch.rand([]    ) ** self.cut_power * (max_size - min_size) + min_size    )
          lowercase							:       str											=      torch.randint(0     ,						sideX - size + 1     ,						()    )
          lowercase							:       Dict											=      torch.randint(0     ,						sideY - size + 1     ,						()    )
          lowercase							:       int											=      pixel_values[:, :, offsety : offsety + size, offsetx : offsetx + size]
          cutouts.append(F.adaptive_avg_poolad(lowercase     ,						self.cut_size    )    )
        return torch.cat(lowercase    )







def 							__UpperCamelCase   (				__UpperCAmelCase     :				Dict   ,      __UpperCAmelCase     :				Any )       ->  Dict:
  lowercase							:       Dict											=      F.normalize(__UpperCAmelCase   ,      dim=-1 )
  lowercase							:       Optional[Any]											=      F.normalize(__UpperCAmelCase   ,      dim=-1 )
  return (x - y).norm(dim=-1 ).div(2 ).arcsin().pow(2 ).mul(2 )







def 							__UpperCamelCase   (				__UpperCAmelCase     :				Optional[int]   ,      __UpperCAmelCase     :				Optional[int] )       ->  List[Any]:
  for param in model.parameters():
    lowercase							:       Any											=      value






class _lowerCAmelCase (  snake_case				):







      def __init__(     self     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase     ,						):
        super().__init__()
        self.register_modules(
            vae=lowercase     ,						text_encoder=lowercase     ,						clip_model=lowercase     ,						tokenizer=lowercase     ,						unet=lowercase     ,						scheduler=lowercase     ,						feature_extractor=lowercase     ,						)

        lowercase							:       List[str]											=      transforms.Normalize(mean=feature_extractor.image_mean     ,						std=feature_extractor.image_std    )
        lowercase							:       Dict											=      (
            feature_extractor.size
            if isinstance(feature_extractor.size     ,						lowercase    )
            else feature_extractor.size['''shortest_edge''']
        )
        lowercase							:       Optional[Any]											=      MakeCutouts(self.cut_out_size    )

        set_requires_grad(self.text_encoder     ,						lowercase    )
        set_requires_grad(self.clip_model     ,						lowercase    )







      def     lowerCamelCase__			(     self     ,						lowercase = "auto"    ):
        if slice_size == "auto":
          # half the attention head size is usually a good trade-off between
          # speed and memory
          lowercase							:       Dict											=      self.unet.config.attention_head_dim // 2
        self.unet.set_attention_slice(lowercase    )







      def     lowerCamelCase__			(     self    ):
        self.enable_attention_slicing(lowercase    )







      def     lowerCamelCase__			(     self    ):
        set_requires_grad(self.vae     ,						lowercase    )







      def     lowerCamelCase__			(     self    ):
        set_requires_grad(self.vae     ,						lowercase    )







      def     lowerCamelCase__			(     self    ):
        set_requires_grad(self.unet     ,						lowercase    )







      def     lowerCamelCase__			(     self    ):
        set_requires_grad(self.unet     ,						lowercase    )







      @torch.enable_grad()
      def     lowerCamelCase__			(     self     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase=True     ,						):
        lowercase							:       List[Any]											=      latents.detach().requires_grad_()

        lowercase							:       int											=      self.scheduler.scale_model_input(lowercase     ,						lowercase    )

        # predict the noise residual
        lowercase							:       Optional[int]											=      self.unet(lowercase     ,						lowercase     ,						encoder_hidden_states=lowercase    ).sample

        if isinstance(self.scheduler     ,						(PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)    ):
          lowercase							:       int											=      self.scheduler.alphas_cumprod[timestep]
          lowercase							:       int											=      1 - alpha_prod_t
          # compute predicted original sample from predicted noise also called
          # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
          lowercase							:       Union[str, Any]											=      (latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

          lowercase							:       Optional[Any]											=      torch.sqrt(lowercase    )
          lowercase							:       Dict											=      pred_original_sample * (fac) + latents * (1 - fac)
        elif isinstance(self.scheduler     ,						lowercase    ):
          lowercase							:       str											=      self.scheduler.sigmas[index]
          lowercase							:       Optional[Any]											=      latents - sigma * noise_pred
        else:
          raise ValueError(F"""scheduler type {type(self.scheduler    )} not supported"""    )

        lowercase							:       Tuple											=      1 / self.vae.config.scaling_factor * sample
        lowercase							:       str											=      self.vae.decode(lowercase    ).sample
        lowercase							:       Optional[Any]											=      (image / 2 + 0.5).clamp(0     ,						1    )

        if use_cutouts:
          lowercase							:       Optional[Any]											=      self.make_cutouts(lowercase     ,						lowercase    )
        else:
          lowercase							:       Union[str, Any]											=      transforms.Resize(self.cut_out_size    )(lowercase    )
        lowercase							:       int											=      self.normalize(lowercase    ).to(latents.dtype    )

        lowercase							:       str											=      self.clip_model.get_image_features(lowercase    )
        lowercase							:       Optional[int]											=      image_embeddings_clip / image_embeddings_clip.norm(p=2     ,						dim=-1     ,						keepdim=lowercase    )

        if use_cutouts:
          lowercase							:       Union[str, Any]											=      spherical_dist_loss(lowercase     ,						lowercase    )
          lowercase							:       Any											=      dists.view([num_cutouts, sample.shape[0], -1]    )
          lowercase							:       List[Any]											=      dists.sum(2    ).mean(0    ).sum() * clip_guidance_scale
        else:
          lowercase							:       Tuple											=      spherical_dist_loss(lowercase     ,						lowercase    ).mean() * clip_guidance_scale

        lowercase							:       str											=      -torch.autograd.grad(lowercase     ,						lowercase    )[0]

        if isinstance(self.scheduler     ,						lowercase    ):
          lowercase							:       Dict											=      latents.detach() + grads * (sigma**2)
          lowercase							:       List[str]											=      noise_pred_original
        else:
          lowercase							:       Tuple											=      noise_pred_original - torch.sqrt(lowercase    ) * grads
        return noise_pred, latents







      @torch.no_grad()
      def __call__(     self     ,						lowercase     ,						lowercase = 512     ,						lowercase = 512     ,						lowercase = 50     ,						lowercase = 7.5     ,						lowercase = 1     ,						lowercase = 0.0     ,						lowercase = 100     ,						lowercase = None     ,						lowercase = 4     ,						lowercase = True     ,						lowercase = None     ,						lowercase = None     ,						lowercase = "pil"     ,						lowercase = True     ,						):
        if isinstance(lowercase     ,						lowercase    ):
          lowercase							:       Tuple											=      1
        elif isinstance(lowercase     ,						lowercase    ):
          lowercase							:       Any											=      len(lowercase    )
        else:
          raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(lowercase    )}"""    )

        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}."""    )

        # get prompt text embeddings
        lowercase							:       List[str]											=      self.tokenizer(
            lowercase     ,						padding='''max_length'''     ,						max_length=self.tokenizer.model_max_length     ,						truncation=lowercase     ,						return_tensors='''pt'''     ,						)
        lowercase							:       List[str]											=      self.text_encoder(text_input.input_ids.to(self.device    )    )[0]
        # duplicate text embeddings for each generation per prompt
        lowercase							:       Union[str, Any]											=      text_embeddings.repeat_interleave(lowercase     ,						dim=0    )

        if clip_guidance_scale > 0:
          if clip_prompt is not None:
            lowercase							:       Optional[Any]											=      self.tokenizer(
                lowercase     ,						padding='''max_length'''     ,						max_length=self.tokenizer.model_max_length     ,						truncation=lowercase     ,						return_tensors='''pt'''     ,						).input_ids.to(self.device    )
          else:
            lowercase							:       List[str]											=      text_input.input_ids.to(self.device    )
          lowercase							:       str											=      self.clip_model.get_text_features(lowercase    )
          lowercase							:       Dict											=      text_embeddings_clip / text_embeddings_clip.norm(p=2     ,						dim=-1     ,						keepdim=lowercase    )
          # duplicate text embeddings clip for each generation per prompt
          lowercase							:       int											=      text_embeddings_clip.repeat_interleave(lowercase     ,						dim=0    )

        # 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.
        lowercase							:       Optional[Any]											=      guidance_scale > 1.0
        # get unconditional embeddings for classifier free guidance
        if do_classifier_free_guidance:
          lowercase							:       List[Any]											=      text_input.input_ids.shape[-1]
          lowercase							:       str											=      self.tokenizer(['''''']     ,						padding='''max_length'''     ,						max_length=lowercase     ,						return_tensors='''pt'''    )
          lowercase							:       int											=      self.text_encoder(uncond_input.input_ids.to(self.device    )    )[0]
          # duplicate unconditional embeddings for each generation per prompt
          lowercase							:       str											=      uncond_embeddings.repeat_interleave(lowercase     ,						dim=0    )

          # 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
          lowercase							:       Any											=      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`.
        lowercase							:       Union[str, Any]											=      (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
        lowercase							:       Any											=      text_embeddings.dtype
        if latents is None:
          if self.device.type == "mps":
            # randn does not work reproducibly on mps
            lowercase							:       List[str]											=      torch.randn(lowercase     ,						generator=lowercase     ,						device='''cpu'''     ,						dtype=lowercase    ).to(
                self.device    )
          else:
            lowercase							:       Tuple											=      torch.randn(lowercase     ,						generator=lowercase     ,						device=self.device     ,						dtype=lowercase    )
        else:
          if latents.shape != latents_shape:
            raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {latents_shape}"""    )
          lowercase							:       List[str]											=      latents.to(self.device    )

        # set timesteps
        lowercase							:       Any											=      '''offset''' in set(inspect.signature(self.scheduler.set_timesteps    ).parameters.keys()    )
        lowercase							:       Dict											=      {}
        if accepts_offset:
          lowercase							:       List[Any]											=      1

        self.scheduler.set_timesteps(lowercase     ,						**lowercase    )

        # Some schedulers like PNDM have timesteps as arrays
        # It's more optimized to move all timesteps to correct device beforehand
        lowercase							:       Dict											=      self.scheduler.timesteps.to(self.device    )

        # scale the initial noise by the standard deviation required by the scheduler
        lowercase							:       Union[str, 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]
        lowercase							:       Optional[Any]											=      '''eta''' in set(inspect.signature(self.scheduler.step    ).parameters.keys()    )
        lowercase							:       str											=      {}
        if accepts_eta:
          lowercase							:       Union[str, Any]											=      eta

        # check if the scheduler accepts generator
        lowercase							:       Union[str, Any]											=      '''generator''' in set(inspect.signature(self.scheduler.step    ).parameters.keys()    )
        if accepts_generator:
          lowercase							:       str											=      generator

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

          # predict the noise residual
          lowercase							:       Dict											=      self.unet(lowercase     ,						lowercase     ,						encoder_hidden_states=lowercase    ).sample

          # perform classifier free guidance
          if do_classifier_free_guidance:
            lowercase							,			lowercase							:       Optional[Any]											=      noise_pred.chunk(2    )
            lowercase							:       Union[str, Any]											=      noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

          # perform clip guidance
          if clip_guidance_scale > 0:
            lowercase							:       Any											=      (
                text_embeddings.chunk(2    )[1] if do_classifier_free_guidance else text_embeddings
            )
            lowercase							,			lowercase							:       Union[str, Any]											=      self.cond_fn(
                lowercase     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase     ,						)

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

        # scale and decode the image latents with vae
        lowercase							:       Optional[Any]											=      1 / self.vae.config.scaling_factor * latents
        lowercase							:       Dict											=      self.vae.decode(lowercase    ).sample

        lowercase							:       int											=      (image / 2 + 0.5).clamp(0     ,						1    )
        lowercase							:       Optional[Any]											=      image.cpu().permute(0     ,						2     ,						3     ,						1    ).numpy()

        if output_type == "pil":
          lowercase							:       Union[str, Any]											=      self.numpy_to_pil(lowercase    )

        if not return_dict:
          return (image, None)

        return StableDiffusionPipelineOutput(images=lowercase     ,						nsfw_content_detected=lowercase    )
'''simple docstring'''
import os
from typing import Optional, Union

import numpy as np
from huggingface_hub import hf_hub_download

from ... import AutoTokenizer
from ...utils import logging


_A		:			List[str]	      ='''block_records.npy'''


_A		:			Optional[int]	      =logging.get_logger(__name__)







def 							__UpperCamelCase   (				__UpperCAmelCase     :				str   ,      __UpperCAmelCase     :				int )       ->  np.ndarray:
  import tensorflow.compat.va as tf

  lowercase							:       Optional[Any]											=      tf.data.TFRecordDataset(__UpperCAmelCase   ,      buffer_size=512 * 1_024 * 1_024 )
  lowercase							:       Optional[Any]											=      blocks_dataset.batch(__UpperCAmelCase   ,      drop_remainder=__UpperCAmelCase )
  lowercase							:       Tuple											=      next(blocks_dataset.take(1 ).as_numpy_iterator() )

  return np_record




class _lowerCAmelCase :







      def __init__(     self     ,						lowercase     ,						lowercase     ,						lowercase=2     ,						lowercase=1000     ,						lowercase=100     ,						lowercase=10_0000     ,						):
        from scann.scann_ops.py.scann_ops_pybind import builder as Builder

        lowercase							:       Any											=      Builder(db=lowercase     ,						num_neighbors=lowercase     ,						distance_measure='''dot_product'''    )
        lowercase							:       Optional[Any]											=      builder.tree(
            num_leaves=lowercase     ,						num_leaves_to_search=lowercase     ,						training_sample_size=lowercase    )
        lowercase							:       Optional[Any]											=      builder.score_ah(dimensions_per_block=lowercase    )

        lowercase							:       Any											=      builder.build()







      def     lowerCamelCase__			(     self     ,						lowercase    ):
        lowercase							,			lowercase							:       str											=      self.searcher.search_batched(question_projection.detach().cpu()    )
        return retrieved_block_ids.astype('''int64'''    )





class _lowerCAmelCase :







      def __init__(     self     ,						lowercase     ,						lowercase    ):
        super().__init__()
        lowercase							:       int											=      block_records
        lowercase							:       List[str]											=      tokenizer







      def __call__(     self     ,						lowercase     ,						lowercase     ,						lowercase     ,						lowercase=None     ,						lowercase="pt"    ):
        lowercase							:       List[str]											=      np.take(self.block_records     ,						indices=lowercase     ,						axis=0    )

        lowercase							:       Union[str, Any]											=      self.tokenizer.decode(question_input_ids[0]     ,						skip_special_tokens=lowercase    )

        lowercase							:       Dict											=      []
        lowercase							:       Dict											=      []
        for retrieved_block in retrieved_blocks:
          text.append(lowercase    )
          text_pair.append(retrieved_block.decode()    )

        lowercase							:       Optional[Any]											=      self.tokenizer(
            lowercase     ,						lowercase     ,						padding=lowercase     ,						truncation=lowercase     ,						return_special_tokens_mask=lowercase     ,						max_length=lowercase    )
        lowercase							:       Optional[int]											=      concat_inputs.convert_to_tensors(lowercase    )

        if answer_ids is not None:
          return self.block_has_answer(lowercase     ,						lowercase    ) + (concat_inputs_tensors,)
        else:
          return (None, None, None, concat_inputs_tensors)







      @classmethod
      def     lowerCamelCase__			(     cls     ,						lowercase     ,						*lowercase     ,						**lowercase    ):
        if os.path.isdir(lowercase    ):
          lowercase							:       Optional[Any]											=      os.path.join(lowercase     ,						_REALM_BLOCK_RECORDS_FILENAME    )
        else:
          lowercase							:       Tuple											=      hf_hub_download(
              repo_id=lowercase     ,						filename=_REALM_BLOCK_RECORDS_FILENAME     ,						**lowercase    )
        lowercase							:       int											=      np.load(lowercase     ,						allow_pickle=lowercase    )

        lowercase							:       List[Any]											=      AutoTokenizer.from_pretrained(lowercase     ,						*lowercase     ,						**lowercase    )

        return cls(lowercase     ,						lowercase    )







      def     lowerCamelCase__			(     self     ,						lowercase    ):
        # save block records
        np.save(os.path.join(lowercase     ,						_REALM_BLOCK_RECORDS_FILENAME    )     ,						self.block_records    )
        # save tokenizer
        self.tokenizer.save_pretrained(lowercase    )







      def     lowerCamelCase__			(     self     ,						lowercase     ,						lowercase    ):
        lowercase							:       Union[str, Any]											=      []
        lowercase							:       Any											=      []
        lowercase							:       List[str]											=      []
        lowercase							:       Optional[int]											=      0

        for input_id in concat_inputs.input_ids:
          lowercase							:       Union[str, Any]											=      input_id.tolist()
          # Check answers between two [SEP] tokens
          lowercase							:       Dict											=      input_id_list.index(self.tokenizer.sep_token_id    )
          lowercase							:       int											=      first_sep_idx + 1 + input_id_list[first_sep_idx + 1 :].index(self.tokenizer.sep_token_id    )

          start_pos.append([]    )
          end_pos.append([]    )
          for answer in answer_ids:
            for idx in range(first_sep_idx + 1     ,						lowercase    ):
              if answer[0] == input_id_list[idx]:
                if input_id_list[idx : idx + len(lowercase    )] == answer:
                  start_pos[-1].append(lowercase    )
                  end_pos[-1].append(idx + len(lowercase    ) - 1    )

          if len(start_pos[-1]    ) == 0:
            has_answers.append(lowercase    )
          else:
            has_answers.append(lowercase    )
            if len(start_pos[-1]    ) > max_answers:
              lowercase							:       int											=      len(start_pos[-1]    )

        # Pad -1 to max_answers
        for start_pos_, end_pos_ in zip(lowercase     ,						lowercase    ):
          if len(start_pos_    ) < max_answers:
            lowercase							:       List[str]											=      [-1] * (max_answers - len(start_pos_    ))
            start_pos_ += padded
            end_pos_ += padded
        return has_answers, start_pos, end_pos