



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


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



class       UpperCAmelCase__(							lowerCamelCase	):

   '''simple docstring'''




   def __init__(						self  :			Any			,				lowerCAmelCase  :			AutoencoderKL			,				lowerCAmelCase  :			CLIPTextModel			,				lowerCAmelCase  :			CLIPTokenizer			,				lowerCAmelCase  :			UNetaDConditionModel			,				lowerCAmelCase  :			Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]			,				lowerCAmelCase  :			StableDiffusionSafetyChecker			,				lowerCAmelCase  :			CLIPImageProcessor			,				)	->				Optional[Any]:




         """simple docstring"""

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


   def 			UpperCAmelCase		(						self  :			List[str]			,				lowerCAmelCase  :			Optional[Union[str, int]] = "auto")	->				List[Any]:




         """simple docstring"""

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


   def 			UpperCAmelCase		(						self  :			Optional[Any])	->				Dict:




         """simple docstring"""

         self.enable_attention_slicing(lowerCAmelCase)



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




         """simple docstring"""

         if isinstance(lowerCAmelCase			,				lowerCAmelCase):
               lowercase__     =  1
         elif isinstance(lowerCAmelCase			,				lowerCAmelCase):
               lowercase__     =  len(lowerCAmelCase)
         else:
               raise ValueError(f'''`prompt` has to be of type `str` or `list` but is {type(lowerCAmelCase)}''')

         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(lowerCAmelCase			,				lowerCAmelCase) or callback_steps <= 0)
         ):
               raise ValueError(
                   f'''`callback_steps` has to be a positive integer but is {callback_steps} of type'''
                   f''' {type(lowerCAmelCase)}.''')

         # get prompt text embeddings
         lowercase__     =  self.tokenizer(
             lowerCAmelCase			,				padding='max_length'			,				max_length=self.tokenizer.model_max_length			,				return_tensors='pt'			,				)
         lowercase__     =  text_inputs.input_ids

         if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
               lowercase__     =  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}''')
               lowercase__     =  text_input_ids[:, : self.tokenizer.model_max_length]

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

         # duplicate text embeddings for each generation per prompt, using mps friendly method
         lowercase__,	lowercase__,	lowercase__     =  text_embeddings.shape
         lowercase__     =  text_embeddings.repeat(1			,				lowerCAmelCase			,				1)
         lowercase__     =  text_embeddings.view(bs_embed * num_images_per_prompt			,				lowerCAmelCase			,				-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.
         lowercase__     =  guidance_scale > 1.0
         # get unconditional embeddings for classifier free guidance
         if do_classifier_free_guidance:
               lowercase__     =  42
               if negative_prompt is None:
                     lowercase__     =  ['']
               elif type(lowerCAmelCase) is not type(lowerCAmelCase):
                     raise TypeError(
                         f'''`negative_prompt` should be the same type to `prompt`, but got {type(lowerCAmelCase)} !='''
                         f''' {type(lowerCAmelCase)}.''')
               elif isinstance(lowerCAmelCase			,				lowerCAmelCase):
                     lowercase__     =  [negative_prompt]
               elif batch_size != len(lowerCAmelCase):
                     raise ValueError(
                         f'''`negative_prompt`: {negative_prompt} has batch size {len(lowerCAmelCase)}, but `prompt`:'''
                         f''' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'''
                         ' the batch size of `prompt`.')
               else:
                     lowercase__     =  negative_prompt

               lowercase__     =  text_input_ids.shape[-1]
               lowercase__     =  self.tokenizer(
                   lowerCAmelCase			,				padding='max_length'			,				max_length=lowerCAmelCase			,				truncation=lowerCAmelCase			,				return_tensors='pt'			,				)
               lowercase__     =  self.text_encoder(uncond_input.input_ids.to(self.device))[0]

               # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
               lowercase__     =  uncond_embeddings.shape[1]
               lowercase__     =  uncond_embeddings.repeat(lowerCAmelCase			,				lowerCAmelCase			,				1)
               lowercase__     =  uncond_embeddings.view(batch_size * num_images_per_prompt			,				lowerCAmelCase			,				-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
               lowercase__     =  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__     =  (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
         lowercase__     =  (batch_size * num_images_per_prompt, self.unet.config.in_channels, 64, 64)
         lowercase__     =  text_embeddings.dtype
         if latents is None:
               if self.device.type == "mps":
                     # randn does not exist on mps
                     lowercase__     =  torch.randn(
                         lowerCAmelCase			,				generator=lowerCAmelCase			,				device='cpu'			,				dtype=lowerCAmelCase).to(self.device)
                     lowercase__     =  torch.randn(lowerCAmelCase			,				generator=lowerCAmelCase			,				device='cpu'			,				dtype=lowerCAmelCase).to(
                         self.device)
               else:
                     lowercase__     =  torch.randn(
                         lowerCAmelCase			,				generator=lowerCAmelCase			,				device=self.device			,				dtype=lowerCAmelCase)
                     lowercase__     =  torch.randn(lowerCAmelCase			,				generator=lowerCAmelCase			,				device=self.device			,				dtype=lowerCAmelCase)
         else:
               if latents_reference.shape != latents_shape:
                     raise ValueError(f'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}''')
               lowercase__     =  latents_reference.to(self.device)
               lowercase__     =  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
         lowercase__     =  (latents_shape[3] - latents_shape_reference[3]) // 2
         lowercase__     =  (latents_shape[2] - latents_shape_reference[2]) // 2
         lowercase__     =  latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
         lowercase__     =  latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
         lowercase__     =  0 if dx < 0 else dx
         lowercase__     =  0 if dy < 0 else dy
         lowercase__     =  max(-dx			,				0)
         lowercase__     =  max(-dy			,				0)
         # import pdb
         # pdb.set_trace()
         lowercase__     =  latents_reference[:, :, dy : dy + h, dx : dx + w]

         # set timesteps
         self.scheduler.set_timesteps(lowerCAmelCase)

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

         # scale the initial noise by the standard deviation required by the scheduler
         lowercase__     =  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__     =  'eta' in set(inspect.signature(self.scheduler.step).parameters.keys())
         lowercase__     =  {}
         if accepts_eta:
               lowercase__     =  eta

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

               # predict the noise residual
               lowercase__     =  self.unet(lowerCAmelCase			,				lowerCAmelCase			,				encoder_hidden_states=lowerCAmelCase).sample

               # perform guidance
               if do_classifier_free_guidance:
                     lowercase__,	lowercase__     =  noise_pred.chunk(2)
                     lowercase__     =  noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

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

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

         lowercase__     =  1 / 0.1_82_15 * latents
         lowercase__     =  self.vae.decode(lowerCAmelCase).sample

         lowercase__     =  (image / 2 + 0.5).clamp(0			,				1)

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

         if self.safety_checker is not None:
               lowercase__     =  self.feature_extractor(self.numpy_to_pil(lowerCAmelCase)			,				return_tensors='pt').to(
                   self.device)
               lowercase__,	lowercase__     =  self.safety_checker(
                   images=lowerCAmelCase			,				clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype))
         else:
               lowercase__     =  None

         if output_type == "pil":
               lowercase__     =  self.numpy_to_pil(lowerCAmelCase)

         if not return_dict:
               return (image, has_nsfw_concept)

         return StableDiffusionPipelineOutput(images=lowerCAmelCase			,				nsfw_content_detected=lowerCAmelCase)





# Imports
import numpy as np



class       UpperCAmelCase__:

   '''simple docstring'''




   def __init__(						self  :			Any			,				lowerCAmelCase  :			Dict=None			,				lowerCAmelCase  :			List[Any]=None			,				lowerCAmelCase  :			List[Any]=None			,				lowerCAmelCase  :			List[str]=None			,				lowerCAmelCase  :			List[str]=None)	->				Dict:




         """simple docstring"""

         self.set_matricies(red=lowerCAmelCase			,				green=lowerCAmelCase			,				blue=lowerCAmelCase			,				red_edge=lowerCAmelCase			,				nir=lowerCAmelCase)


   def 			UpperCAmelCase		(						self  :			Dict			,				lowerCAmelCase  :			Dict=None			,				lowerCAmelCase  :			Union[str, Any]=None			,				lowerCAmelCase  :			Tuple=None			,				lowerCAmelCase  :			str=None			,				lowerCAmelCase  :			str=None)	->				int:




         """simple docstring"""

         if red is not None:
               lowercase__     =  red
         if green is not None:
               lowercase__     =  green
         if blue is not None:
               lowercase__     =  blue
         if red_edge is not None:
               lowercase__     =  red_edge
         if nir is not None:
               lowercase__     =  nir
         return True


   def 			UpperCAmelCase		(						self  :			Optional[int]			,				lowerCAmelCase  :			Union[str, Any]=""			,				lowerCAmelCase  :			Tuple=None			,				lowerCAmelCase  :			Optional[Any]=None			,				lowerCAmelCase  :			Optional[Any]=None			,				lowerCAmelCase  :			List[Any]=None			,				lowerCAmelCase  :			Dict=None)	->				Union[str, Any]:




         """simple docstring"""

         self.set_matricies(red=lowerCAmelCase			,				green=lowerCAmelCase			,				blue=lowerCAmelCase			,				red_edge=lowerCAmelCase			,				nir=lowerCAmelCase)
         lowercase__     =  {
             'ARVI2': self.arvaa,
             'CCCI': self.ccci,
             'CVI': self.cvi,
             'GLI': self.gli,
             'NDVI': self.ndvi,
             'BNDVI': self.bndvi,
             'redEdgeNDVI': self.red_edge_ndvi,
             'GNDVI': self.gndvi,
             'GBNDVI': self.gbndvi,
             'GRNDVI': self.grndvi,
             'RBNDVI': self.rbndvi,
             'PNDVI': self.pndvi,
             'ATSAVI': self.atsavi,
             'BWDRVI': self.bwdrvi,
             'CIgreen': self.ci_green,
             'CIrededge': self.ci_rededge,
             'CI': self.ci,
             'CTVI': self.ctvi,
             'GDVI': self.gdvi,
             'EVI': self.evi,
             'GEMI': self.gemi,
             'GOSAVI': self.gosavi,
             'GSAVI': self.gsavi,
             'Hue': self.hue,
             'IVI': self.ivi,
             'IPVI': self.ipvi,
             'I': self.i,
             'RVI': self.rvi,
             'MRVI': self.mrvi,
             'MSAVI': self.m_savi,
             'NormG': self.norm_g,
             'NormNIR': self.norm_nir,
             'NormR': self.norm_r,
             'NGRDI': self.ngrdi,
             'RI': self.ri,
             'S': self.s,
             'IF': self._if,
             'DVI': self.dvi,
             'TVI': self.tvi,
             'NDRE': self.ndre,
         }

         try:
               return funcs[index]()
         except KeyError:
               print('Index not in the list!')
               return False


   def 			UpperCAmelCase		(						self  :			Optional[int])	->				List[str]:




         """simple docstring"""

         return -0.18 + (1.17 * ((self.nir - self.red) / (self.nir + self.red)))


   def 			UpperCAmelCase		(						self  :			int)	->				Any:




         """simple docstring"""

         return ((self.nir - self.redEdge) / (self.nir + self.redEdge)) / (
             (self.nir - self.red) / (self.nir + self.red)
         )


   def 			UpperCAmelCase		(						self  :			str)	->				Optional[int]:




         """simple docstring"""

         return self.nir * (self.red / (self.green**2))


   def 			UpperCAmelCase		(						self  :			List[str])	->				Optional[int]:




         """simple docstring"""

         return (2 * self.green - self.red - self.blue) / (
             2 * self.green + self.red + self.blue
         )


   def 			UpperCAmelCase		(						self  :			Tuple)	->				Any:




         """simple docstring"""

         return (self.nir - self.red) / (self.nir + self.red)


   def 			UpperCAmelCase		(						self  :			Union[str, Any])	->				str:




         """simple docstring"""

         return (self.nir - self.blue) / (self.nir + self.blue)


   def 			UpperCAmelCase		(						self  :			List[Any])	->				Optional[int]:




         """simple docstring"""

         return (self.redEdge - self.red) / (self.redEdge + self.red)


   def 			UpperCAmelCase		(						self  :			Optional[Any])	->				Optional[int]:




         """simple docstring"""

         return (self.nir - self.green) / (self.nir + self.green)


   def 			UpperCAmelCase		(						self  :			Dict)	->				int:




         """simple docstring"""

         return (self.nir - (self.green + self.blue)) / (
             self.nir + (self.green + self.blue)
         )


   def 			UpperCAmelCase		(						self  :			Union[str, Any])	->				Optional[int]:




         """simple docstring"""

         return (self.nir - (self.green + self.red)) / (
             self.nir + (self.green + self.red)
         )


   def 			UpperCAmelCase		(						self  :			Optional[Any])	->				Dict:




         """simple docstring"""

         return (self.nir - (self.blue + self.red)) / (self.nir + (self.blue + self.red))


   def 			UpperCAmelCase		(						self  :			Any)	->				Union[str, Any]:




         """simple docstring"""

         return (self.nir - (self.green + self.red + self.blue)) / (
             self.nir + (self.green + self.red + self.blue)
         )


   def 			UpperCAmelCase		(						self  :			Optional[int]			,				lowerCAmelCase  :			List[Any]=0.08			,				lowerCAmelCase  :			Optional[int]=1.22			,				lowerCAmelCase  :			int=0.03)	->				List[Any]:




         """simple docstring"""

         return a * (
             (self.nir - a * self.red - b)
             / (a * self.nir + self.red - a * b + x * (1 + a**2))
         )


   def 			UpperCAmelCase		(						self  :			Tuple)	->				Any:




         """simple docstring"""

         return (0.1 * self.nir - self.blue) / (0.1 * self.nir + self.blue)


   def 			UpperCAmelCase		(						self  :			int)	->				Tuple:




         """simple docstring"""

         return (self.nir / self.green) - 1


   def 			UpperCAmelCase		(						self  :			Any)	->				str:




         """simple docstring"""

         return (self.nir / self.redEdge) - 1


   def 			UpperCAmelCase		(						self  :			Any)	->				List[str]:




         """simple docstring"""

         return (self.red - self.blue) / self.red


   def 			UpperCAmelCase		(						self  :			Any)	->				Optional[int]:




         """simple docstring"""

         lowercase__     =  self.ndvi()
         return ((ndvi + 0.5) / (abs(ndvi + 0.5))) * (abs(ndvi + 0.5) ** (1 / 2))


   def 			UpperCAmelCase		(						self  :			List[Any])	->				str:




         """simple docstring"""

         return self.nir - self.green


   def 			UpperCAmelCase		(						self  :			Tuple)	->				List[Any]:




         """simple docstring"""

         return 2.5 * (
             (self.nir - self.red) / (self.nir + 6 * self.red - 7.5 * self.blue + 1)
         )


   def 			UpperCAmelCase		(						self  :			Any)	->				Union[str, Any]:




         """simple docstring"""

         lowercase__     =  (2 * (self.nir**2 - self.red**2) + 1.5 * self.nir + 0.5 * self.red) / (
             self.nir + self.red + 0.5
         )
         return n * (1 - 0.25 * n) - (self.red - 0.1_25) / (1 - self.red)


   def 			UpperCAmelCase		(						self  :			int			,				lowerCAmelCase  :			int=0.16)	->				Dict:




         """simple docstring"""

         return (self.nir - self.green) / (self.nir + self.green + y)


   def 			UpperCAmelCase		(						self  :			str			,				lowerCAmelCase  :			Optional[int]=0.5)	->				Union[str, Any]:




         """simple docstring"""

         return ((self.nir - self.green) / (self.nir + self.green + n)) * (1 + n)


   def 			UpperCAmelCase		(						self  :			str)	->				int:




         """simple docstring"""

         return np.arctan(
             ((2 * self.red - self.green - self.blue) / 30.5) * (self.green - self.blue))


   def 			UpperCAmelCase		(						self  :			Union[str, Any]			,				lowerCAmelCase  :			int=None			,				lowerCAmelCase  :			List[str]=None)	->				Tuple:




         """simple docstring"""

         return (self.nir - b) / (a * self.red)


   def 			UpperCAmelCase		(						self  :			int)	->				Dict:




         """simple docstring"""

         return (self.nir / ((self.nir + self.red) / 2)) * (self.ndvi() + 1)


   def 			UpperCAmelCase		(						self  :			Optional[int])	->				Optional[Any]:




         """simple docstring"""

         return (self.red + self.green + self.blue) / 30.5


   def 			UpperCAmelCase		(						self  :			int)	->				str:




         """simple docstring"""

         return self.nir / self.red


   def 			UpperCAmelCase		(						self  :			Optional[int])	->				Optional[Any]:




         """simple docstring"""

         return (self.rvi() - 1) / (self.rvi() + 1)


   def 			UpperCAmelCase		(						self  :			Optional[int])	->				Optional[int]:




         """simple docstring"""

         return (
             (2 * self.nir + 1)
             - ((2 * self.nir + 1) ** 2 - 8 * (self.nir - self.red)) ** (1 / 2)
         ) / 2


   def 			UpperCAmelCase		(						self  :			Tuple)	->				Any:




         """simple docstring"""

         return self.green / (self.nir + self.red + self.green)


   def 			UpperCAmelCase		(						self  :			Any)	->				Optional[Any]:




         """simple docstring"""

         return self.nir / (self.nir + self.red + self.green)


   def 			UpperCAmelCase		(						self  :			List[Any])	->				Dict:




         """simple docstring"""

         return self.red / (self.nir + self.red + self.green)


   def 			UpperCAmelCase		(						self  :			Optional[Any])	->				Any:




         """simple docstring"""

         return (self.green - self.red) / (self.green + self.red)


   def 			UpperCAmelCase		(						self  :			Dict)	->				Tuple:




         """simple docstring"""

         return (self.red - self.green) / (self.red + self.green)


   def 			UpperCAmelCase		(						self  :			str)	->				int:




         """simple docstring"""

         lowercase__     =  np.max([np.max(self.red), np.max(self.green), np.max(self.blue)])
         lowercase__     =  np.min([np.min(self.red), np.min(self.green), np.min(self.blue)])
         return (max_value - min_value) / max_value


   def 			UpperCAmelCase		(						self  :			Optional[int])	->				Tuple:




         """simple docstring"""

         return (2 * self.red - self.green - self.blue) / (self.green - self.blue)


   def 			UpperCAmelCase		(						self  :			int)	->				Optional[Any]:




         """simple docstring"""

         return self.nir / self.red


   def 			UpperCAmelCase		(						self  :			Dict)	->				Dict:




         """simple docstring"""

         return (self.ndvi() + 0.5) ** (1 / 2)



   def 			UpperCAmelCase		(						self  :			str)	->				List[Any]:




         """simple docstring"""

         return (self.nir - self.redEdge) / (self.nir + self.redEdge)
