
"""simple docstring"""







import inspect
from typing import Optional, Union

import numpy as np
import PIL
import torch
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPFeatureExtractor, CLIPModel, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    DPMSolverMultistepScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
from diffusers.utils import (
    PIL_INTERPOLATION,
    randn_tensor,
)
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )  ->     Dict:







							"""simple docstring"""


							if isinstance(SCREAMING_SNAKE_CASE							,    torch.Tensor ):
														return image
							elif isinstance(SCREAMING_SNAKE_CASE							,    PIL.Image.Image ):
														__snake_case     			=       [image]

							if isinstance(image[0]							,    PIL.Image.Image ):
														__snake_case     			=       [np.array(i.resize((w, h)							,    resample=PIL_INTERPOLATION["lanczos"] ) )[None, :] for i in image]
														__snake_case     			=       np.concatenate(SCREAMING_SNAKE_CASE							,    axis=0 )
														__snake_case     			=       np.array(SCREAMING_SNAKE_CASE ).astype(np.floataa ) / 255.0
														__snake_case     			=       image.transpose(0							,    3							,    1							,    2 )
														__snake_case     			=       2.0 * image - 1.0
														__snake_case     			=       torch.from_numpy(SCREAMING_SNAKE_CASE )
							elif isinstance(image[0]							,    torch.Tensor ):
														__snake_case     			=       torch.cat(SCREAMING_SNAKE_CASE							,    dim=0 )
							return image
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE=0.9_995 )  ->     Union[str, Any]:







							"""simple docstring"""


							if not isinstance(SCREAMING_SNAKE_CASE							,    np.ndarray ):
														__snake_case     			=       True
														__snake_case     			=       va.device
														__snake_case     			=       va.cpu().numpy()
														__snake_case     			=       va.cpu().numpy()

							__snake_case     			=       np.sum(va * va / (np.linalg.norm(SCREAMING_SNAKE_CASE ) * np.linalg.norm(SCREAMING_SNAKE_CASE )) )
							if np.abs(SCREAMING_SNAKE_CASE ) > DOT_THRESHOLD:
														__snake_case     			=       (1 - t) * va + t * va
							else:
														__snake_case     			=       np.arccos(SCREAMING_SNAKE_CASE )
														__snake_case     			=       np.sin(SCREAMING_SNAKE_CASE )
														__snake_case     			=       theta_a * t
														__snake_case     			=       np.sin(SCREAMING_SNAKE_CASE )
														__snake_case     			=       np.sin(theta_a - theta_t ) / sin_theta_a
														__snake_case     			=       sin_theta_t / sin_theta_a
														__snake_case     			=       sa * va + sa * va

							if inputs_are_torch:
														__snake_case     			=       torch.from_numpy(SCREAMING_SNAKE_CASE ).to(SCREAMING_SNAKE_CASE )

							return va
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )  ->     int:







							"""simple docstring"""


							__snake_case     			=       F.normalize(SCREAMING_SNAKE_CASE							,    dim=-1 )
							__snake_case     			=       F.normalize(SCREAMING_SNAKE_CASE							,    dim=-1 )
							return (x - y).norm(dim=-1 ).div(2 ).arcsin().pow(2 ).mul(2 )
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )  ->     int:







							"""simple docstring"""


							for param in model.parameters():
														__snake_case     			=       value







class 							__magic_name__     (		lowercase__   ):

					'''simple docstring'''





					def __init__(				self   :							Tuple							,						snake_case_   :							AutoencoderKL							,						snake_case_   :							CLIPTextModel							,						snake_case_   :							CLIPModel							,						snake_case_   :							CLIPTokenizer							,						snake_case_   :							UNetaDConditionModel							,						snake_case_   :							Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler]							,						snake_case_   :							CLIPFeatureExtractor							,						snake_case_   :							List[str]=None							,						snake_case_   :							Optional[int]=None							,						snake_case_   :							Dict=None							,						):
												super().__init__()
												self.register_modules(
												    vae=snake_case_							,						text_encoder=snake_case_							,						clip_model=snake_case_							,						tokenizer=snake_case_							,						unet=snake_case_							,						scheduler=snake_case_							,						feature_extractor=snake_case_							,						coca_model=snake_case_							,						coca_tokenizer=snake_case_							,						coca_transform=snake_case_							,						)
												__snake_case     			=       (
												    feature_extractor.size
												    if isinstance(feature_extractor.size							,						snake_case_    )
												    else feature_extractor.size["shortest_edge"]
												)
												__snake_case     			=       transforms.Normalize(mean=feature_extractor.image_mean							,						std=feature_extractor.image_std    )
												set_requires_grad(self.text_encoder							,						snake_case_    )
												set_requires_grad(self.clip_model							,						snake_case_    )





					def 							lowerCAmelCase      (				self   :							Union[str, Any]							,						snake_case_   :							Optional[Union[str, int]] = "auto"    ):
												if slice_size == "auto":
																			# half the attention head size is usually a good trade-off between
																			# speed and memory
																			__snake_case     			=       self.unet.config.attention_head_dim // 2
												self.unet.set_attention_slice(snake_case_    )





					def 							lowerCAmelCase      (				self   :							Dict    ):
												self.enable_attention_slicing(snake_case_    )





					def 							lowerCAmelCase      (				self   :							Optional[int]    ):
												set_requires_grad(self.vae							,						snake_case_    )





					def 							lowerCAmelCase      (				self   :							Union[str, Any]    ):
												set_requires_grad(self.vae							,						snake_case_    )





					def 							lowerCAmelCase      (				self   :							Dict    ):
												set_requires_grad(self.unet							,						snake_case_    )





					def 							lowerCAmelCase      (				self   :							Union[str, Any]    ):
												set_requires_grad(self.unet							,						snake_case_    )





					def 							lowerCAmelCase      (				self   :							Tuple							,						snake_case_   :							Any							,						snake_case_   :							str							,						snake_case_   :							int    ):
												# get the original timestep using init_timestep
												__snake_case     			=       min(int(num_inference_steps * strength    )							,						snake_case_    )

												__snake_case     			=       max(num_inference_steps - init_timestep							,						0    )
												__snake_case     			=       self.scheduler.timesteps[t_start:]

												return timesteps, num_inference_steps - t_start





					def 							lowerCAmelCase      (				self   :							Optional[int]							,						snake_case_   :							Optional[int]							,						snake_case_   :							Optional[Any]							,						snake_case_   :							Dict							,						snake_case_   :							str							,						snake_case_   :							Any							,						snake_case_   :							int=None    ):
												if not isinstance(snake_case_							,						torch.Tensor    ):
																			raise ValueError(F'''`image` has to be of type `torch.Tensor` but is {type(snake_case_    )}'''    )

												__snake_case     			=       image.to(device=snake_case_							,						dtype=snake_case_    )

												if isinstance(snake_case_							,						snake_case_    ):
																			__snake_case     			=       [
																			    self.vae.encode(image[i : i + 1]    ).latent_dist.sample(generator[i]    ) for i in range(snake_case_    )
																			]
																			__snake_case     			=       torch.cat(snake_case_							,						dim=0    )
												else:
																			__snake_case     			=       self.vae.encode(snake_case_    ).latent_dist.sample(snake_case_    )

												# Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
												__snake_case     			=       0.18215 * init_latents
												__snake_case     			=       init_latents.repeat_interleave(snake_case_							,						dim=0    )

												__snake_case     			=       randn_tensor(init_latents.shape							,						generator=snake_case_							,						device=snake_case_							,						dtype=snake_case_    )

												# get latents
												__snake_case     			=       self.scheduler.add_noise(snake_case_							,						snake_case_							,						snake_case_    )
												__snake_case     			=       init_latents

												return latents





					def 							lowerCAmelCase      (				self   :							Any							,						snake_case_   :							Optional[int]    ):
												__snake_case     			=       self.coca_transform(snake_case_    ).unsqueeze(0    )
												with torch.no_grad(), torch.cuda.amp.autocast():
																			__snake_case     			=       self.coca_model.generate(transformed_image.to(device=self.device							,						dtype=self.coca_model.dtype    )    )
												__snake_case     			=       self.coca_tokenizer.decode(generated[0].cpu().numpy()    )
												return generated.split("<end_of_text>"    )[0].replace("<start_of_text>"							,						""    ).rstrip(" .,"    )





					def 							lowerCAmelCase      (				self   :							Any							,						snake_case_   :							Tuple							,						snake_case_   :							Optional[Any]    ):
												__snake_case     			=       self.feature_extractor.preprocess(snake_case_    )
												__snake_case     			=       torch.from_numpy(clip_image_input["pixel_values"][0]    ).unsqueeze(0    ).to(self.device    ).half()
												__snake_case     			=       self.clip_model.get_image_features(snake_case_    )
												__snake_case     			=       image_embeddings_clip / image_embeddings_clip.norm(p=2							,						dim=-1							,						keepdim=snake_case_    )
												__snake_case     			=       image_embeddings_clip.repeat_interleave(snake_case_							,						dim=0    )
												return image_embeddings_clip





					@torch.enable_grad()
					def 							lowerCAmelCase      (				self   :							Any							,						snake_case_   :							List[str]							,						snake_case_   :							Tuple							,						snake_case_   :							Union[str, Any]							,						snake_case_   :							List[str]							,						snake_case_   :							Optional[int]							,						snake_case_   :							Tuple							,						snake_case_   :							List[Any]							,						):
												__snake_case     			=       latents.detach().requires_grad_()

												__snake_case     			=       self.scheduler.scale_model_input(snake_case_							,						snake_case_    )

												# predict the noise residual
												__snake_case     			=       self.unet(snake_case_							,						snake_case_							,						encoder_hidden_states=snake_case_    ).sample

												if isinstance(self.scheduler							,						(PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)    ):
																			__snake_case     			=       self.scheduler.alphas_cumprod[timestep]
																			__snake_case     			=       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
																			__snake_case     			=       (latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

																			__snake_case     			=       torch.sqrt(snake_case_    )
																			__snake_case     			=       pred_original_sample * (fac) + latents * (1 - fac)
												elif isinstance(self.scheduler							,						snake_case_    ):
																			__snake_case     			=       self.scheduler.sigmas[index]
																			__snake_case     			=       latents - sigma * noise_pred
												else:
																			raise ValueError(F'''scheduler type {type(self.scheduler    )} not supported'''    )

												# Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
												__snake_case     			=       1 / 0.18215 * sample
												__snake_case     			=       self.vae.decode(snake_case_    ).sample
												__snake_case     			=       (image / 2 + 0.5).clamp(0							,						1    )

												__snake_case     			=       transforms.Resize(self.feature_extractor_size    )(snake_case_    )
												__snake_case     			=       self.normalize(snake_case_    ).to(latents.dtype    )

												__snake_case     			=       self.clip_model.get_image_features(snake_case_    )
												__snake_case     			=       image_embeddings_clip / image_embeddings_clip.norm(p=2							,						dim=-1							,						keepdim=snake_case_    )

												__snake_case     			=       spherical_dist_loss(snake_case_							,						snake_case_    ).mean() * clip_guidance_scale

												__snake_case     			=       -torch.autograd.grad(snake_case_							,						snake_case_    )[0]

												if isinstance(self.scheduler							,						snake_case_    ):
																			__snake_case     			=       latents.detach() + grads * (sigma**2)
																			__snake_case     			=       noise_pred_original
												else:
																			__snake_case     			=       noise_pred_original - torch.sqrt(snake_case_    ) * grads
												return noise_pred, latents





					@torch.no_grad()
					def __call__(				self   :							Dict							,						snake_case_   :							Union[torch.FloatTensor, PIL.Image.Image]							,						snake_case_   :							Union[torch.FloatTensor, PIL.Image.Image]							,						snake_case_   :							Optional[str] = None							,						snake_case_   :							Optional[str] = None							,						snake_case_   :							Optional[int] = 512							,						snake_case_   :							Optional[int] = 512							,						snake_case_   :							float = 0.6							,						snake_case_   :							Optional[int] = 50							,						snake_case_   :							Optional[float] = 7.5							,						snake_case_   :							Optional[int] = 1							,						snake_case_   :							float = 0.0							,						snake_case_   :							Optional[float] = 100							,						snake_case_   :							Optional[torch.Generator] = None							,						snake_case_   :							Optional[str] = "pil"							,						snake_case_   :							bool = True							,						snake_case_   :							float = 0.8							,						snake_case_   :							float = 0.1							,						snake_case_   :							float = 0.1							,						):
												if isinstance(snake_case_							,						snake_case_    ) and len(snake_case_    ) != batch_size:
																			raise ValueError(F'''You have passed {batch_size} batch_size, but only {len(snake_case_    )} generators.'''    )

												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 isinstance(snake_case_							,						torch.Generator    ) and batch_size > 1:
																			__snake_case     			=       [generator] + [None] * (batch_size - 1)

												__snake_case     			=       [
												    ("model", self.coca_model is None),
												    ("tokenizer", self.coca_tokenizer is None),
												    ("transform", self.coca_transform is None),
												]
												__snake_case     			=       [x[0] for x in coca_is_none if x[1]]
												__snake_case     			=       ", ".join(snake_case_    )
												# generate prompts with coca model if prompt is None
												if content_prompt is None:
																			if len(snake_case_    ):
																										raise ValueError(
																										    F'''Content prompt is None and CoCa [{coca_is_none_str}] is None.'''
																										    F'''Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.'''    )
																			__snake_case     			=       self.get_image_description(snake_case_    )
												if style_prompt is None:
																			if len(snake_case_    ):
																										raise ValueError(
																										    F'''Style prompt is None and CoCa [{coca_is_none_str}] is None.'''
																										    F''' Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.'''    )
																			__snake_case     			=       self.get_image_description(snake_case_    )

												# get prompt text embeddings for content and style
												__snake_case     			=       self.tokenizer(
												    snake_case_							,						padding="max_length"							,						max_length=self.tokenizer.model_max_length							,						truncation=snake_case_							,						return_tensors="pt"							,						)
												__snake_case     			=       self.text_encoder(content_text_input.input_ids.to(self.device    )    )[0]

												__snake_case     			=       self.tokenizer(
												    snake_case_							,						padding="max_length"							,						max_length=self.tokenizer.model_max_length							,						truncation=snake_case_							,						return_tensors="pt"							,						)
												__snake_case     			=       self.text_encoder(style_text_input.input_ids.to(self.device    )    )[0]

												__snake_case     			=       slerp(snake_case_							,						snake_case_							,						snake_case_    )

												# duplicate text embeddings for each generation per prompt
												__snake_case     			=       text_embeddings.repeat_interleave(snake_case_							,						dim=0    )

												# set timesteps
												__snake_case     			=       "offset" in set(inspect.signature(self.scheduler.set_timesteps    ).parameters.keys()    )
												__snake_case     			=       {}
												if accepts_offset:
																			__snake_case     			=       1

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

												__snake_case			,   __snake_case     			=       self.get_timesteps(snake_case_							,						snake_case_							,						self.device    )
												__snake_case     			=       timesteps[:1].repeat(snake_case_    )

												# Preprocess image
												__snake_case     			=       preprocess(snake_case_							,						snake_case_							,						snake_case_    )
												__snake_case     			=       self.prepare_latents(
												    snake_case_							,						snake_case_							,						snake_case_							,						text_embeddings.dtype							,						self.device							,						snake_case_    )

												__snake_case     			=       preprocess(snake_case_							,						snake_case_							,						snake_case_    )
												__snake_case     			=       self.prepare_latents(
												    snake_case_							,						snake_case_							,						snake_case_							,						text_embeddings.dtype							,						self.device							,						snake_case_    )

												__snake_case     			=       slerp(snake_case_							,						snake_case_							,						snake_case_    )

												if clip_guidance_scale > 0:
																			__snake_case     			=       self.get_clip_image_embeddings(snake_case_							,						snake_case_    )
																			__snake_case     			=       self.get_clip_image_embeddings(snake_case_							,						snake_case_    )
																			__snake_case     			=       slerp(
																			    snake_case_							,						snake_case_							,						snake_case_    )

												# 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.
												__snake_case     			=       guidance_scale > 1.0
												# get unconditional embeddings for classifier free guidance
												if do_classifier_free_guidance:
																			__snake_case     			=       content_text_input.input_ids.shape[-1]
																			__snake_case     			=       self.tokenizer([""]							,						padding="max_length"							,						max_length=snake_case_							,						return_tensors="pt"    )
																			__snake_case     			=       self.text_encoder(uncond_input.input_ids.to(self.device    )    )[0]
																			# duplicate unconditional embeddings for each generation per prompt
																			__snake_case     			=       uncond_embeddings.repeat_interleave(snake_case_							,						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
																			__snake_case     			=       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`.
												__snake_case     			=       (batch_size, self.unet.config.in_channels, height // 8, width // 8)
												__snake_case     			=       text_embeddings.dtype
												if latents is None:
																			if self.device.type == "mps":
																										# randn does not work reproducibly on mps
																										__snake_case     			=       torch.randn(snake_case_							,						generator=snake_case_							,						device="cpu"							,						dtype=snake_case_    ).to(
																										    self.device    )
																			else:
																										__snake_case     			=       torch.randn(snake_case_							,						generator=snake_case_							,						device=self.device							,						dtype=snake_case_    )
												else:
																			if latents.shape != latents_shape:
																										raise ValueError(F'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}'''    )
																			__snake_case     			=       latents.to(self.device    )

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

												# check if the scheduler accepts generator
												__snake_case     			=       "generator" in set(inspect.signature(self.scheduler.step    ).parameters.keys()    )
												if accepts_generator:
																			__snake_case     			=       generator

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

																										# predict the noise residual
																										__snake_case     			=       self.unet(snake_case_							,						snake_case_							,						encoder_hidden_states=snake_case_    ).sample

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

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

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

        # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
												__snake_case     			=       1 / 0.18215 * latents
												__snake_case     			=       self.vae.decode(snake_case_    ).sample

												__snake_case     			=       (image / 2 + 0.5).clamp(0							,						1    )
												__snake_case     			=       image.cpu().permute(0							,						2							,						3							,						1    ).numpy()

												if output_type == "pil":
																			__snake_case     			=       self.numpy_to_pil(snake_case_    )

												if not return_dict:
																			return (image, None)

												return StableDiffusionPipelineOutput(images=snake_case_							,						nsfw_content_detected=snake_case_    )


"""simple docstring"""







from __future__ import annotations

from typing import Any
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE )  ->     int:







							"""simple docstring"""


							if not postfix_notation:
														return 0

							__snake_case     			=       {"+", "-", "*", "/"}
							__snake_case     			=       []

							for token in postfix_notation:
														if token in operations:
																					__snake_case			,   __snake_case     			=       stack.pop(), stack.pop()
																					if token == "+":
																												stack.append(a + b )
																					elif token == "-":
																												stack.append(a - b )
																					elif token == "*":
																												stack.append(a * b )
																					else:
																												if a * b < 0 and a % b != 0:
																																			stack.append(a // b + 1 )
																												else:
																																			stack.append(a // b )
														else:
																					stack.append(int(SCREAMING_SNAKE_CASE ) )

							return stack.pop()


if __name__ == "__main__":
		import doctest

		doctest.testmod()
