

"""simple docstring"""


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						   =			logging.get_logger(__name__)  # pylint: disable=invalid-name





class     snake_case	(				__snake_case  ):

				def __init__(	self							:				List[str]     ,  UpperCamelCase__							:				AutoencoderKL     ,  UpperCamelCase__							:				CLIPTextModel     ,  UpperCamelCase__							:				CLIPTokenizer     ,  UpperCamelCase__							:				UNetaDConditionModel     ,  UpperCamelCase__							:				Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]     ,  UpperCamelCase__							:				StableDiffusionSafetyChecker     ,  UpperCamelCase__							:				CLIPImageProcessor     ,  )->   Optional[int]:



						'''simple docstring'''


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

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



						'''simple docstring'''


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

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



						'''simple docstring'''


						self.enable_attention_slicing(UpperCamelCase__)




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



						'''simple docstring'''


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

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

						# get prompt text embeddings
						__lowerCAmelCase:       List[Any]          =			self.tokenizer(
						    UpperCamelCase__     ,  padding="max_length"     ,  max_length=self.tokenizer.model_max_length     ,  return_tensors="pt"     ,  )
						__lowerCAmelCase:       Tuple          =			text_inputs.input_ids

						if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
								__lowerCAmelCase:       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}")
								__lowerCAmelCase:       List[str]          =			text_input_ids[:, : self.tokenizer.model_max_length]

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

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

								__lowerCAmelCase:       List[Any]          =			text_input_ids.shape[-1]
								__lowerCAmelCase:       Dict          =			self.tokenizer(
								    UpperCamelCase__     ,  padding="max_length"     ,  max_length=UpperCamelCase__     ,  truncation=UpperCamelCase__     ,  return_tensors="pt"     ,  )
								__lowerCAmelCase:       List[Any]          =			self.text_encoder(uncond_input.input_ids.to(self.device))[0]

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

						# set timesteps
						self.scheduler.set_timesteps(UpperCamelCase__)

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

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

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

								# predict the noise residual
								__lowerCAmelCase:       Union[str, Any]          =			self.unet(UpperCamelCase__     ,  UpperCamelCase__     ,  encoder_hidden_states=UpperCamelCase__).sample

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

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

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

						__lowerCAmelCase:       Union[str, Any]          =			1 / 0.18215 * latents
						__lowerCAmelCase:       Any          =			self.vae.decode(UpperCamelCase__).sample

						__lowerCAmelCase:       int          =			(image / 2 + 0.5).clamp(0     ,  1)

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

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

						if output_type == "pil":
								__lowerCAmelCase:       List[Any]          =			self.numpy_to_pil(UpperCamelCase__)

						if not return_dict:
								return (image, has_nsfw_concept)

						return StableDiffusionPipelineOutput(images=UpperCamelCase__     ,  nsfw_content_detected=UpperCamelCase__)





"""simple docstring"""


from ...configuration_utils import PretrainedConfig
from ...utils import logging


__A						   =			logging.get_logger(__name__)

__A						   =			{
    "google/switch-base-8": "https://huggingface.co/google/switch-base-8/blob/main/config.json",
}





class     snake_case	(				__snake_case  ):
				SCREAMING_SNAKE_CASE_   :      Union[str, Any]         =  """switch_transformers"""
				SCREAMING_SNAKE_CASE_   :      Tuple         =  ["""past_key_values"""]
				SCREAMING_SNAKE_CASE_   :      Any         =  {"""hidden_size""": """d_model""", """num_attention_heads""": """num_heads""", """num_hidden_layers""": """num_layers"""}

				def __init__(	self							:				List[str]     ,  UpperCamelCase__							:				List[str]=3_2_1_2_8     ,  UpperCamelCase__							:				str=7_6_8     ,  UpperCamelCase__							:				Union[str, Any]=6_4     ,  UpperCamelCase__							:				Optional[int]=2_0_4_8     ,  UpperCamelCase__							:				Dict=6_4     ,  UpperCamelCase__							:				List[str]=1_2     ,  UpperCamelCase__							:				Dict=3     ,  UpperCamelCase__							:				Any=1_2     ,  UpperCamelCase__							:				Dict=3     ,  UpperCamelCase__							:				Dict=1_2     ,  UpperCamelCase__							:				List[str]=8     ,  UpperCamelCase__							:				Optional[Any]=False     ,  UpperCamelCase__							:				Optional[int]=0.01     ,  UpperCamelCase__							:				Optional[int]="float32"     ,  UpperCamelCase__							:				Any=False     ,  UpperCamelCase__							:				Dict=3_2     ,  UpperCamelCase__							:				Union[str, Any]=1_2_8     ,  UpperCamelCase__							:				Union[str, Any]=0.1     ,  UpperCamelCase__							:				Optional[int]=1e-6     ,  UpperCamelCase__							:				Optional[Any]=0.001     ,  UpperCamelCase__							:				Dict=0.001     ,  UpperCamelCase__							:				int=1.0     ,  UpperCamelCase__							:				str="relu"     ,  UpperCamelCase__							:				int=True     ,  UpperCamelCase__							:				Any=False     ,  UpperCamelCase__							:				str=True     ,  UpperCamelCase__							:				Tuple=0     ,  UpperCamelCase__							:				str=1     ,  **UpperCamelCase__							:				Tuple     ,  )->   Optional[Any]:



						'''simple docstring'''


						__lowerCAmelCase:       int          =			vocab_size
						__lowerCAmelCase:       str          =			d_model
						__lowerCAmelCase:       str          =			d_kv
						__lowerCAmelCase:       str          =			d_ff

						__lowerCAmelCase:       List[str]          =			num_sparse_encoder_layers

						__lowerCAmelCase:       List[Any]          =			num_layers
						__lowerCAmelCase:       Optional[Any]          =			(
						    num_decoder_layers if num_decoder_layers is not None else self.num_layers
						)  # default = symmetry
						__lowerCAmelCase:       Tuple          =			num_sparse_decoder_layers

						# This tells us, each how many encoder layer we'll have to set a sparse layer.
						if self.num_sparse_encoder_layers > 0:
								__lowerCAmelCase:       int          =			self.num_layers // self.num_sparse_encoder_layers
						else:
								__lowerCAmelCase:       Union[str, Any]          =			self.num_layers  # HACK: this will create 0 sparse layers

						# This tells us, each how many encoder layer we'll have to set a sparse layer.
						if self.num_sparse_decoder_layers > 0:
								__lowerCAmelCase:       Dict          =			self.num_decoder_layers // self.num_sparse_decoder_layers
						else:
								__lowerCAmelCase:       Any          =			self.num_decoder_layers  # HACK: this will create 0 sparse layers

						__lowerCAmelCase:       Dict          =			num_heads
						__lowerCAmelCase:       Dict          =			num_experts
						__lowerCAmelCase:       Any          =			expert_capacity
						__lowerCAmelCase:       List[Any]          =			router_bias
						__lowerCAmelCase:       int          =			router_jitter_noise
						if router_dtype not in ["float32", "float16", "bfloat16"]:
								raise ValueError(f"`router_dtype` must be one of 'float32', 'float16' or 'bfloat16', got {router_dtype}")
						__lowerCAmelCase:       Dict          =			router_dtype

						__lowerCAmelCase:       Optional[Any]          =			router_ignore_padding_tokens
						__lowerCAmelCase:       Union[str, Any]          =			relative_attention_num_buckets
						__lowerCAmelCase:       str          =			relative_attention_max_distance

						__lowerCAmelCase:       Optional[int]          =			dropout_rate
						__lowerCAmelCase:       Optional[Any]          =			layer_norm_epsilon
						__lowerCAmelCase:       int          =			initializer_factor
						__lowerCAmelCase:       Tuple          =			feed_forward_proj
						__lowerCAmelCase:       int          =			use_cache
						__lowerCAmelCase:       int          =			add_router_probs

						__lowerCAmelCase:       Optional[Any]          =			router_z_loss_coef
						__lowerCAmelCase:       Dict          =			router_aux_loss_coef

						__lowerCAmelCase:       Union[str, Any]          =			self.feed_forward_proj.split("-")
						__lowerCAmelCase:       Tuple          =			act_info[-1]
						__lowerCAmelCase:       str          =			act_info[0] == "gated"

						if len(UpperCamelCase__) > 1 and act_info[0] != "gated" or len(UpperCamelCase__) > 2:
								raise ValueError(
								    f"`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."
								    "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. "
								    "'gated-gelu' or 'relu'")

						# for backwards compatibility
						if feed_forward_proj == "gated-gelu":
								__lowerCAmelCase:       List[str]          =			"gelu_new"

						super().__init__(
						    pad_token_id=UpperCamelCase__     ,  eos_token_id=UpperCamelCase__     ,  is_encoder_decoder=UpperCamelCase__     ,  **UpperCamelCase__     ,  )


